pax_global_header00006660000000000000000000000064150675266560014534gustar00rootroot0000000000000052 comment=7e16503e403de460d4ed7dfd6c90a656628f5149 bitsandbytes/000077500000000000000000000000001506752665600135735ustar00rootroot00000000000000bitsandbytes/.clang-format000066400000000000000000000016141506752665600161500ustar00rootroot00000000000000--- BasedOnStyle: LLVM AlignAfterOpenBracket: BlockIndent BinPackArguments: true BinPackParameters: true BracedInitializerIndentWidth: 4 ColumnLimit: 120 Cpp11BracedListStyle: true IndentWidth: 4 IndentWrappedFunctionNames: true PointerAlignment: Left SeparateDefinitionBlocks: Always Standard: c++17 StatementMacros: - 'MAKE_PreconditionOptimizer32bit1State' - 'MAKE_PreconditionOptimizer32bit2State' - 'MAKE_PreconditionStatic8bit1State' - 'MAKE_PreconditionStatic8bit2State' - 'MAKE_Optimizer32bit1State' - 'MAKE_optimizerStatic8bit1State' - 'MAKE_optimizerStatic8bit2State' - 'MAKE_OptimizerStatic8bit1StateBlockwise' - 'MAKE_OptimizerStatic8bit2StateBlockwise' - 'MAKE_kQuantizeBlockwise' - 'MAKE_BLOCKWISE8' - 'MAKE_ELEMENTWISE_FUNC' - 'CMAKE_ELEMENTWISE_FUNC' - 'MAKE_FUNC8' - 'MAKE_FUNC32' - 'MAKE_CBLOCKWISE8' - 'MAKE_CFUNC8' - 'MAKE_CFUNC32' UseTab: Never ... bitsandbytes/.editorconfig000066400000000000000000000001001506752665600162370ustar00rootroot00000000000000[*] trim_trailing_whitespace = true insert_final_newline = true bitsandbytes/.git-blame-ignore-revs000066400000000000000000000011101506752665600176640ustar00rootroot00000000000000# ran black and isort for coherent code formatting bfa0e33294f2b1dc25e65a33be2397f989824298 # reran black with linelength 80 for greater readability ea7c14f8ef64924f2d0ff80df3cdabf2c7299848 # Remove f-prefix from strings that don't use formatting 7727fa4c8c6c1ef2b109120aff4196a0a6bf3ed6 # format tests/linear_4bit.py 34735ba89de8235ea9da6ef409f814dcea9e2038 # Reformat with ruff-format 5a4263f4dc05fe8f78f4111beab9f68a81deeab1 # CHANGELOG: to reverse chron order + mdformat 4743ff0d43e04e4cc3e5d8b9e7cd016c0defa36d # Apply clang-format 4955d136ae083c2be1236d8915913166e1790aad bitsandbytes/.github/000077500000000000000000000000001506752665600151335ustar00rootroot00000000000000bitsandbytes/.github/FUNDING.yml000066400000000000000000000000361506752665600167470ustar00rootroot00000000000000open_collective: bitsandbytes bitsandbytes/.github/ISSUE_TEMPLATE/000077500000000000000000000000001506752665600173165ustar00rootroot00000000000000bitsandbytes/.github/ISSUE_TEMPLATE/bug-report.yml000066400000000000000000000017221506752665600221310ustar00rootroot00000000000000name: "\U0001F41B Bug Report" description: Submit a bug report to help us improve bitsandbytes body: - type: textarea id: system-info attributes: label: System Info description: Please share your relevant system information with us placeholder: platform, python version, hardware, ... validations: required: true - type: textarea id: reproduction validations: required: true attributes: label: Reproduction description: | Please provide a code sample that reproduces the problem you ran into. It can be a Colab link or just a code snippet. Please provide the simplest reproducer as possible so that we can quickly fix the issue. placeholder: | Reproducer: - type: textarea id: expected-behavior validations: required: true attributes: label: Expected behavior description: "A clear and concise description of what you would expect to happen." bitsandbytes/.github/ISSUE_TEMPLATE/feature-request.yml000066400000000000000000000014241506752665600231630ustar00rootroot00000000000000name: "\U0001F680 Feature request" description: Submit a proposal/request for a new feature labels: ["feature"] body: - type: textarea id: feature-request validations: required: true attributes: label: Feature request description: | A clear and concise description of the feature proposal. - type: textarea id: motivation validations: required: true attributes: label: Motivation description: | Please outline the motivation for the proposal. Is your feature request related to a problem? - type: textarea id: contribution validations: required: true attributes: label: Your contribution description: | Is there any way that you could help, e.g. by submitting a PR? bitsandbytes/.github/dependabot.yml.disabled000066400000000000000000000003301506752665600215250ustar00rootroot00000000000000version: 2 updates: - package-ecosystem: pip directory: "/" schedule: interval: "weekly" groups: major: update-types: [major] minor-patch: update-types: [minor, patch] bitsandbytes/.github/scripts/000077500000000000000000000000001506752665600166225ustar00rootroot00000000000000bitsandbytes/.github/scripts/auditwheel_show.py000077500000000000000000000013431506752665600223730ustar00rootroot00000000000000import argparse import subprocess def main(): ap = argparse.ArgumentParser() ap.add_argument("wheels", nargs="*") args = ap.parse_args() if not args.wheels: ap.error("At least one wheel must be provided.") for whl in args.wheels: print(f"### `{whl}`") audit_wheel_output = subprocess.run( ["auditwheel", "show", whl], capture_output=True, text=True, errors="backslashreplace", ) if audit_wheel_output.stdout: print(audit_wheel_output.stdout) if audit_wheel_output.stderr: print(f"**Error:**\n```\n{audit_wheel_output.stderr}\n```") print("---") if __name__ == "__main__": main() bitsandbytes/.github/scripts/build-cpu.sh000066400000000000000000000007011506752665600210400ustar00rootroot00000000000000#!/bin/bash declare build_arch declare build_os set -xeuo pipefail pip install cmake==3.28.3 if [ "${build_os:0:5}" == macos ] && [ "${build_arch}" == aarch64 ]; then cmake -DCMAKE_OSX_ARCHITECTURES=arm64 -DCOMPUTE_BACKEND=cpu . else cmake -DCOMPUTE_BACKEND=cpu . fi cmake --build . --config Release output_dir="output/${build_os}/${build_arch}" mkdir -p "${output_dir}" (shopt -s nullglob && cp bitsandbytes/*.{so,dylib,dll} "${output_dir}") bitsandbytes/.github/scripts/build-cuda.sh000066400000000000000000000037251506752665600211760ustar00rootroot00000000000000#!/bin/bash declare build_arch declare build_os declare cuda_version declare cuda_targets set -xeuo pipefail if [[ -v cuda_targets ]]; then build_capability="${cuda_targets}" elif [ "${build_arch}" = "aarch64" ]; then build_capability="75;80;90" # CUDA 12.8-12.9: Add sm100/sm120 [[ "${cuda_version}" == 12.8.* || "${cuda_version}" == 12.9.* ]] && build_capability="75;80;90;100;120" # CUDA 13.0+: Add sm100/sm110/sm120 [[ "${cuda_version}" == 13.*.* ]] && build_capability="75;80;90;100;110;120" else # By default, target Pascal through Hopper. build_capability="60;70;75;80;86;89;90" # CUDA 12.8+: Add sm100 and sm120; remove < sm70 to align with PyTorch 2.8+cu128 minimum [[ "${cuda_version}" == 12.8.* || "${cuda_version}" == 12.9.* ]] && build_capability="70;75;80;86;89;90;100;120" # CUDA 13.0+: Remove < sm75 to align with PyTorch 2.9+cu130 minimum [[ "${cuda_version}" == 13.*.* ]] && build_capability="75;80;86;89;90;100;120" fi [[ "${build_os}" = windows-* ]] && python3 -m pip install ninja if [ "${build_os:0:6}" == ubuntu ]; then # We'll use Rocky Linux 8 in order to maintain manylinux 2.24 compatibility. image="nvidia/cuda:${cuda_version}-devel-rockylinux8" echo "Using image $image" docker run -i -w /src -v "$PWD:/src" "$image" bash -c \ "dnf -y --refresh update --security \ && dnf -y install cmake gcc-toolset-11 --setopt=install_weak_deps=False --setopt=tsflags=nodocs \ && source scl_source enable gcc-toolset-11 \ && cmake -DCOMPUTE_BACKEND=cuda -DCOMPUTE_CAPABILITY=\"${build_capability}\" . \ && cmake --build . --config Release" else pip install cmake==3.28.3 cmake -G Ninja -DCOMPUTE_BACKEND=cuda -DCOMPUTE_CAPABILITY="${build_capability}" -DCMAKE_BUILD_TYPE=Release -S . cmake --build . --config Release fi output_dir="output/${build_os}/${build_arch}" mkdir -p "${output_dir}" (shopt -s nullglob && cp bitsandbytes/*.{so,dylib,dll} "${output_dir}") bitsandbytes/.github/scripts/build-rocm.sh000066400000000000000000000013021506752665600212070ustar00rootroot00000000000000#!/bin/bash declare build_arch declare build_os declare rocm_version set -xeuo pipefail bnb_rocm_arch="gfx90a;gfx942;gfx1100" if [ "${build_os:0:6}" == ubuntu ]; then image=rocm/dev-ubuntu-22.04:${rocm_version}-complete echo "Using image $image" docker run --rm --platform "linux/$build_arch" -i \ -w /src -v "$PWD:/src" "$image" sh -c \ "apt-get update \ && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends cmake \ && cmake -DCOMPUTE_BACKEND=hip -DBNB_ROCM_ARCH=\"${bnb_rocm_arch}\" . \ && cmake --build ." fi output_dir="output/${build_os}/${build_arch}" mkdir -p "${output_dir}" (shopt -s nullglob && cp bitsandbytes/*.{so,dylib,dll} "${output_dir}") bitsandbytes/.github/scripts/build-xpu.sh000077500000000000000000000014051506752665600210720ustar00rootroot00000000000000#!/bin/bash declare build_os set -xeuo pipefail # We currently only build XPU on Linux. if [ "${build_os:0:6}" == ubuntu ]; then # TODO: We might want to pre-build this as our own customized image in the future. image=intel/deep-learning-essentials:2025.1.3-0-devel-ubuntu22.04 echo "Using image $image" docker run --rm -i \ -w /src -v "$PWD:/src" "$image" sh -c \ "apt-get update \ && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ cmake bison intel-fw-gpu intel-ocloc \ && cmake -DCOMPUTE_BACKEND=xpu . \ && cmake --build . --config Release" fi output_dir="output/${build_os}/x86_64" mkdir -p "${output_dir}" (shopt -s nullglob && cp bitsandbytes/*.{so,dylib,dll} "${output_dir}") bitsandbytes/.github/scripts/set_platform_tag.py000066400000000000000000000015701506752665600225310ustar00rootroot00000000000000import argparse import platform import sys def get_platform_tag(architecture): system = platform.system() if system == "Linux": tag = "manylinux_2_24_x86_64" if architecture == "x86_64" else "manylinux_2_24_aarch64" elif system == "Darwin": tag = "macosx_13_1_x86_64" if architecture == "x86_64" else "macosx_13_1_arm64" elif system == "Windows": tag = "win_amd64" if architecture == "x86_64" else "win_arm64" else: sys.exit(f"Unsupported system: {system}") return tag def main(): parser = argparse.ArgumentParser(description="Determine platform tag.") parser.add_argument("arch", type=str, help="Architecture (e.g., x86_64, aarch64)") args = parser.parse_args() tag = get_platform_tag(args.arch) print(tag) # This will be captured by the GitHub Actions workflow if __name__ == "__main__": main() bitsandbytes/.github/workflows/000077500000000000000000000000001506752665600171705ustar00rootroot00000000000000bitsandbytes/.github/workflows/build_documentation.yml000066400000000000000000000011241506752665600237410ustar00rootroot00000000000000name: Build documentation on: push: branches: - main - doc-builder* - v*-release jobs: build: uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@main with: commit_sha: ${{ github.sha }} package: bitsandbytes repo_owner: bitsandbytes-foundation # avoid /src suffix leading to wrong links, like bitsandbytes/blob/main/src/bitsandbytes/nn/ version_tag_suffix: '' # defaults to '/src' custom_container: huggingface/transformers-doc-builder secrets: hf_token: ${{ secrets.HUGGINGFACE_PUSH }} bitsandbytes/.github/workflows/build_pr_documentation.yml000066400000000000000000000013261506752665600244460ustar00rootroot00000000000000name: Build PR Documentation on: pull_request: concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} cancel-in-progress: true jobs: build: if: github.repository == 'bitsandbytes-foundation/bitsandbytes' uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@main with: commit_sha: ${{ github.event.pull_request.head.sha }} pr_number: ${{ github.event.number }} package: bitsandbytes repo_owner: bitsandbytes-foundation # avoid /src suffix leading to wrong links, like bitsandbytes/blob/main/src/bitsandbytes/nn/ version_tag_suffix: '' # defaults to '/src' custom_container: huggingface/transformers-doc-builder bitsandbytes/.github/workflows/lint.yml000066400000000000000000000005101506752665600206550ustar00rootroot00000000000000name: Lint on: push: branches: - main pull_request: jobs: Lint: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v4 with: python-version: "3.12" - uses: pre-commit/action@v3.0.0 env: RUFF_OUTPUT_FORMAT: github bitsandbytes/.github/workflows/python-package.yml000066400000000000000000000334251506752665600226340ustar00rootroot00000000000000name: Python package on: push: {} pull_request: branches: [main] paths: - ".github/workflows/python-package.yml" - ".github/scripts/**" - "bitsandbytes/**" - "csrc/**" - "include/**" - "tests/**" - "CMakeLists.txt" - "MANIFEST.in" - "setup.py" - "pyproject.toml" release: types: [published] workflow_dispatch: {} # Allow manual trigger workflow_call: {} # Allow triggering from other worfkflows concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: true jobs: ## # This job matrix builds the CPU versions of the libraries for all supported platforms. ## build-cpu: strategy: matrix: include: - os: ubuntu-22.04 arch: x86_64 - os: ubuntu-22.04-arm arch: aarch64 - os: windows-2025 arch: x86_64 - os: macos-15 arch: arm64 runs-on: ${{ matrix.os }} steps: - uses: actions/checkout@v4 - name: Setup MSVC if: startsWith(matrix.os, 'windows') uses: ilammy/msvc-dev-cmd@v1.13.0 # to use cl - name: Build C++ run: bash .github/scripts/build-cpu.sh env: build_os: ${{ matrix.os }} build_arch: ${{ matrix.arch }} - name: Upload build artifact uses: actions/upload-artifact@v4 with: name: shared_library_${{ matrix.os }}_${{ matrix.arch }} path: output/* retention-days: 7 ## # This job matrix builds the CUDA versions of the libraries for platforms that support CUDA (Linux x64/aarch64 + Windows x64) ## build-cuda: strategy: fail-fast: false matrix: os: [ubuntu-22.04, ubuntu-22.04-arm, windows-2025] include: - os: ubuntu-22.04 arch: x86_64 - os: ubuntu-22.04-arm arch: aarch64 - os: windows-2025 arch: x86_64 cuda_version: ["11.8.0", "12.0.1", "12.1.1", "12.2.2", "12.3.2", "12.4.1", "12.5.1", "12.6.3", "12.8.1", "12.9.1", "13.0.1"] runs-on: ${{ matrix.os }} steps: - uses: actions/checkout@v4 # Windows: We install Cuda on the agent (slow) - uses: N-Storm/cuda-toolkit@d68ba29a800229200a2c3f572f9e816d7f67cdb4 # v0.2.24m if: startsWith(matrix.os, 'windows') id: cuda-toolkit with: # Temporary: Use CUDA 13.0.0 for Windows until 13.0.1 is supported with this action. cuda: ${{ matrix.cuda_version == '13.0.1' && '13.0.0' || matrix.cuda_version }} method: "network" # The "crt" "nvvm" and "nvptxcompiler" components are added for CUDA 13. sub-packages: ${{ format('["nvcc"{0},"cudart","cusparse","cublas","thrust","cublas_dev","cusparse_dev"]', startsWith(matrix.cuda_version, '13.') && ',"crt","nvvm","nvptxcompiler"' || '') }} use-github-cache: false use-local-cache: false log-file-suffix: ${{matrix.os}}-${{matrix.cuda_version}}.txt - name: Setup MSVC if: startsWith(matrix.os, 'windows') uses: ilammy/msvc-dev-cmd@v1.13.0 # to use cl - name: Build C++ run: bash .github/scripts/build-cuda.sh env: build_os: ${{ matrix.os }} build_arch: ${{ matrix.arch }} cuda_version: ${{ matrix.cuda_version }} - name: Upload build artifact uses: actions/upload-artifact@v4 with: name: shared_library_cuda_${{ matrix.os }}_${{ matrix.arch }}_${{ matrix.cuda_version }} path: output/* retention-days: 7 build-xpu: strategy: matrix: os: [ubuntu-22.04] runs-on: ${{ matrix.os }} steps: - uses: actions/checkout@v4 - name: Build C++ run: bash .github/scripts/build-xpu.sh env: build_os: ${{ matrix.os }} - name: Upload build artifact uses: actions/upload-artifact@v4 with: name: shared_library_xpu_${{ matrix.os }}_x86_64 path: output/* retention-days: 7 build-rocm: if: false # Temporary strategy: matrix: os: [ubuntu-22.04] arch: [x86_64] rocm_version: ["6.1.2", "6.2.4", "6.3.4", "6.4.4", "7.0"] runs-on: ${{ matrix.os }} steps: - uses: actions/checkout@v4 - name: Clean up disk space run: | sudo rm -rf \ /usr/share/dotnet \ /opt/ghc \ "/usr/local/share/boost" \ "$AGENT_TOOLSDIRECTORY" \ /opt/hostedtoolcache \ /opt/google/chrome \ /opt/microsoft/msedge \ /opt/microsoft/powershell \ /opt/pipx \ /usr/lib/mono \ /usr/local/julia* \ /usr/local/lib/android \ /usr/local/lib/node_modules \ /usr/local/share/chromium \ /usr/local/share/powershell \ /usr/share/swift - name: Build C++ run: bash .github/scripts/build-rocm.sh env: build_os: ${{ matrix.os }} build_arch: ${{ matrix.arch }} rocm_version: ${{ matrix.rocm_version }} - name: Upload build artifact uses: actions/upload-artifact@v4 with: name: shared_library_rocm_${{ matrix.os }}_${{ matrix.arch }}_${{ matrix.rocm_version }} path: output/* retention-days: 7 build-wheels: needs: - build-cpu - build-cuda #- build-rocm - build-xpu strategy: matrix: os: [ubuntu-22.04, ubuntu-22.04-arm, windows-2025, macos-15] include: - os: ubuntu-22.04 arch: x86_64 - os: ubuntu-22.04-arm arch: aarch64 - os: windows-2025 arch: x86_64 - os: macos-15 arch: arm64 # The specific Python version is irrelevant in this context as we are only packaging non-C extension # code. This ensures compatibility across Python versions, as compatibility is # dictated by the packaged code itself, not the Python version used for packaging. python-version: ["3.10"] runs-on: ${{ matrix.os }} steps: - uses: actions/checkout@v4 - name: Download build artifacts uses: actions/download-artifact@v4 with: merge-multiple: true pattern: "shared_library*_${{ matrix.os }}_${{ matrix.arch }}*" path: output/ - name: Copy correct platform shared library shell: bash run: | ls -lR output/ cp output/${{ matrix.os }}/${{ matrix.arch }}/* bitsandbytes/ - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} cache: pip - run: pip install build wheel - run: python -m build . - name: Determine and Set Platform Tag, then Tag Wheel shell: bash run: | PLATFORM_TAG=$(python .github/scripts/set_platform_tag.py "${{ matrix.arch }}") echo "PLATFORM_TAG=$PLATFORM_TAG" wheel tags --remove --abi-tag=none --python-tag=py3 --platform-tag=$PLATFORM_TAG dist/bitsandbytes-*.whl - name: Upload build artifact uses: actions/upload-artifact@v4 with: name: bdist_wheel_${{ matrix.os }}_${{ matrix.arch }} path: dist/bitsandbytes-*.whl retention-days: 7 upload-pre-release-wheels: name: Create release and upload artifacts runs-on: ubuntu-latest if: github.ref_name == 'main' permissions: contents: write needs: - build-wheels steps: - name: Download and rename artifacts uses: actions/download-artifact@v4 with: path: tmp/ pattern: "bdist_wheel_*" merge-multiple: true - name: Inspect tmp directory after downloading artifacts run: | ls -alFR tmp/ WHEEL_COUNT=$(find tmp/ -type f -name "*.whl" | wc -l) echo "Found $WHEEL_COUNT wheel files" if [ "$WHEEL_COUNT" -eq 0 ]; then echo "::error::No wheel files found in tmp directory! Cannot proceed with release." exit 1 fi - name: Move and rename wheel files with pattern replacement run: | mkdir -p wheels/ # The whole point of the continuous release is to have a stable download link and the only way to have a PEP 440–compliant wheel name # is to use a stable placeholder version. Otherwise, pip won't let you install the wheel. The cool thing is that we can now install the # wheel directly from the GH pre-release which gets updated continuously, e.g. # `pip install https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl` STABLE_PLACEHOLDER_VERSION="1.33.7.preview" # exclude macos wheels for now find tmp/ -type f -name '*.whl' ! -name '*macos*' -print0 | while IFS= read -r -d '' wheel; do wheel_filename=$(basename "$wheel") # Strip off the original version rest=${wheel_filename#bitsandbytes-*-} new_name="bitsandbytes-${STABLE_PLACEHOLDER_VERSION}-${rest}" echo "Renaming $wheel_filename → $new_name" mv "$wheel" "wheels/${new_name}" done - name: Inspect wheels directory after renaming files run: ls -alFR wheels/ - uses: actions/checkout@v4 with: path: repo - name: Delete old pre-release (if exists) run: | cd repo && gh release delete continuous-release_main --cleanup-tag -y env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - name: Ensure tag exists run: | cd repo git tag -f continuous-release_main git push -f origin continuous-release_main env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - name: Generate pip install commands for release body run: | cat > body.md << 'ENDOFMARKDOWN' ## Latest `main` pre-release wheel This pre-release contains the latest development wheels for all supported platforms, rebuilt automatically on every commit to the `main` branch. **How to install:** Pick the correct command for your platform and run it in your terminal: ENDOFMARKDOWN for whl in wheels/*.whl; do fname=$(basename "$whl") url="https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/$fname" if [[ "$fname" == *"manylinux_2_24_x86_64"* ]]; then echo "### Linux (x86_64)" >> body.md elif [[ "$fname" == *"manylinux_2_24_aarch64"* ]]; then echo "### Linux (ARM/aarch64)" >> body.md elif [[ "$fname" == *"win_amd64"* ]]; then echo "### Windows (x86_64)" >> body.md else echo "### Other platform" >> body.md fi echo "\`\`\`sh" >> body.md echo "pip install --force-reinstall $url" >> body.md echo "\`\`\`" >> body.md echo "" >> body.md done cat >> body.md << 'ENDOFMARKDOWN' > **Note:** > These wheels are updated automatically with every commit to `main` and become available as soon as the [python-package.yml](.github/workflows/python-package.yml) workflow finishes. The version number is replaced with 1.33.7-preview in order to keep the link stable, this however does not affect the installed version at all: ``` > pip install https://.../bitsandbytes-1.33.7-preview-py3-none-manylinux_2_24_x86_64.whl Collecting bitsandbytes==1.33.7rc0 ... Successfully installed bitsandbytes-0.46.0.dev0 ``` ENDOFMARKDOWN # for debugging: cat body.md - name: Create new pre-release and upload artifacts uses: softprops/action-gh-release@v2.2.1 with: files: wheels/*.whl prerelease: true name: Latest `main` wheel body_path: body.md tag_name: continuous-release_main make_latest: false draft: false audit-wheels: needs: build-wheels strategy: matrix: os: [ubuntu-22.04, ubuntu-22.04-arm] include: - os: ubuntu-22.04 arch: x86_64 - os: ubuntu-22.04-arm arch: aarch64 runs-on: ${{ matrix.os }} env: PIP_DISABLE_PIP_VERSION_CHECK: 1 steps: - uses: actions/checkout@v4 - name: Download wheel uses: actions/download-artifact@v4 with: name: bdist_wheel_${{ matrix.os }}_${{ matrix.arch }} path: wheels/ - name: Set up Python uses: actions/setup-python@v5 with: python-version: "3.12" - run: pip install auditwheel - run: python ./.github/scripts/auditwheel_show.py wheels/* | tee $GITHUB_STEP_SUMMARY publish-wheels: name: Publish wheels to PyPI needs: [build-wheels, audit-wheels] runs-on: ubuntu-latest if: | github.repository == 'bitsandbytes-foundation/bitsandbytes' && github.event_name == 'push' && startsWith(github.ref, 'refs/tags') environment: name: release url: https://pypi.org/p/bitsandbytes permissions: id-token: write steps: - name: Download distribution artifacts uses: actions/download-artifact@v4 with: path: dist/ pattern: "bdist_wheel_*" merge-multiple: true - name: Remove macOS wheels run: rm dist/*macos* - name: Publish to PyPI uses: pypa/gh-action-pypi-publish@release/v1 with: print-hash: true bitsandbytes/.github/workflows/stale.yml.disabled000066400000000000000000000010511506752665600225660ustar00rootroot00000000000000name: Stale Bot on: schedule: - cron: "0 15 * * *" jobs: close_stale_issues: name: Close Stale Issues if: github.repository == 'TimDettmers/bitsandbytes' runs-on: ubuntu-latest env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} steps: - uses: actions/checkout@v3 - name: Setup Python uses: actions/setup-python@v4 with: python-version: 3.8 - name: Install requirements run: | pip install PyGithub - name: Close stale issues run: | python scripts/stale.py bitsandbytes/.github/workflows/tests.yml000066400000000000000000000320571506752665600210640ustar00rootroot00000000000000name: Nightly Tests on: workflow_dispatch: schedule: # Every day at 02:15 AM UTC - cron: "15 2 * * *" concurrency: group: ${{ github.workflow }}-${{ github.ref }} cancel-in-progress: true jobs: build-cpu: strategy: matrix: os: [ubuntu-22.04, ubuntu-22.04-arm, windows-2025, macos-15] include: - os: ubuntu-22.04 arch: x86_64 - os: ubuntu-22.04-arm arch: aarch64 - os: windows-2025 arch: x86_64 - os: macos-15 arch: arm64 runs-on: ${{ matrix.os }} steps: - uses: actions/checkout@v4 - name: Setup MSVC if: startsWith(matrix.os, 'windows') uses: ilammy/msvc-dev-cmd@v1.13.0 # to use cl - name: Build C++ run: bash .github/scripts/build-cpu.sh env: build_os: ${{ matrix.os }} build_arch: ${{ matrix.arch }} - name: Upload build artifact uses: actions/upload-artifact@v4 with: name: lib_cpu_${{ matrix.os }}_${{ matrix.arch }} path: output/${{ matrix.os }}/${{ matrix.arch }}/* retention-days: 7 build-cuda: strategy: matrix: # TODO: Add 13.0.1 when we have runners with new enough drivers. cuda_version: ["11.8.0", "12.6.3", "12.8.1", "12.9.1"] os: [ubuntu-22.04, ubuntu-22.04-arm] include: - os: ubuntu-22.04 arch: x86_64 - os: ubuntu-22.04-arm arch: aarch64 - os: windows-2025 arch: x86_64 cuda_version: "11.8.0" runs-on: ${{ matrix.os }} steps: - uses: actions/checkout@v4 - name: Install CUDA Toolkit uses: Jimver/cuda-toolkit@c35baa1a18fd1fc9dcf47c5bd839bf30559c0bc3 # v0.2.24 if: startsWith(matrix.os, 'windows') id: cuda-toolkit with: cuda: ${{ matrix.cuda_version }} method: "network" sub-packages: '["nvcc","cudart","cusparse","cublas","thrust","nvrtc_dev","cublas_dev","cusparse_dev"]' use-github-cache: false - name: Setup MSVC if: startsWith(matrix.os, 'windows') uses: ilammy/msvc-dev-cmd@v1.13.0 # to use cl # Target just the architectures used for our runners. - name: Build C++ / CUDA run: bash .github/scripts/build-cuda.sh env: build_os: ${{ matrix.os }} build_arch: ${{ matrix.arch }} cuda_version: ${{ matrix.cuda_version }} cuda_targets: "75;89" - name: Upload build artifact uses: actions/upload-artifact@v4 with: name: lib_cuda_${{matrix.cuda_version}}_${{ matrix.os }}_${{ matrix.arch }} path: output/${{ matrix.os }}/${{ matrix.arch }}/* retention-days: 7 test-cpu: if: github.repository == 'bitsandbytes-foundation/bitsandbytes' needs: build-cpu strategy: fail-fast: false matrix: os: [ubuntu-22.04, ubuntu-22.04-arm, windows-2025, macos-15] # Test with the oldest supported torch version, the newest two stable/RC. torch_version: ["2.3.1", "2.7.1", "2.8.0"] include: - os: ubuntu-22.04 arch: x86_64 runner: banb-aws-general-8-plus-use1-public-80 - os: ubuntu-22.04-arm arch: aarch64 - os: ubuntu-22.04-arm arch: aarch64 torch_version: "2.5.1" # Higher minimum requirement for aarch64 - os: windows-2025 arch: x86_64 - os: macos-15 arch: arm64 exclude: - os: ubuntu-22.04-arm torch_version: "2.3.1" runs-on: ${{ matrix.runner || matrix.os }} env: BNB_TEST_DEVICE: cpu steps: - uses: actions/checkout@v4 - name: Download build artifact uses: actions/download-artifact@v4 with: name: lib_cpu_${{ matrix.os }}_${{ matrix.arch }} path: bitsandbytes/ merge-multiple: true - name: Setup Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Setup MSVC if: startsWith(matrix.os, 'windows') uses: ilammy/msvc-dev-cmd@v1.13.0 # to use cl for torch.compile - name: Install dependencies run: | pip install torch==${{ matrix.torch_version }} --index-url https://download.pytorch.org/whl/cpu pip install -e ".[test]" pip install pytest-cov # We need to downgrade to numpy<2 for torch<2.4.1 compatibility on Windows # See: https://github.com/pytorch/pytorch/issues/131668 - name: Downgrade NumPy if: startsWith(matrix.os, 'windows') && startsWith(matrix.torch_version, '2.3.') run: pip install "numpy<2" - name: Show installed packages run: pip list - name: Show environment information run: python -m torch.utils.collect_env - name: Run tests run: pytest --durations=100 test-cpu-intel: if: github.repository == 'bitsandbytes-foundation/bitsandbytes' needs: build-cpu runs-on: banb-aws-general-8-plus-use1-public-80 env: BNB_TEST_DEVICE: cpu steps: - uses: actions/checkout@v4 - name: Download build artifact uses: actions/download-artifact@v4 with: name: lib_cpu_ubuntu-22.04_x86_64 path: bitsandbytes/ merge-multiple: true - name: Setup Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Install dependencies run: | pip install torch==2.7.1 --index-url https://download.pytorch.org/whl/cpu pip install -e ".[test]" pip install pytest-cov - name: Show installed packages run: pip list - name: Show environment information run: python -m torch.utils.collect_env - name: Run tests run: pytest --durations=100 # test-cuda-aarch64: # if: github.repository == 'bitsandbytes-foundation/bitsandbytes' # needs: build-cuda # strategy: # fail-fast: false # matrix: # os: [ubuntu-22.04-arm] # arch: [aarch64] # torch_version: ["2.7.0"] # cuda_version: ["11.8.0", "12.8.1"] # runs-on: bandb-aws-g5g-4xlarge-plus-use1-public-80 # env: # BNB_TEST_DEVICE: cuda # steps: # - name: Show GPU Information # run: nvidia-smi # - name: Show pip packages # run: pip list test-hpu: if: false # github.repository == 'bitsandbytes-foundation/bitsandbytes' needs: build-cpu strategy: fail-fast: false matrix: torch_version: ["2.6.0"] runs-on: group: bandb-itac-bmemr-gaudi3-1gaudi env: BNB_TEST_DEVICE: hpu container: image: vault.habana.ai/gaudi-docker/1.21.1/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest options: --runtime=habana --shm-size=64G --env HABANA_VISIBLE_DEVICES --env HABANA_VISIBLE_MODULES env: OMPI_MCA_btl_vader_single_copy_mechanism: none BNB_TEST_DEVICE: hpu steps: - name: Show system information run: | echo "OS: $(uname -a)" echo "CPU: $(lscpu | grep 'Model name')" echo "Memory: $(free -h)" - name: Show HPU Information run: | hl-smi - uses: actions/checkout@v4 - name: Download build artifact uses: actions/download-artifact@v4 with: name: lib_cpu_ubuntu-22.04_x86_64 path: bitsandbytes/ merge-multiple: true - name: Show installed packages run: pip list - name: Install dependencies run: | pip install -e ".[test]" pip install pytest-cov - name: Show installed packages run: pip list - name: Show environment information run: | python -m torch.utils.collect_env python -m bitsandbytes - name: Run tests run: pytest --durations=100 test-xpu: if: false # github.repository == 'bitsandbytes-foundation/bitsandbytes' needs: build-cpu strategy: fail-fast: false matrix: torch_version: ["2.7.1"] #["2.6.0", "2.7.1"] runs-on: group: bandb-itac-bmsprpvc1550-8-1gpu env: BNB_TEST_DEVICE: xpu steps: - name: Show system information run: | echo "OS: $(uname -a)" echo "CPU: $(lscpu | grep 'Model name')" echo "Memory: $(free -h)" - name: Show XPU Information run: | xpu-smi discovery sudo xpu-smi discovery sudo apt-get install -y hwinfo hwinfo --display - uses: actions/checkout@v4 - name: Download build artifact uses: actions/download-artifact@v4 with: name: lib_cpu_ubuntu-22.04_x86_64 path: bitsandbytes/ merge-multiple: true - name: Setup Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Install PyTorch run: pip install torch==${{ matrix.torch_version }} --index-url https://download.pytorch.org/whl/xpu - name: Install dependencies run: | pip install -e ".[test]" pip install pytest-cov - name: Show installed packages run: pip list - name: Show environment information run: | python -m torch.utils.collect_env python -m bitsandbytes # - name: Run tests # run: pytest --durations=100 test-cuda: if: github.repository == 'bitsandbytes-foundation/bitsandbytes' needs: build-cuda strategy: fail-fast: false matrix: os: [ubuntu-22.04, windows-2025] arch: [x86_64] gpu: [T4, L40S] cuda_version: ["11.8.0", "12.6.3", "12.8.1", "12.9.1"] #, "13.0.1"] include: - cuda_version: "11.8.0" torch_version: "2.3.1" pypi_index: "https://download.pytorch.org/whl/cu118" - cuda_version: "12.6.3" torch_version: "2.6.0" pypi_index: "https://download.pytorch.org/whl/cu126" - cuda_version: "12.9.1" torch_version: "2.8.0" pypi_index: "https://download.pytorch.org/whl/cu129" - cuda_version: "12.8.1" torch_version: "2.9.0" pypi_index: "https://download.pytorch.org/whl/test/cu128" # Note: Currently our runners do not have new enough drivers for CUDA 13. # Add this when supported. # - cuda_version: "13.0.1" # torch_version: "2.9.0" # pypi_index: "https://download.pytorch.org/whl/test/cu130" # Linux L40S runners - os: ubuntu-22.04 gpu: L40S runner: bandb-aws-g6e-4xlarge-plus-use1-public-80 # Linux T4 runners - os: ubuntu-22.04 gpu: T4 runner: bandb-aws-g4dn-4xlarge-plus-use1-public-80 # Specific Windows runners using cu118 - os: windows-2025 arch: x86_64 gpu: T4 runner: CUDA-Windows-x64 cuda_version: "11.8.0" torch_version: "2.3.1" pypi_index: "https://download.pytorch.org/whl/cu118" - os: windows-2025 arch: x86_64 gpu: T4 runner: CUDA-Windows-x64 cuda_version: "11.8.0" torch_version: "2.6.0" pypi_index: "https://download.pytorch.org/whl/cu118" - os: windows-2025 arch: x86_64 gpu: T4 runner: CUDA-Windows-x64 cuda_version: "11.8.0" torch_version: "2.7.1" # Note: this is the last PyTorch release supporting CUDA 11.8. pypi_index: "https://download.pytorch.org/whl/cu118" exclude: # Our current T4 Windows runner has a driver too old (471.11) # and cannot support CUDA 12+. Skip for now. - os: windows-2025 cuda_version: "13.0.1" - os: windows-2025 cuda_version: "12.9.1" - os: windows-2025 cuda_version: "12.8.1" - os: windows-2025 cuda_version: "12.6.3" # No Windows L40S runners. - os: windows-2025 gpu: L40S runs-on: ${{ matrix.runner }} env: BNB_TEST_DEVICE: cuda steps: - name: Show GPU Information run: nvidia-smi - uses: actions/checkout@v4 - name: Download build artifact uses: actions/download-artifact@v4 with: name: lib_cuda_${{ matrix.cuda_version }}_${{ matrix.os }}_${{ matrix.arch }} path: bitsandbytes/ merge-multiple: true - name: Setup Python uses: actions/setup-python@v5 with: python-version: '3.10' - name: Install dependencies run: | pip install --pre torch~=${{ matrix.torch_version }}.dev0 --index-url ${{ matrix.pypi_index }} pip install -e ".[test]" pip install pytest-cov - name: Show installed packages run: pip list - name: Show environment information run: python -m torch.utils.collect_env - name: Run tests run: pytest --durations=100 bitsandbytes/.github/workflows/upload_pr_documentation.yml000066400000000000000000000007351506752665600246360ustar00rootroot00000000000000name: Upload PR Documentation on: workflow_run: workflows: ["Build PR Documentation"] types: - completed permissions: contents: read pull-requests: write # Allows posting comments on pull requests jobs: build: uses: huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml@main with: package_name: bitsandbytes secrets: hf_token: ${{ secrets.HUGGINGFACE_PUSH }} comment_bot_token: ${{ secrets.GITHUB_TOKEN }} bitsandbytes/.gitignore000066400000000000000000000040331506752665600155630ustar00rootroot00000000000000# Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class *.so *.dll *.dylib *.o *.obj *.air *.metallib # CMake generated files CMakeCache.txt CMakeScripts/ cmake_install.cmake Makefile CMakeFiles/ *.sln *.vcxproj* *.xcodeproj/ bitsandbytes.dir/ Debug/ Release/ cmake-build-*/ # IDE local files .vs/ .idea/ # Distribution / packaging .Python build/ develop-eggs/ dist/ downloads/ eggs/ .eggs/ lib/ lib64/ parts/ sdist/ var/ wheels/ pip-wheel-metadata/ share/python-wheels/ *.egg-info/ .installed.cfg *.egg MANIFEST # PyInstaller # Usually these files are written by a python script from a template # before PyInstaller builds the exe, so as to inject date/other infos into it. *.manifest *.spec # Installer logs pip-log.txt pip-delete-this-directory.txt # Unit test / coverage reports htmlcov/ .tox/ .nox/ .coverage .coverage.* .cache nosetests.xml coverage.xml *.cover *.py,cover .hypothesis/ .pytest_cache/ # Translations *.mo *.pot # Django stuff: *.log local_settings.py db.sqlite3 db.sqlite3-journal # Flask stuff: instance/ .webassets-cache # Scrapy stuff: .scrapy # Sphinx documentation docs/_build/ # PyBuilder target/ # Jupyter Notebook .ipynb_checkpoints # IPython profile_default/ ipython_config.py # pyenv .python-version # pipenv # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. # However, in case of collaboration, if having platform-specific dependencies or dependencies # having no cross-platform support, pipenv may install dependencies that don't work, or not # install all needed dependencies. #Pipfile.lock # PEP 582; used by e.g. github.com/David-OConnor/pyflow __pypackages__/ # Celery stuff celerybeat-schedule celerybeat.pid # SageMath parsed files *.sage.py # Environments .env .venv env/ venv/ ENV/ env.bak/ venv.bak/ # Spyder project settings .spyderproject .spyproject # Rope project settings .ropeproject # mkdocs documentation /site # mypy .mypy_cache/ .dmypy.json dmypy.json # Pyre type checker .pyre/ # vim *.swp dependencies cuda_build output/ bitsandbytes/.pre-commit-config.yaml000066400000000000000000000013151506752665600200540ustar00rootroot00000000000000repos: - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.11.2 hooks: - id: ruff args: - --fix - id: ruff-format - repo: https://github.com/pre-commit/pre-commit-hooks rev: v5.0.0 hooks: - id: check-merge-conflict - id: check-yaml - id: end-of-file-fixer - id: fix-byte-order-marker - id: trailing-whitespace - id: mixed-line-ending args: - --fix=lf - repo: https://github.com/crate-ci/typos rev: v1.26.0 hooks: - id: typos - repo: https://github.com/pre-commit/mirrors-clang-format rev: v20.1.6 hooks: - id: clang-format types_or: [c++, c, cuda] files: ^csrc/ bitsandbytes/.vscode/000077500000000000000000000000001506752665600151345ustar00rootroot00000000000000bitsandbytes/.vscode/extensions.json000066400000000000000000000001621506752665600202250ustar00rootroot00000000000000{ "recommendations": [ "ms-python.python", "charliermarsh.ruff", "twxs.cmake" ] } bitsandbytes/.vscode/settings.json000066400000000000000000000002061506752665600176650ustar00rootroot00000000000000{ "ruff.fixAll": true, "ruff.lint.run": "onType", "editor.codeActionsOnSave": { "source.fixAll": "always" } } bitsandbytes/CHANGELOG.md000066400000000000000000000637021506752665600154140ustar00rootroot00000000000000### v0.45.1 #### Improvements: * Compatibility for `triton>=3.2.0` * Moved package configuration to `pyproject.toml` * Build system: initial support for NVIDIA Blackwell B100 GPUs, RTX 50 Blackwell series GPUs and Jetson Thor Blackwell. * Note: Binaries built for these platforms are not included in this release. They will be included in future releases upon the availability of the upcoming CUDA Toolkit 12.7 and 12.8. #### Bug Fixes: * Packaging: wheels will no longer include unit tests. (#1478) #### Dependencies: * Sets the minimum PyTorch version to 2.0.0. ### 0.45.0 This is a significant release, bringing support for LLM.int8() to NVIDIA Hopper GPUs such as the H100. As part of the compatibility enhancements, we've rebuilt much of the LLM.int8() code in order to simplify for future compatibility and maintenance. We no longer use the col32 or architecture-specific tensor layout formats while maintaining backwards compatibility. We additionally bring performance improvements targeted for inference scenarios. #### Performance Improvements This release includes broad performance improvements for a wide variety of inference scenarios. See [this X thread](https://x.com/Tim_Dettmers/status/1864706051171287069) for a detailed explanation. #### Breaking Changes 🤗[PEFT](https://github.com/huggingface/peft) users wishing to merge adapters with 8-bit weights will need to upgrade to `peft>=0.14.0`. #### Packaging Improvements * The size of our wheel has been reduced by ~43.5% from 122.4 MB to 69.1 MB! This results in an on-disk size decrease from ~396MB to ~224MB. * Binaries built with CUDA Toolkit 12.6.2 are now included in the PyPI distribution. * The CUDA 12.5.0 build has been updated to CUDA Toolkit 12.5.1. #### Deprecations * A number of public API functions have been marked for deprecation and will emit `FutureWarning` when used. These functions will become unavailable in future releases. This should have minimal impact on most end-users. * The k-bit quantization features are deprecated in favor of blockwise quantization. For all optimizers, using `block_wise=False` is not recommended and support will be removed in a future release. * As part of the refactoring process, we've implemented many new 8bit operations. These operations no longer use specialized data layouts. #### Full Changelog * refine docs for multi-backend alpha release by @Titus-von-Koeller in #1380 * README: Replace special Unicode text symbols with regular characters by @akx in #1385 * Update CI tools & fix typos by @akx in #1386 * Fix invalid escape sequence warning in Python 3.12 by @oshiteku in #1420 * [Build] Add CUDA 12.6.2 build; update 12.5.0 to 12.5.1 by @matthewdouglas in #1431 * LLM.int8() Refactoring: Part 1 by @matthewdouglas in #1401 ### 0.44.1 #### Bug fixes: * Fix optimizer support for Python <= 3.9 by @matthewdouglas in #1379 ### 0.44.0 #### New: AdEMAMix Optimizer The [AdEMAMix](https://hf.co/papers/2409.03137) optimizer is a modification to AdamW which proposes tracking two EMAs to better leverage past gradients. This allows for faster convergence with less training data and improved resistance to forgetting. We've implemented 8bit and paged variations: `AdEMAMix`, `AdEMAMix8bit`, `PagedAdEMAMix`, and `PagedAdEMAMix8bit`. These can be used with a similar API to existing optimizers. #### Improvements: * **8-bit Optimizers**: The block size for all 8-bit optimizers has been reduced from 2048 to 256 in this release. This is a change from the original implementation proposed in [the paper](https://hf.co/papers/2110.02861) which improves accuracy. * **CUDA Graphs support**: A fix to enable [CUDA Graphs](https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/) capture of kernel functions was made in #1330. This allows for performance improvements with inference frameworks like vLLM. Thanks @jeejeelee! #### Full Changelog: * Embedding4bit and Embedding8bit implementation by @galqiwi in #1292 * Bugfix: Load correct nocublaslt library variant when BNB_CUDA_VERSION override is set by @matthewdouglas in #1318 * Enable certain CUDA kernels to accept specified cuda stream by @jeejeelee in #1330 * Initial support for ppc64le by @mgiessing in #1316 * Cuda source cleanup , refactor and fixes by @abhilash1910 in #1328 * Update for VS2022 17.11 compatibility with CUDA < 12.4 by @matthewdouglas in #1341 * Bump the minor-patch group with 3 updates by @dependabot in #1362 * Update matplotlib requirement from ~=3.9.1 to ~=3.9.2 in the major group by @dependabot in #1361 * docs: add internal reference to multi-backend guide by @Titus-von-Koeller in #1352 * Add move_to_device kwarg to the optimizer's load_state_dict by @koute in #1344 * Add AdEMAMix optimizer by @matthewdouglas in #1360 * Change 8bit optimizer blocksize 2048->256; additional bf16 support by @matthewdouglas in #1365 ### 0.43.3 #### Improvements: - FSDP: Enable loading prequantized weights with bf16/fp16/fp32 quant_storage - Background: This update, linked to [Transformer PR #32276](https://github.com/huggingface/transformers/pull/32276), allows loading prequantized weights with alternative storage formats. Metadata is tracked similarly to `Params4bit.__new__` post PR #970. It supports models exported with non-default `quant_storage`, such as [this NF4 model with BF16 storage](https://huggingface.co/hugging-quants/Meta-Llama-3.1-405B-BNB-NF4-BF16). - Special thanks to @winglian and @matthewdouglas for enabling FSDP+QLoRA finetuning of Llama 3.1 405B on a single 8xH100 or 8xA100 node with as little as 256GB system RAM. ### 0.43.2 This release is quite significant as the QLoRA bug fix big implications for higher `seqlen` and batch sizes. For each sequence (i.e. batch size increase of one) we expect memory savings of: - 405B: 39GB for `seqlen=1024`, and 4888GB for `seqlen=128,00` - 70B: 10.1GB for `seqlen=1024` and 1258GB for `seqlen=128,00` This was due to activations being unnecessary for frozen parameters, yet the memory for them was still erroneously allocated due to the now fixed bug. #### Improvements: - docs: FSDP+QLoRA and CPU install guide (#1211 #1227, thanks @stevhliu) - Add CUDA 12.5 and update 12.4 builds (#1284) #### Bug Fixes - 4bit getstate and 8bit deepcopy (#1230 #1231, thanks @BenjaminBossan) - missing optimizers in `str2optimizer32bit` (#1222, thanks @EtienneDosSantos) - CUDA 12.5 build issue (#1273, thanks @HennerM) - fix for min_8bit_size functionality in Optimizer base classes (#1286, thanks @Edenzzzz) - QLoRA mem bug (#1270, thanks @Ther-nullptr) - tests for cpu only platforms (#1259, thanks @galqiwi) - restoration of quant_storage for CPU offloading (#1279) - optim update error with non-contiguous grads/params (deepspeed) (#1187) ### 0.43.1 #### Improvements: - Improved the serialization format for 8-bit weights; this change is fully backwards compatible. (#1164, thanks to @younesbelkada for the contributions and @akx for the review). - Added CUDA 12.4 support to the Linux x86-64 build workflow, expanding the library's compatibility with the latest CUDA versions. (#1171, kudos to @matthewdouglas for this addition). - Docs enhancement: Improved the instructions for installing the library from source. (#1149, special thanks to @stevhliu for the enhancements). #### Bug Fixes - Fix 4bit quantization with blocksize = 4096, where an illegal memory access was encountered. (#1160, thanks @matthewdouglas for fixing and @YLGH for reporting) #### Internal Improvements: - Tests: improve memory usage (#1147, thanks @matthewdouglas) - Add CUDA 12.4 to docs/install helper (#1136, thanks @matthewdouglas) - Minor type/doc fixes (#1128, thanks @akx) - Reformat Python code with Ruff (#1081, thanks @akx) - Rework of CUDA/native-library setup and diagnostics (#1041, thanks @akx) ### 0.43.0 #### Improvements and New Features: - QLoRA + FSDP official support is now live! https://github.com/TimDettmers/bitsandbytes/pull/970 by @warner-benjamin and team - with FSDP you can train very large models (70b scale) on multiple 24GB consumer-type GPUs. See https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html for more details. - Introduced improvements to the CI process for enhanced performance and efficiency during builds, specifically enabling more effective cross-compilation on Linux platforms. This was accomplished by deprecating Make and migrating to Cmake, as well as implementing new corresponding workflows. Huge thanks go to @wkpark, @rickardp, @matthewdouglas and @younesbelkada; #1055, #1050, #1111. - Windows should be officially supported in bitsandbytes if you install the library from source. See: https://huggingface.co/docs/bitsandbytes/main/en/index for more details - Updated installation instructions to provide more comprehensive guidance for users. This includes clearer explanations and additional tips for various setup scenarios, making the library more accessible to a broader audience (@rickardp, #1047). - Enhanced the library's compatibility and setup process, including fixes for CPU-only installations and improvements in CUDA setup error messaging. This effort aims to streamline the installation process and improve user experience across different platforms and setups (@wkpark, @akx, #1038, #996, #1012). - Setup a new documentation at https://huggingface.co/docs/bitsandbytes/main with extensive new sections and content to help users better understand and utilize the library. Especially notable are the new API docs. (big thanks to @stevhliu and @mishig25 from HuggingFace #1012). The API docs have been also addressed in #1075. #### Bug Fixes: - Addressed a race condition in kEstimateQuantiles, enhancing the reliability of quantile estimation in concurrent environments (@pnunna93, #1061). - Fixed various minor issues, including typos in code comments and documentation, to improve code clarity and prevent potential confusion (@Brian Vaughan, #1063). #### Backwards Compatibility - After upgrading from `v0.42` to `v0.43`, when using 4bit quantization, models may generate slightly different outputs (approximately up to the 2nd decimal place) due to a fix in the code. For anyone interested in the details, [see this comment](https://github.com/TimDettmers/bitsandbytes/discussions/1094#discussioncomment-8984069). #### Internal and Build System Enhancements: - Implemented several enhancements to the internal and build systems, including adjustments to the CI workflows, portability improvements, and build artifact management. These changes contribute to a more robust and flexible development process, ensuring the library's ongoing quality and maintainability (@rickardp, @akx, @wkpark, @matthewdouglas; #949, #1053, #1045, #1037). #### Contributors: This release is made possible thanks to the many active contributors that submitted PRs and many others who contributed to discussions, reviews, and testing. Your efforts greatly enhance the library's quality and user experience. It's truly inspiring to work with such a dedicated and competent group of volunteers and professionals! We give a special thanks to @TimDettmers for managing to find a little bit of time for valuable consultations on critical topics, despite preparing for and touring the states applying for professor positions. We wish him the utmost success! We also extend our gratitude to the broader community for your continued support, feedback, and engagement, which play a crucial role in driving the library's development forward. ### 0.42.0 Features: - 4-bit serialization now supported. This enables 4-bit load/store. Thank you @poedator #753 - the bitsandbytes library now has a version attribute: `bitsandbytes.__version__` @rasbt #710 Bug fixes: - Fixed bugs in dynamic exponent data type creation. Thank you @RossM, @KohakuBlueleaf, @ArrowM #659 #227 #262 #152 - Fixed an issue where 4-bit serialization would fail for layers without double quantization #868. Thank you, @poedator - Fixed an issue where calling .to() or .cuda() on a 4-bit layer twice would result in an error #867. Thank you, @jph00 - Fixed a bug where a missing access permission in a path searched for CUDA would lead to an error @osma #677 - Fixed a bug where the GOOGLE_VM_CONFIG_LOCK_FILE variable could cause errors in colab environments @akrentsel @xaptronic #715 #883 #622 - Fixed a bug where kgetColRowStats (LLM.int8()) would fail for certain dimensions @LucQueen @905 - Fixed a bug where the adjusted regular Embedding layer was not available via bnb.nn.Embedding @neel04 #563 - Fixed added missing scipy requirement @dulalbert #525 ### 0.41.3 Bug fixes: - Fixed an issue where 4-bit serialization would fail for layers without double quantization #868. Thank you, @poedator - Fixed an issue where calling .to() or .cuda() on a 4-bit layer twice would result in an error #867. Thank you, @jph00 ### 0.41.2 Feature: - 4-bit serialization now supported. This enables 4-bit load/store. Thank you @poedator #753 ### 0.41.1 Bug fixes: - Fixed bugs in dynamic exponent data type creation. Thank you @RossM, @KohakuBlueleaf, @ArrowM #659 #227 #262 #152 ### 0.41.0 Features: - Added precompiled CUDA 11.8 binaries to support H100 GPUs without compilation #571 - CUDA SETUP now no longer looks for libcuda and libcudart and relies PyTorch CUDA libraries. To manually override this behavior see: how_to_use_nonpytorch_cuda.md. Thank you @rapsealk Bug fixes: - Fixed a bug where the default type of absmax was undefined which leads to errors if the default type is different than torch.float32. # 553 - Fixed a missing scipy dependency in requirements.txt. #544 - Fixed a bug, where a view operation could cause an error in 8-bit layers. - Fixed a bug where CPU bitsandbytes would during the import. #593 Thank you @bilelomrani - Fixed a but where a non-existent LD_LIBRARY_PATH variable led to a failure in python -m bitsandbytes #588 - Removed outdated get_cuda_lib_handle calls that lead to errors. #595 Thank you @ihsanturk - Fixed bug where read-permission was assumed for a file. #497 - Fixed a bug where prefetchAsync lead to errors on GPUs that do not support unified memory but not prefetching (Maxwell, SM52). #470 #451 #453 #477 Thank you @jllllll and @stoperro Documentation: - Improved documentation for GPUs that do not support 8-bit matmul. #529 - Added description and pointers for the NF4 data type. #543 User experience: - Improved handling of default compute_dtype for Linear4bit Layers, so that compute_dtype = input_dtype if the input data type is stable enough (float32, bfloat16, but not float16). Performance: - improved 4-bit inference performance for A100 GPUs. This degraded performance for A40/RTX3090 and RTX 4090 GPUs slightly. ### 0.40.2 Bug fixes: - Fixed a but where a non-existent LD_LIBRARY_PATH variable led to a failure in python -m bitsandbytes #588 - Removed outdated get_cuda_lib_handle calls that lead to errors. #595 Thank you @ihsanturk - Fixed bug where read-permission was assumed for a file. #497 - Fixed a bug where prefetchAsync lead to errors on GPUs that do not support unified memory but not prefetching (Maxwell, SM52). #470 #451 #453 #477 Thank you @jllllll and @stoperro ### 0.40.1 Features: - Added precompiled CUDA 11.8 binaries to support H100 GPUs without compilation #571 - CUDA SETUP now no longer looks for libcuda and libcudart and relies PyTorch CUDA libraries. To manually override this behavior see: how_to_use_nonpytorch_cuda.md. Thank you @rapsealk Bug fixes: - Fixed a bug where the default type of absmax was undefined which leads to errors if the default type is different than torch.float32. # 553 - Fixed a missing scipy dependency in requirements.txt. #544 - Fixed a bug, where a view operation could cause an error in 8-bit layers. - Fixed a bug where CPU bitsandbytes would during the import. #593 Thank you @bilelomrani Documentation: - Improved documentation for GPUs that do not support 8-bit matmul. #529 - Added description and pointers for the NF4 data type. #543 ### 0.40.0 Features: - Added 4-bit inference kernels for batch size=1. Currently support are the NF4, FP4 data types. - Added support for quantizations of bfloat16 input data. Bug fixes: - Added `device` variable for bitsandbytes layers to be compatible with PyTorch layers. Deprecated: - Binaries for CUDA 11.2, 11.6 no longer ship with `pip install bitsandbytes` and need to be compiled from source. ### 0.39.0 Features: - 4-bit matrix multiplication for Float4 and NormalFloat4 data types. - Added 4-bit quantization routines - Doubled quantization routines for 4-bit quantization - Paged optimizers for Adam and Lion. - bfloat16 gradient / weight support for Adam and Lion with 8 or 32-bit states. Bug fixes: - Fixed a bug where 8-bit models consumed twice the memory as expected after serialization Deprecated: - Kepler binaries (GTX 700s and Tesla K40/K80) are not longer provided via pip and need to be compiled from source. Kepler support might be fully removed in the future. ### 0.38.1 Features: - Added Int8 SwitchBack layers - Added Fake FP8 layers for research purposes (available under `bnb.research.nn. ...`) ### 0.38.0 #### 8-bit Lion, Load/Store 8-bit Models directly from/to HF Hub Features: - Support for 32 and 8-bit Lion has been added. Thank you @lucidrains - Support for serialization of Linear8bitLt layers (LLM.int8()). This allows to store and load 8-bit weights directly from the HuggingFace Hub. Thank you @myrab - New bug report features `python -m bitsandbytes` now gives extensive debugging details to debug CUDA setup failures. Bug fixes: - Fixed a bug where some bitsandbytes methods failed in a model-parallel setup on multiple GPUs. Thank you @tonylins - Fixed a bug where cudart.so libraries could not be found in newer PyTorch releases. Improvements: - Improved the CUDA Setup procedure by doing a more extensive search for CUDA libraries Deprecated: - Devices with compute capability 3.0 (GTX 700s, K10) and 3.2 (Tegra K1, Jetson TK1) are now deprecated and support will be removed in 0.39.0. - Support for CUDA 10.0 and 10.2 will be removed in bitsandbytes 0.39.0 ### 0.37.0 #### Int8 Matmul + backward support for all GPUs Features: - Int8 MatmulLt now supports backward through inversion of the ColTuring/ColAmpere format. Slow, but memory efficient. Big thanks to @borzunov - Int8 now supported on all GPUs. On devices with compute capability \< 7.5, the Int weights are cast to 16/32-bit for the matrix multiplication. Contributed by @borzunov Improvements: - Improved logging for the CUDA detection mechanism. ### 0.36.0 #### Improvements, Ada/Hopper support, fake k-bit quantization. Features: - CUDA 11.8 and 12.0 support added - support for Ada and Hopper GPUs added (compute capability 8.9 and 9.0) - support for fake k-bit block-wise quantization for Int, Float, quantile quantization, and dynamic exponent data types added - Added CUDA instruction generator to fix some installations. - Added additional block sizes for quantization {64, 128, 256, 512, 1024} - Added SRAM Quantile algorithm to quickly estimate less than 256 quantiles - Added option to suppress the bitsandbytes welcome message (@Cyberes) Regression: - Compute capability 3.0 removed: GTX 600s and 700s series is no longer supported (except GTX 780 and GTX 780 Ti) Bug fixes: - fixed a bug where too long directory names would crash the CUDA SETUP #35 (@tomaarsen) - fixed a bug where CPU installations on Colab would run into an error #34 (@tomaarsen) - fixed an issue where the default CUDA version with fast-DreamBooth was not supported #52 - fixed a bug where the CUDA setup failed due to a wrong function call. - fixed a bug in the CUDA Setup which led to an incomprehensible error if no GPU was detected. - fixed a bug in the CUDA Setup failed with the cuda runtime was found, but not the cuda library. - fixed a bug where not finding the cuda runtime led to an incomprehensible error. - fixed a bug where with missing CUDA the default was an error instead of the loading the CPU library - fixed a bug where the CC version of the GPU was not detected appropriately (@BlackHC) - fixed a bug in CPU quantization which lead to errors when the input buffer exceeded 2^31 elements Improvements: - multiple improvements in formatting, removal of unused imports, and slight performance improvements (@tomaarsen) - StableEmbedding layer now has device and dtype parameters to make it 1:1 replaceable with regular Embedding layers (@lostmsu) - runtime performance of block-wise quantization slightly improved - added error message for the case multiple libcudart.so are installed and bitsandbytes picks the wrong one ### 0.35.4 Bug fixes: - Fixed a bug in the CUDA Setup failed with the cuda runtime was found, but not the cuda library. - Fixed a bug where not finding the cuda runtime led to an incomprehensible error. ### 0.35.3 Bug fixes: - Fixed a bug in the CUDA Setup which led to an incomprehensible error if no GPU was detected. ### 0.35.2 Bug fixes: - Fixed a bug where the CUDA setup failed due to a wrong function call. ### 0.35.1 Features: - Added CUDA instruction generator to fix some installations. Bug fixes: - Fixed a problem where warning messages would be displayed even though everything worked correctly. ### 0.35.0 #### CUDA 11.8 support and bug fixes Features: - CUDA 11.8 support added and binaries added to the PyPI release. Bug fixes: - fixed a bug where too long directory names would crash the CUDA SETUP #35 (thank you @tomaarsen) - fixed a bug where CPU installations on Colab would run into an error #34 (thank you @tomaarsen) - fixed an issue where the default CUDA version with fast-DreamBooth was not supported #52 ### 0.34.0 #### Bug fixes and memory efficient backprop Features: - Linear8bitLt layer now supports `memory_efficient_backward=True` which enables backprop of gradients through frozen weights. Bug fixes: - fixed an issue where too many threads were created in blockwise quantization on the CPU for large tensors ### 0.33.0 #### Various bug fixes Features: - CPU quantization now supports a variable `blocksize` variable to enhance quantization speed or precision. Bug fixes: - fixed an issue in CPU quantization where tensors with more than 2^31 elements would fail 19a7adca7a6c9bf7061a384d7e9d9b13676a1a88 - fixed a bug where cpu binaries would fail if no GPU would be detected eab4d8232d558f2e6bd7f7cc3d00e2e6e94f4e80 - fixed an issue where cpu binaries cause additional stdout messages 92a3363096e10ad6a5c4e944af898bd1186d806a - fixed an import of bnb.utils 2e630b55f51d454f3bd723dffda68a07ef93190c We thank @mryab, @mbrukman, @chessgecko, @dbaranchuk for pull request with bug fixes and new features. ### 0.32.0 #### 8-bit Inference Performance Enhancements We added performance enhancements for small models. This makes small models about 2x faster for LLM.int8() inference. Features: - Int32 dequantization now supports fused biases. - Linear8bitLt now uses a fused bias implementation. - Change `.data.storage().data_ptr()` to `.data.data_ptr()` to enhance inference performance. Bug fixes: - Now throws and error if LLM.int8() is used on a GPU that is not supported. - Enhances error messaging if CUDA SETUP fails. ### 0.31.0 #### 8-bit Inference and Packaging Update Features: - added direct outlier extraction. This enables outlier extraction without fp16 weights without performance degradation. - Added automatic CUDA SETUP procedure and packaging all binaries into a single bitsandbytes package. ### 0.30.0 #### 8-bit Inference Update Features: - Added 8-bit matrix multiplication form cuBLAS, and cuBLASLt as well as multiple GEMM kernels (GEMM, GEMMEx, GEMMLt) - Added 8-bit Linear layers with 8-bit Params that perform memory efficient inference with an option for 8-bit mixed precision matrix decomposition for inference without performance degradation - Added quantization methods for "fake" quantization as well as optimized kernels vector-wise quantization and equalization as well as optimized cuBLASLt transformations - CPU only build now available (Thank you, @mryab) Deprecated: - Pre-compiled release for CUDA 9.2, 10.0, 10.2 no longer available ### 0.26.0: Features: - Added Adagrad (without grad clipping) as 32-bit and 8-bit block-wise optimizer. - Added AdamW (copy of Adam with weight decay init 1e-2). #10 - Introduced ModuleConfig overrides which can be seamlessly be used at initialization time of a module. - Added `bnb.nn.Embedding` layer which runs at 32-bit but without the layernorm. This works well if you need to fine-tune pretrained models that do not have a embedding layer norm. #19 Bug fixes: - Fixed a bug where weight decay was incorrectly applied to 32-bit Adam. #13 - Fixed an unsafe use of eval. #8 - Fixed a bug where the StableEmbedding layer 32-bit optimizer override would not work without registering the whole model first (`bnb.optim.GlobalOptimManager.get_instance().register_parameters(model.parameters())`). #13 #15 Docs: - Added instructions how to solve "\_\_fatbinwrap\_" errors. ### 0.0.25: Features: - Added `skip_zeros` for block-wise and 32-bit optimizers. This ensures correct updates for sparse gradients and sparse models. - Added support for Kepler GPUs. (#4) - Added Analysis Adam to track 8-bit vs 32-bit quantization errors over time. - Make compilation more user friendly. Bug fixes: - fixed "undefined symbol: \_\_fatbinwrap_38" error for P100 GPUs on CUDA 10.1 (#5) Docs: - Added docs with instructions to compile from source. ### 0.0.24: - Fixed a bug where a float/half conversion led to a compilation error for CUDA 11.1 on Turning GPUs. - removed Apex dependency for bnb LAMB ### 0.0.23: Bugs: - Unified quantization API: each quantization function now returns `Q, S` where `Q` is the quantized tensor and `S` the quantization state which may hold absolute max values, a quantization map or more. For dequantization all functions now accept the inputs `Q, S` so that `Q` is dequantized with the quantization state `S`. - Fixed an issue where the CUDA 11.1 binary was not compiled with the right headers API changes: - Block-wise quantization for optimizers now enabled by default Features: - Block-wise quantization routines now support CPU Tensors. ### 0.0.22: - Fixed an error where a `reset_parameters()` call on the `StableEmbedding` would lead to an error in older PyTorch versions (from 1.7.0). ### 0.0.21 - Ampere, RTX 30 series GPUs now compatible with the library. bitsandbytes/CMakeLists.txt000066400000000000000000000340671506752665600163450ustar00rootroot00000000000000# This CMake config hopefully makes it easier to compile. # Ensure the CUDA Toolkit is available on your path. Then run: # For GCC: `cmake -B build . && cmake --build build` # For MSVC: `cmake -B build . && cmake --build build --config Release` # You can also use the following options and variables # - COMPUTE_BACKEND: Set to `cpu`, `cuda`, or `mps` to select the backend # - CUDA_VERSION: The expected CUDA version, for sanity checking. The actual version # is whatever CMake finds on your path. # - COMPUTE_CAPABILITY: Which GPU Arch/Compute codes to provide to NVCC. # Separate by semicolons, i.e. `-DCOMPUTE_CAPABILITY=89;90;100;120` # Check your compute capability here: https://developer.nvidia.com/cuda-gpus # - PTXAS_VERBOSE: Pass the `-v` option to the PTX Assembler cmake_minimum_required(VERSION 3.22.1) project(bitsandbytes LANGUAGES CXX) # If run without specifying a build type, default to using the Release configuration: # optimizing the generated binaries for performance and also adds the `-DNDEBUG` flag, # which turns off a bunch of asserts which seem to link to new symbols in libstdc++, # worsening our many_linux compliance.. if(NOT CMAKE_BUILD_TYPE) set(CMAKE_BUILD_TYPE Release) endif() # Define included source files set(CPP_FILES csrc/common.cpp csrc/cpu_ops.cpp csrc/pythonInterface.cpp) set(CUDA_FILES csrc/ops.cu csrc/kernels.cu) set(HIP_FILES csrc/ops.hip csrc/kernels.hip) set(MPS_FILES csrc/mps_ops.mm) set(METAL_FILES csrc/mps_kernels.metal) set(XPU_FILES csrc/xpu_ops.cpp csrc/xpu_kernels.cpp) # C++ sources are always included list(APPEND SRC_FILES ${CPP_FILES}) set(COMPUTE_BACKEND "cpu" CACHE STRING "The compute backend to use (cpu, cuda, hip, mps, xpu)") set_property(CACHE COMPUTE_BACKEND PROPERTY STRINGS cpu cuda hip mps xpu) option(PTXAS_VERBOSE "Pass through -v flag to PTX Assembler" OFF) if(APPLE) set(CMAKE_OSX_DEPLOYMENT_TARGET 13.1) endif() set(BNB_OUTPUT_NAME "bitsandbytes") message(STATUS "Configuring ${PROJECT_NAME} (Backend: ${COMPUTE_BACKEND})") if(${COMPUTE_BACKEND} STREQUAL "cuda") if(APPLE) message(FATAL_ERROR "CUDA is not supported on macOS" ) endif() set(BUILD_CUDA ON) set(BUILD_HIP OFF) set(BUILD_MPS OFF) elseif(${COMPUTE_BACKEND} STREQUAL "hip") if(APPLE) message(FATAL_ERROR "HIP is not supported on macOS" ) endif() set(BUILD_CUDA OFF) set(BUILD_HIP ON) set(BUILD_MPS OFF) elseif(${COMPUTE_BACKEND} STREQUAL "mps") if(NOT APPLE) message(FATAL_ERROR "MPS is only supported on macOS" ) endif() set(BUILD_CUDA OFF) set(BUILD_HIP OFF) set(BUILD_MPS ON) elseif(${COMPUTE_BACKEND} STREQUAL "xpu") if(APPLE) message(FATAL_ERROR "XPU is not supported on macOS" ) endif() set(BUILD_CUDA OFF) set(BUILD_HIP OFF) set(BUILD_MPS OFF) set(BUILD_XPU ON) else() set(BUILD_CUDA OFF) set(BUILD_HIP OFF) set(BUILD_MPS OFF) set(BUILD_XPU OFF) endif() if(BUILD_CUDA) # NVCC normally will only work with MSVC up to 1939. VS2022 17.10+ starts using versions 1940+. # Workaround: use --allow-unsupported-compiler # This needs to be added *before* we try to enable the CUDA language so CMake's compiler check passes. if(MSVC AND MSVC_VERSION VERSION_GREATER_EQUAL 1940) string(APPEND CMAKE_CUDA_FLAGS " --allow-unsupported-compiler") # This is needed to build with VS2022 17.11+ and CUDA < 12.4. if (MSVC_VERSION VERSION_GREATER_EQUAL 1941) string(APPEND CMAKE_CUDA_FLAGS " -D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") endif() endif() enable_language(CUDA) # This will fail if CUDA is not found find_package(CUDAToolkit REQUIRED) # Convert the CUDA version from X.Y.z to XY. There's probably a shorter way of doing this string(REGEX MATCH "^[0-9]+.[0-9]+" _CUDA_VERSION_FIRST_TWO "${CMAKE_CUDA_COMPILER_VERSION}") string(REPLACE "." "" CUDA_VERSION_SHORT "${_CUDA_VERSION_FIRST_TWO}") # Expose a cache variable that the user can set to ensure the correct version of CUDA is found set(CUDA_VERSION "${CUDA_VERSION_SHORT}" CACHE STRING "Expected CUDA Version Shortcode") message(STATUS "CUDA Version: ${CUDA_VERSION_SHORT} (${CMAKE_CUDA_COMPILER_VERSION})") message(STATUS "CUDA Compiler: ${CMAKE_CUDA_COMPILER}") # It should match the discovered version if(NOT CUDA_VERSION STREQUAL "${CUDA_VERSION_SHORT}") message(FATAL_ERROR "You've specified CUDA version ${CUDA_VERSION} however the CUDA compiler found is ${CUDA_VERSION_SHORT}." " Ensure the desired CUDA compiler is the first one available on your PATH." ) endif() if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS "11.8") message(FATAL_ERROR "CUDA Version < 11.8 is not supported") elseif(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "14.0") message(FATAL_ERROR "CUDA Version > 13 is not supported") endif() # CMake < 3.23.0 does not define CMAKE_CUDA_ARCHITECTURES_ALL. if(CMAKE_VERSION VERSION_LESS "3.23.0") message(STATUS "CMake < 3.23.0; determining CUDA architectures supported...") if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "13.0") # Starting in CUDA 13.0, Thor Blackwell is renamed to SM110. # Support for architectures older than Turing (SM75) is removed. list(APPEND CMAKE_CUDA_ARCHITECTURES_ALL 75 80 86 87 88 89 90 100 103 110 120 121) list(APPEND CMAKE_CUDA_ARCHITECTURES_ALL_MAJOR 80 90 100 110 120) else() # 11.8-12.9 supports these at a minimum. set(CMAKE_CUDA_ARCHITECTURES_ALL 50 52 53 60 61 62 70 72 75 80 86 87 89 90) set(CMAKE_CUDA_ARCHITECTURES_ALL_MAJOR 50 60 70 80 90) # CUDA 12.8 adds support for Blackwell. if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "12.8") list(APPEND CMAKE_CUDA_ARCHITECTURES_ALL 100 101 120 121) list(APPEND CMAKE_CUDA_ARCHITECTURES_ALL_MAJOR 100 120) endif() # CUDA 12.9 adds SM103 (Blackwell B300). if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "12.9") list(APPEND CMAKE_CUDA_ARCHITECTURES_ALL 103) endif() endif() endif() string(APPEND CMAKE_CUDA_FLAGS " --use_fast_math") if(PTXAS_VERBOSE) # Verbose? Outputs register usage information, and other things... string(APPEND CMAKE_CUDA_FLAGS " -Xptxas=-v") endif() foreach(capability ${CMAKE_CUDA_ARCHITECTURES_ALL}) # Most of the items here are like: `xx-real`, so we just extract the `xx` portion string(REGEX MATCH "[0-9]+" capability_id "${capability}") if(capability_id GREATER 0) list(APPEND POSSIBLE_CAPABILITIES ${capability_id}) endif() endforeach() # This can be changed via -D argument to CMake # By default all possible capabilities are compiled set(COMPUTE_CAPABILITY "${POSSIBLE_CAPABILITIES}" CACHE STRING "Compute Capabilities Targeted") message(STATUS "CUDA Capabilities Available: ${POSSIBLE_CAPABILITIES}") message(STATUS "CUDA Capabilities Selected: ${COMPUTE_CAPABILITY}") # Use the "real" option to build native cubin for all selections. # Ensure we build the PTX for the latest version. # This behavior of adding a PTX (virtual) target for the highest architecture # is similar to how the "all" and "all-major" options would behave in CMake >= 3.23. # TODO: Consider bumping CMake requirement and using CMAKE_CUDA_ARCHITECTURES=[all | native] by default list(REMOVE_DUPLICATES COMPUTE_CAPABILITY) list(SORT COMPUTE_CAPABILITY COMPARE NATURAL) list(POP_BACK COMPUTE_CAPABILITY _LATEST_CAPABILITY) list(TRANSFORM COMPUTE_CAPABILITY APPEND "-real" OUTPUT_VARIABLE CMAKE_CUDA_ARCHITECTURES) list(APPEND CMAKE_CUDA_ARCHITECTURES ${_LATEST_CAPABILITY}) message(STATUS "CUDA Targets: ${CMAKE_CUDA_ARCHITECTURES}") message(STATUS "CUDA NVCC Flags: ${CMAKE_CUDA_FLAGS}") list(APPEND SRC_FILES ${CUDA_FILES}) string(APPEND BNB_OUTPUT_NAME "_cuda${CUDA_VERSION_SHORT}") add_compile_definitions(BUILD_CUDA) elseif(BUILD_HIP) enable_language(HIP) message(STATUS "HIP Compiler: ${CMAKE_HIP_COMPILER}") if(DEFINED BNB_ROCM_ARCH) set(CMAKE_HIP_ARCHITECTURES ${BNB_ROCM_ARCH}) else() if (NOT AMDGPU_TARGETS AND NOT CMAKE_HIP_ARCHITECTURES) set(CMAKE_HIP_ARCHITECTURES "gfx90a;gfx942;gfx1100") elseif (AMDGPU_TARGETS AND NOT CMAKE_HIP_ARCHITECTURES) set(CMAKE_HIP_ARCHITECTURES ${AMDGPU_TARGETS}) endif() endif() message(STATUS "HIP Targets: ${CMAKE_HIP_ARCHITECTURES}") list(APPEND SRC_FILES ${HIP_FILES}) string(APPEND BNB_OUTPUT_NAME "_rocm") # get hip version execute_process(COMMAND hipconfig --version OUTPUT_VARIABLE HIP_CONFIG_VERSION) string(REGEX MATCH "[0-9]+\\.[0-9]+" HIP_VERSION "${HIP_CONFIG_VERSION}") string(REPLACE "." "" HIP_VERSION_SHORT "${HIP_VERSION}") string(APPEND BNB_OUTPUT_NAME "${HIP_VERSION_SHORT}") add_compile_definitions(__HIP_PLATFORM_AMD__) add_compile_definitions(__HIP_PLATFORM_HCC__) add_compile_definitions(BUILD_HIP) elseif(BUILD_MPS) if(NOT APPLE) message(FATAL_ERROR "MPS is only supported on macOS" ) endif() enable_language(OBJCXX) list(APPEND SRC_FILES ${MPS_FILES}) string(APPEND BNB_OUTPUT_NAME "_mps") add_compile_definitions(BUILD_MPS) file(MAKE_DIRECTORY "build") add_custom_command(OUTPUT "bitsandbytes/bitsandbytes.metallib" COMMAND xcrun metal -c -o "build/bitsandbytes.air" ${METAL_FILES} COMMAND xcrun metallib "build/bitsandbytes.air" -o "bitsandbytes/bitsandbytes.metallib" DEPENDS "${METAL_FILES}" COMMENT "Compiling Metal kernels" VERBATIM) add_custom_target(metallib DEPENDS "bitsandbytes/bitsandbytes.metallib") elseif(BUILD_XPU) list(APPEND SRC_FILES ${XPU_FILES}) string(APPEND BNB_OUTPUT_NAME "_xpu") add_compile_definitions(BUILD_XPU) set(CMAKE_C_COMPILER icx) set(CMAKE_CXX_COMPILER icpx) if(WIN32) set(CMAKE_CXX_COMPILER icx) endif() else() string(APPEND BNB_OUTPUT_NAME "_cpu") set(GPU_SOURCES) endif() if(WIN32) # Export all symbols set(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS ON) endif() if(MSVC) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /arch:AVX2 /fp:fast") endif() set_source_files_properties(${CPP_FILES} PROPERTIES LANGUAGE CXX) add_library(bitsandbytes SHARED ${SRC_FILES}) target_compile_features(bitsandbytes PUBLIC cxx_std_17) target_include_directories(bitsandbytes PUBLIC csrc include) if(BUILD_CUDA) target_include_directories(bitsandbytes PUBLIC ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) target_link_libraries(bitsandbytes PUBLIC CUDA::cudart CUDA::cublas CUDA::cublasLt CUDA::cusparse) set_target_properties(bitsandbytes PROPERTIES CUDA_SEPARABLE_COMPILATION ON ) endif() if(BUILD_HIP) if(NOT DEFINED ENV{ROCM_PATH}) set(ROCM_PATH /opt/rocm) else() set(ROCM_PATH $ENV{ROCM_PATH}) endif() list(APPEND CMAKE_PREFIX_PATH ${ROCM_PATH}) macro(find_package_and_print_version PACKAGE_NAME) find_package("${PACKAGE_NAME}" ${ARGN}) message("${PACKAGE_NAME} VERSION: ${${PACKAGE_NAME}_VERSION}") endmacro() find_package_and_print_version(hipblas REQUIRED) find_package_and_print_version(hiprand REQUIRED) find_package_and_print_version(hipsparse REQUIRED) ## hacky way of excluding hip::amdhip64 (with it linked many tests unexpectedly fail e.g. adam8bit because of inaccuracies) set_target_properties(hip::host PROPERTIES INTERFACE_LINK_LIBRARIES "") set_target_properties(hip-lang::host PROPERTIES INTERFACE_LINK_LIBRARIES "") set(CMAKE_HIP_IMPLICIT_LINK_LIBRARIES "") target_include_directories(bitsandbytes PRIVATE ${CMAKE_SOURCE_DIR} ${CMAKE_SOURCE_DIR}/include ${ROCM_PATH}/include /include) target_link_directories(bitsandbytes PRIVATE ${ROCM_PATH}/lib /lib) target_link_libraries(bitsandbytes PUBLIC roc::hipblas hip::hiprand roc::hipsparse) target_compile_definitions(bitsandbytes PUBLIC BNB_USE_HIP) set_source_files_properties(${HIP_FILES} PROPERTIES LANGUAGE HIP) set_target_properties(bitsandbytes PROPERTIES LINKER_LANGUAGE CXX) if(HIP_VERSION VERSION_LESS "6.1") target_compile_definitions(bitsandbytes PUBLIC NO_HIPBLASLT) else() find_package(hipblaslt) target_link_libraries(bitsandbytes PUBLIC roc::hipblaslt) endif() endif() if(BUILD_MPS) add_dependencies(bitsandbytes metallib) target_link_libraries(bitsandbytes objc "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" "-framework MetalPerformanceShadersGraph") endif() if(BUILD_XPU) set(SYCL_LINK_FLAGS "-fsycl;--offload-compress;-fsycl-targets=spir64_gen,spir64;-Xs;-device pvc,xe-lpg,ats-m150 -options ' -cl-intel-enable-auto-large-GRF-mode -cl-poison-unsupported-fp64-kernels -cl-intel-greater-than-4GB-buffer-required'") set(SYCL_COMPILE_FLAGS "-fsycl;-fhonor-nans;-fhonor-infinities;-fno-associative-math;-fno-approx-func;-fno-sycl-instrument-device-code;--offload-compress;-fsycl-targets=spir64_gen,spir64;") set_property(TARGET bitsandbytes PROPERTY CXX_STANDARD 20) target_compile_options(bitsandbytes PRIVATE ${SYCL_COMPILE_FLAGS}) target_link_options(bitsandbytes PRIVATE ${SYCL_LINK_FLAGS}) endif() if(WIN32) set_target_properties(bitsandbytes PROPERTIES PREFIX "lib") endif() set_target_properties(bitsandbytes PROPERTIES OUTPUT_NAME ${BNB_OUTPUT_NAME}) if(MSVC) set_target_properties(bitsandbytes PROPERTIES LIBRARY_OUTPUT_DIRECTORY_RELEASE "${PROJECT_SOURCE_DIR}/bitsandbytes") set_target_properties(bitsandbytes PROPERTIES LIBRARY_OUTPUT_DIRECTORY_DEBUG "${PROJECT_SOURCE_DIR}/bitsandbytes") set_target_properties(bitsandbytes PROPERTIES RUNTIME_OUTPUT_DIRECTORY_RELEASE "${PROJECT_SOURCE_DIR}/bitsandbytes") set_target_properties(bitsandbytes PROPERTIES RUNTIME_OUTPUT_DIRECTORY_DEBUG "${PROJECT_SOURCE_DIR}/bitsandbytes") endif() set_target_properties(bitsandbytes PROPERTIES LIBRARY_OUTPUT_DIRECTORY "${PROJECT_SOURCE_DIR}/bitsandbytes") bitsandbytes/CODE_OF_CONDUCT.md000066400000000000000000000067251506752665600164040ustar00rootroot00000000000000# Code of Conduct ## Our Pledge In the interest of fostering an open and welcoming environment, we as contributors and maintainers pledge to make participation in our project and our community a harassment-free experience for everyone, regardless of age, body size, disability, ethnicity, sex characteristics, gender identity and expression, level of experience, education, socio-economic status, nationality, personal appearance, race, religion, or sexual identity and orientation. ## Our Standards Examples of behavior that contributes to creating a positive environment include: * Using welcoming and inclusive language * Being respectful of differing viewpoints and experiences * Gracefully accepting constructive criticism * Focusing on what is best for the community * Showing empathy towards other community members Examples of unacceptable behavior by participants include: * The use of sexualized language or imagery and unwelcome sexual attention or advances * Trolling, insulting/derogatory comments, and personal or political attacks * Public or private harassment * Publishing others' private information, such as a physical or electronic address, without explicit permission * Other conduct which could reasonably be considered inappropriate in a professional setting ## Our Responsibilities Project maintainers are responsible for clarifying the standards of acceptable behavior and are expected to take appropriate and fair corrective action in response to any instances of unacceptable behavior. Project maintainers have the right and responsibility to remove, edit, or reject comments, commits, code, wiki edits, issues, and other contributions that are not aligned to this Code of Conduct, or to ban temporarily or permanently any contributor for other behaviors that they deem inappropriate, threatening, offensive, or harmful. ## Scope This Code of Conduct applies within all project spaces, and it also applies when an individual is representing the project or its community in public spaces. Examples of representing a project or community include using an official project e-mail address, posting via an official social media account, or acting as an appointed representative at an online or offline event. Representation of a project may be further defined and clarified by project maintainers. This Code of Conduct also applies outside the project spaces when there is a reasonable belief that an individual's behavior may have a negative impact on the project or its community. ## Enforcement Instances of abusive, harassing, or otherwise unacceptable behavior may be reported by contacting the project team at . All complaints will be reviewed and investigated and will result in a response that is deemed necessary and appropriate to the circumstances. The project team is obligated to maintain confidentiality with regard to the reporter of an incident. Further details of specific enforcement policies may be posted separately. Project maintainers who do not follow or enforce the Code of Conduct in good faith may face temporary or permanent repercussions as determined by other members of the project's leadership. ## Attribution This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4, available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html [homepage]: https://www.contributor-covenant.org For answers to common questions about this code of conduct, see https://www.contributor-covenant.org/faq bitsandbytes/CONTRIBUTING.md000066400000000000000000000015331506752665600160260ustar00rootroot00000000000000# Contributing to bitsandbytes We want to make contributing to this project as easy and transparent as possible. ## Pull Requests We actively welcome your pull requests. 1. Fork the repo and create your branch from `main`. 2. If you've added code that should be tested, add tests. 3. If you've changed APIs, update the documentation. 4. Ensure the test suite passes. 5. Make sure your code lints, install the [pre-commit hooks as documented here](https://huggingface.co/docs/bitsandbytes/main/en/contributing#setup-pre-commit-hooks). ## Issues We use GitHub issues to track public bugs. Please ensure your description is clear and has sufficient instructions to be able to reproduce the issue. ## License By contributing to bitsandbytes, you agree that your contributions will be licensed under the LICENSE file in the root directory of this source tree. bitsandbytes/LICENSE000066400000000000000000000020761506752665600146050ustar00rootroot00000000000000MIT License Copyright (c) Facebook, Inc. and its affiliates. Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. bitsandbytes/MANIFEST.in000066400000000000000000000000601506752665600153250ustar00rootroot00000000000000include CMakeLists.txt graft csrc graft include bitsandbytes/NOTICE.md000066400000000000000000000004111506752665600150720ustar00rootroot00000000000000The majority of bitsandbytes is licensed under MIT, however portions of the project are available under separate license terms: Pytorch is licensed under the BSD license. We thank Fabio Cannizzo for this work on FastBinarySearch which is included in this project. bitsandbytes/README.md000066400000000000000000000171641506752665600150630ustar00rootroot00000000000000

bitsandbytes

License Downloads Nightly Unit Tests GitHub Release PyPI - Python Version

`bitsandbytes` enables accessible large language models via k-bit quantization for PyTorch. We provide three main features for dramatically reducing memory consumption for inference and training: * 8-bit optimizers uses block-wise quantization to maintain 32-bit performance at a small fraction of the memory cost. * LLM.int8() or 8-bit quantization enables large language model inference with only half the required memory and without any performance degradation. This method is based on vector-wise quantization to quantize most features to 8-bits and separately treating outliers with 16-bit matrix multiplication. * QLoRA or 4-bit quantization enables large language model training with several memory-saving techniques that don't compromise performance. This method quantizes a model to 4-bits and inserts a small set of trainable low-rank adaptation (LoRA) weights to allow training. The library includes quantization primitives for 8-bit & 4-bit operations, through `bitsandbytes.nn.Linear8bitLt` and `bitsandbytes.nn.Linear4bit` and 8-bit optimizers through `bitsandbytes.optim` module. ## System Requirements bitsandbytes has the following minimum requirements for all platforms: * Python 3.9+ * [PyTorch](https://pytorch.org/get-started/locally/) 2.3+ * _Note: While we aim to provide wide backwards compatibility, we recommend using the latest version of PyTorch for the best experience._ #### Accelerator support: Note: this table reflects the status of the current development branch. For the latest stable release, see the [document in the 0.48.0 tag](https://github.com/bitsandbytes-foundation/bitsandbytes/blob/0.48.0/README.md#accelerator-support). ##### Legend: 🚧 = In Development, 〰️ = Partially Supported, ✅ = Supported, ❌ = Not Supported
Platform Accelerator Hardware Requirements LLM.int8() QLoRA 4-bit 8-bit Optimizers
🐧 Linux, glibc >= 2.24
x86-64 ◻️ CPU AVX2 ✅ ✅ ❌
🟩 NVIDIA GPU
cuda		 SM60+ minimum
SM75+ recommended		 ✅ ✅ ✅
🟥 AMD GPU
cuda		 CDNA: gfx90a, gfx942
RDNA: gfx1100 		✅ 〰️ ✅
🟦 Intel GPU
xpu		 Data Center GPU Max Series
Arc A-Series (Alchemist)
Arc B-Series (Battlemage) 		✅ ✅ 〰️
🟪 Intel Gaudi
hpu		 Gaudi2, Gaudi3 ✅ 〰️ ❌
aarch64 ◻️ CPU ✅ ✅ ❌
🟩 NVIDIA GPU
cuda		 SM75+ ✅ ✅ ✅
🪟 Windows 11 / Windows Server 2019+
x86-64 ◻️ CPU AVX2 ✅ ✅ ❌
🟩 NVIDIA GPU
cuda		 SM60+ minimum
SM75+ recommended		 ✅ ✅ ✅
🟦 Intel GPU
xpu		 Arc A-Series (Alchemist)
Arc B-Series (Battlemage) 		✅ ✅ 〰️
🍎 macOS 14+
arm64 ◻️ CPU Apple M1+ 🚧 🚧 ❌
⬜ Metal
mps		 Apple M1+ 🚧 🚧 ❌
## :book: Documentation * [Official Documentation](https://huggingface.co/docs/bitsandbytes/main) * 🤗 [Transformers](https://huggingface.co/docs/transformers/quantization/bitsandbytes) * 🤗 [Diffusers](https://huggingface.co/docs/diffusers/quantization/bitsandbytes) * 🤗 [PEFT](https://huggingface.co/docs/peft/developer_guides/quantization#quantize-a-model) ## :heart: Sponsors The continued maintenance and development of `bitsandbytes` is made possible thanks to the generous support of our sponsors. Their contributions help ensure that we can keep improving the project and delivering valuable updates to the community. Hugging Face   Intel ## License `bitsandbytes` is MIT licensed. We thank Fabio Cannizzo for his work on [FastBinarySearch](https://github.com/fabiocannizzo/FastBinarySearch) which we use for CPU quantization. ## How to cite us If you found this library useful, please consider citing our work: ### QLoRA ```bibtex @article{dettmers2023qlora, title={Qlora: Efficient finetuning of quantized llms}, author={Dettmers, Tim and Pagnoni, Artidoro and Holtzman, Ari and Zettlemoyer, Luke}, journal={arXiv preprint arXiv:2305.14314}, year={2023} } ``` ### LLM.int8() ```bibtex @article{dettmers2022llmint8, title={LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale}, author={Dettmers, Tim and Lewis, Mike and Belkada, Younes and Zettlemoyer, Luke}, journal={arXiv preprint arXiv:2208.07339}, year={2022} } ``` ### 8-bit Optimizers ```bibtex @article{dettmers2022optimizers, title={8-bit Optimizers via Block-wise Quantization}, author={Dettmers, Tim and Lewis, Mike and Shleifer, Sam and Zettlemoyer, Luke}, journal={9th International Conference on Learning Representations, ICLR}, year={2022} } ``` bitsandbytes/_typos.toml000066400000000000000000000010371506752665600160060ustar00rootroot00000000000000[files] # Skip these files in typo checks extend-exclude = [ "csrc/xpu_ops.h", "csrc/xpu_ops.cpp", "csrc/xpu_kernels.h", "csrc/xpu_kernels.cpp" ] [default] extend-ignore-re = [ "@Ther-nul", # valid Github user ] extend-ignore-identifiers-re = [ ".*arange.*", ".*ARANGE.*", ] [type.py.extend-words] "BA" = "BA" # used as a commented-out variable in tests [type.cuda.extend-words] "subtile" = "subtile" "subtiles" = "subtiles" "transation" = "transation" # TODO: is this transition, transaction, translation..? bitsandbytes/benchmarking/000077500000000000000000000000001506752665600162235ustar00rootroot00000000000000bitsandbytes/benchmarking/README.md000066400000000000000000000251541506752665600175110ustar00rootroot00000000000000# Benchmarking ## Inference End-to-end inference benchmarking can be performed using the 🤗 [`optimum-benchmark`](https://github.com/huggingface/optimum-benchmark) library. See the example script in [inference_benchmark.py](inference_benchmark.py). ### Results (as of v0.45.0) Our overall benchmarking results compared with v0.44.1 provide the following insights: #### LLM.int8() * **Turing/Ampere/Ada**: The observed per-token throughput is improved by 60-85%, while latency is decreased by 40-45%. * **H100**: With our benchmarking of Llama 3.1 70B, we observed the new LLM.int8() to consistently outperform NF4 at batch size >= 8. #### NF4/FP4 * **Turing/Ampere/Ada**: With batch size of 1, per-token throughput is _improved by 10-25%_ and per-token latency is _decreased by 10-20%_. * **H100**: Across all batch sizes, per-token throughput is _improved by up to 28%_ and per-token latency is _decreased by up to 22%_. Summaries with the benchmarking results are provided below. #### NVIDIA T4 16GB
Qwen 2.5 3B Instruct | | Batch Size | Mean Latency (s) v0.45.0.dev | Throughput v0.45.0.dev | Mean Latency (s) v0.44.1 | Latency Improvement | Throughput v0.44.1 | Throughput Improvement | |----------------------|------------|------------------------------|------------------------|--------------------------|---------------------|--------------------|------------------------| | FP16 | 1 | 0.0390 | 25.66 | 0.0390 | 1.00 | 25.66 | 1.000x | | NF4 | 1 | 0.0608 | 16.45 | 0.0710 | 1.14 | 14.08 | 1.168x | | NF4+DQ | 1 | 0.0736 | 13.58 | 0.0905 | 1.19 | 11.05 | 1.229x | | INT8 | 1 | 0.0902 | 11.08 | 0.1609 | 1.44 | 6.21 | 1.784x | | INT8+Decomp | 1 | 0.1672 | 5.98 | 0.2994 | 1.44 | 3.34 | 1.790x | | FP16 | 8 | 0.0422 | 189.56 | 0.0422 | 1.00 | 189.56 | 1.000x | | NF4 | 8 | 0.0960 | 83.37 | 0.1010 | 1.05 | 79.17 | 1.053x | | NF4+DQ | 8 | 0.1042 | 76.80 | 0.1156 | 1.10 | 69.18 | 1.110x | | INT8 | 8 | 0.0919 | 87.01 | 0.1640 | 1.44 | 48.78 | 1.784x | | INT8+Decomp | 8 | 0.1812 | 44.15 | 0.3296 | 1.45 | 24.28 | 1.818x | | FP16 | 32 | 0.0601 | 532.30 | 0.0601 | 1.00 | 532.30 | 1.000x | | NF4 | 32 | 0.1150 | 278.32 | 0.1182 | 1.03 | 270.71 | 1.028x | | NF4+DQ | 32 | 0.1215 | 263.36 | 0.1297 | 1.06 | 246.76 | 1.067x | | INT8 | 32 | 0.0943 | 339.21 | 0.1640 | 1.42 | 195.14 | 1.738x | | INT8+Decomp | 32 | 0.1912 | 167.37 | 0.3413 | 1.44 | 93.75 | 1.785x |
#### NVIDIA RTX 4090 24GB
Llama 3.1 8B | | Batch Size | Mean Latency (s) v0.45.0.dev | Throughput v0.45.0.dev | Mean Latency (s) v0.44.1 | Latency Improvement | Throughput v0.44.1 | Throughput Improvement | |----------------------|------------|------------------------------|------------------------|--------------------------|---------------------|--------------------|------------------------| | BF16 | 1 | 0.0211 | 47.46 | 0.0211 | 1.00 | 47.46 | 1.000x | | NF4 | 1 | 0.0148 | 67.71 | 0.0164 | 1.10 | 61.08 | 1.109x | | NF4+DQ | 1 | 0.0175 | 57.08 | 0.0208 | 1.16 | 48.15 | 1.185x | | INT8 | 1 | 0.0220 | 45.39 | 0.0395 | 1.44 | 25.32 | 1.793x | | INT8+Decomp | 1 | 0.0449 | 22.26 | 0.0743 | 1.40 | 13.45 | 1.655x | | BF16 | 8 | 0.0239 | 334.64 | 0.0239 | 1.00 | 334.64 | 1.000x | | NF4 | 8 | 0.0425 | 188.08 | 0.0422 | 0.99 | 189.50 | 0.993x | | NF4+DQ | 8 | 0.0443 | 180.68 | 0.0437 | 0.99 | 183.02 | 0.987x | | INT8 | 8 | 0.0221 | 361.61 | 0.0389 | 1.43 | 205.82 | 1.757x | | INT8+Decomp | 8 | 0.0478 | 164.55 | 0.0777 | 1.38 | 103.01 | 1.597x | | BF16 | 32 | 0.0304 | 1054.35 | 0.0304 | 1.00 | 1054.35 | 1.000x | | NF4 | 32 | 0.0461 | 694.60 | 0.0466 | 1.01 | 686.90 | 1.011x | | NF4+DQ | 32 | 0.0471 | 678.73 | 0.0480 | 1.02 | 666.33 | 1.019x | | INT8 | 32 | 0.0230 | 1390.54 | 0.0390 | 1.41 | 819.99 | 1.696x | | INT8+Decomp | 32 | 0.0512 | 624.94 | 0.0835 | 1.39 | 383.18 | 1.631x |
Qwen 2.5 14B Instruct | | Batch Size | Mean Latency (s) v0.45.0.dev | Throughput v0.45.0.dev | Mean Latency (s) v0.44.1 | Latency Improvement | Throughput v0.44.1 | Throughput Improvement | |----------------------|------------|------------------------------|------------------------|--------------------------|---------------------|--------------------|------------------------| | NF4 | 1 | 0.0214 | 46.74 | 0.0256 | 1.16 | 39.10 | 1.195x | | NF4+DQ | 1 | 0.0256 | 39.03 | 0.0318 | 1.19 | 31.46 | 1.241x | | INT8 | 1 | 0.0326 | 30.68 | 0.0596 | 1.45 | 16.79 | 1.827x | | INT8+Decomp | 1 | 0.0648 | 15.44 | 0.1105 | 1.41 | 9.05 | 1.706x | | NF4 | 8 | 0.0696 | 114.95 | 0.0697 | 1.00 | 114.78 | 1.001x | | NF4+DQ | 8 | 0.0719 | 111.29 | 0.0723 | 1.01 | 110.70 | 1.005x | | INT8 | 8 | 0.0325 | 246.22 | 0.0596 | 1.45 | 134.21 | 1.835x | | INT8+Decomp | 8 | 0.0721 | 110.95 | 0.1201 | 1.40 | 66.62 | 1.665x |
#### NVIDIA H100 80GB SXM
Llama 3.1 8B | | Batch Size | Mean Latency (s) v0.45.0.dev | Throughput v0.45.0.dev | Mean Latency (s) v0.44.1 | Latency Improvement | Throughput v0.44.1 | Throughput Improvement | |----------------------|------------|------------------------------|------------------------|--------------------------|---------------------|--------------------|------------------------| | BF16 | 1 | 0.0244 | 40.99 | 0.0244 | 1.00 | 40.99 | 1.000x | | NF4 | 1 | 0.0331 | 30.14 | 0.0391 | 1.15 | 25.60 | 1.177x | | NF4+DQ | 1 | 0.0411 | 24.34 | 0.0528 | 1.22 | 18.92 | 1.286x | | INT8 | 1 | 0.0522 | 19.17 | N/A | N/A | N/A | N/A | | INT8+Decomp | 1 | 0.0817 | 12.24 | N/A | N/A | N/A | N/A | | BF16 | 8 | 0.0255 | 313.90 | 0.0255 | 1.00 | 313.90 | 1.000x | | NF4 | 8 | 0.0476 | 168.05 | 0.0551 | 1.14 | 145.13 | 1.158x | | NF4+DQ | 8 | 0.0566 | 141.27 | 0.0663 | 1.15 | 120.67 | 1.171x | | INT8 | 8 | 0.0515 | 155.44 | N/A | N/A | N/A | N/A | | INT8+Decomp | 8 | 0.0853 | 93.79 | N/A | N/A | N/A | N/A | | BF16 | 32 | 0.0261 | 1227.96 | 0.0261 | 1.00 | 1227.96 | 1.000x | | NF4 | 32 | 0.0486 | 658.65 | 0.0546 | 1.11 | 585.91 | 1.124x | | NF4+DQ | 32 | 0.0577 | 555.06 | 0.0665 | 1.13 | 481.04 | 1.154x | | INT8 | 32 | 0.0545 | 586.26 | N/A | N/A | N/A | N/A | | INT8+Decomp | 32 | 0.0864 | 370.51 | N/A | N/A | N/A | N/A |
Qwen 2.5 32B Instruct | | Batch Size | Mean Latency (s) v0.45.0.dev | Throughput v0.45.0.dev | |-------------|------------|-----------------------------------------|-----------------------------------| | BF16 | 1 | 0.0508 | 19.67 | | NF4 | 1 | 0.0707 | 14.14 | | NF4+DQ | 1 | 0.0860 | 11.63 | | INT8 | 1 | 0.1031 | 9.70 | | INT8+Decomp | 1 | 0.1820 | 5.49 | | BF16 | 8 | 0.0525 | 152.50 | | NF4 | 8 | 0.1154 | 69.35 | | NF4+DQ | 8 | 0.1209 | 66.19 | | INT8 | 8 | 0.1078 | 74.24 | | INT8+Decomp | 8 | 0.1958 | 40.87 | | BF16 | 32 | 0.0547 | 584.54 | | NF4 | 32 | 0.1246 | 256.84 | | NF4+DQ | 32 | 0.1298 | 246.47 | | INT8 | 32 | 0.1056 | 302.96 | | INT8+Decomp | 32 | 0.2027 | 157.83 |
Llama 3.1 70B | | Batch Size | Mean Latency (s) v0.45.0.dev | Throughput v0.45.0.dev | |-------------|------------|-----------------------------------------|-----------------------------------| | NF4 | 1 | 0.0833 | 12.00 | | NF4+DQ | 1 | 0.1052 | 9.50 | | INT8 | 1 | 0.1294 | 7.73 | | INT8+Decomp | 1 | 0.1985 | 5.04 | | NF4 | 8 | 0.2348 | 34.07 | | NF4+DQ | 8 | 0.2423 | 33.01 | | INT8 | 8 | 0.1313 | 60.94 | | INT8+Decomp | 8 | 0.2052 | 38.99 | | NF4 | 32 | 0.2491 | 128.46 | | NF4+DQ | 32 | 0.2580 | 124.04 | | INT8 | 32 | 0.1314 | 243.45 | | INT8+Decomp | 32 | 0.2189 | 146.19 |
#### Software Configuration We focus on the default PyTorch CUDA backend in 🤗 [`optimum-benchmark`](https://github.com/huggingface/optimum-benchmark). We used commit [`6e6b1036`](https://github.com/huggingface/optimum-benchmark/commit/6e6b10363f3ac65926881f2c6a6113b6cefc06cd). For all hardware configurations, we used the following dependencies: * `transformers==4.46.3` * `accelerate==1.1.1` * `tokenizers==0.20.3` * `torch==2.5.1` * `bitsandbytes==0.44.1` * `bitsandbytes==0.45.0.dev` In the RTX 4090 setting, the CUDA 12.4 build of PyTorch is used. In the other settings we used the CUDA 12.1 build. bitsandbytes/benchmarking/inference_benchmark.py000066400000000000000000000135361506752665600225550ustar00rootroot00000000000000""" Inference benchmarking tool. Requirements: transformers accelerate bitsandbytes optimum-benchmark Usage: python inference_benchmark.py model_id options: -h, --help show this help message and exit --configs {bf16,fp16,nf4,nf4-dq,int8,int8-decomp} [{bf16,fp16,nf4,nf4-dq,int8,int8-decomp} ...] --bf16 --fp16 --nf4 --nf4-dq --int8 --int8-decomp --batches BATCHES [BATCHES ...] --input-length INPUT_LENGTH --out-dir OUT_DIR --iterations ITERATIONS --warmup-runs WARMUP_RUNS --output-length OUTPUT_LENGTH """ import argparse from pathlib import Path from optimum_benchmark import Benchmark, BenchmarkConfig, InferenceConfig, ProcessConfig, PyTorchConfig from optimum_benchmark.logging_utils import setup_logging import torch torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True BFLOAT16_SUPPORT = torch.cuda.get_device_capability()[0] >= 8 WEIGHTS_CONFIGS = { "fp16": {"torch_dtype": "float16", "quantization_scheme": None, "quantization_config": {}}, "bf16": {"torch_dtype": "bfloat16", "quantization_scheme": None, "quantization_config": {}}, "nf4": { "torch_dtype": "bfloat16" if BFLOAT16_SUPPORT else "float16", "quantization_scheme": "bnb", "quantization_config": { "load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_use_double_quant": False, "bnb_4bit_compute_dtype": torch.bfloat16 if BFLOAT16_SUPPORT else "float16", }, }, "nf4-dq": { "torch_dtype": "bfloat16" if BFLOAT16_SUPPORT else "float16", "quantization_scheme": "bnb", "quantization_config": { "load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_use_double_quant": True, "bnb_4bit_compute_dtype": torch.bfloat16 if BFLOAT16_SUPPORT else "float16", }, }, "int8-decomp": { "torch_dtype": "float16", "quantization_scheme": "bnb", "quantization_config": { "load_in_8bit": True, "llm_int8_threshold": 6.0, }, }, "int8": { "torch_dtype": "float16", "quantization_scheme": "bnb", "quantization_config": { "load_in_8bit": True, "llm_int8_threshold": 0.0, }, }, } def parse_args(): parser = argparse.ArgumentParser(description="bitsandbytes inference benchmark tool") parser.add_argument("model_id", type=str, help="The model checkpoint to use.") parser.add_argument( "--configs", nargs="+", choices=["bf16", "fp16", "nf4", "nf4-dq", "int8", "int8-decomp"], default=["nf4", "int8", "int8-decomp"], ) parser.add_argument("--bf16", dest="configs", action="append_const", const="bf16") parser.add_argument("--fp16", dest="configs", action="append_const", const="fp16") parser.add_argument("--nf4", dest="configs", action="append_const", const="nf4") parser.add_argument("--nf4-dq", dest="configs", action="append_const", const="nf4-dq") parser.add_argument("--int8", dest="configs", action="append_const", const="int8") parser.add_argument("--int8-decomp", dest="configs", action="append_const", const="int8-decomp") parser.add_argument("--batches", nargs="+", type=int, default=[1, 8, 16, 32]) parser.add_argument("--input-length", type=int, default=64) parser.add_argument("--out-dir", type=str, default="reports") parser.add_argument("--iterations", type=int, default=10, help="Number of iterations for each benchmark run") parser.add_argument( "--warmup-runs", type=int, default=10, help="Number of warmup runs to discard before measurement" ) parser.add_argument( "--output-length", type=int, default=64, help="If set, `max_new_tokens` and `min_new_tokens` will be set to this value.", ) return parser.parse_args() def run_benchmark(args, config, batch_size): launcher_config = ProcessConfig(device_isolation=True, device_isolation_action="warn", start_method="spawn") scenario_config = InferenceConfig( latency=True, memory=True, input_shapes={"batch_size": batch_size, "sequence_length": args.input_length}, iterations=args.iterations, warmup_runs=args.warmup_runs, # set duration to 0 to disable the duration-based stopping criterion # this is IMPORTANT to ensure that all benchmarks run the same number of operations, regardless of hardware speed/bottlenecks duration=0, # for consistent results, set a fixed min and max for output tokens generate_kwargs={"min_new_tokens": args.output_length, "max_new_tokens": args.output_length}, forward_kwargs={"min_new_tokens": args.output_length, "max_new_tokens": args.output_length}, ) backend_config = PyTorchConfig( device="cuda", device_ids="0", device_map="auto", no_weights=False, model=args.model_id, **WEIGHTS_CONFIGS[config], ) test_name = ( f"benchmark-{config}" f"-bsz-{batch_size}" f"-isz-{args.input_length}" f"-osz-{args.output_length}" f"-iter-{args.iterations}" f"-wrmup-{args.warmup_runs}" ) benchmark_config = BenchmarkConfig( name=test_name, scenario=scenario_config, launcher=launcher_config, backend=backend_config, ) out_path = out_dir / (test_name + ".json") print(f"[{test_name}] Starting:") benchmark_report = Benchmark.launch(benchmark_config) benchmark_report.save_json(out_path) if __name__ == "__main__": setup_logging(level="INFO") args = parse_args() out_dir = Path(args.out_dir) out_dir.mkdir(parents=True, exist_ok=True) for batch_size in args.batches: for config in args.configs: run_benchmark(args, config, batch_size) bitsandbytes/benchmarking/int8/000077500000000000000000000000001506752665600171055ustar00rootroot00000000000000bitsandbytes/benchmarking/int8/int8_benchmark.py000066400000000000000000000031731506752665600223570ustar00rootroot00000000000000""" Basic benchmark for text generation. Usage: python benchmarking/int8/int8_benchmark.py """ import time import torch from torch.profiler import ProfilerActivity, profile from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig MAX_NEW_TOKENS = 128 model_name = "meta-llama/Llama-3.1-8B" text = "Below is a question. I need an answer.\n\nExplain machine learning: " tokenizer = AutoTokenizer.from_pretrained(model_name) input_ids = tokenizer([text] * 8, return_tensors="pt").input_ids.to(0) model = AutoModelForCausalLM.from_pretrained( model_name, device_map="auto", quantization_config=BitsAndBytesConfig( load_in_8bit=True, llm_int8_threshold=6.0, ), attn_implementation="sdpa", torch_dtype=torch.float16, ) print(model) # warmup print("Warmup...") for i in range(3): generated_ids = model.generate(input_ids, max_new_tokens=MAX_NEW_TOKENS) print("Profiler starting...") with profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], with_modules=True, with_stack=True, ) as prof: model.generate(input_ids, max_new_tokens=1) print( prof.key_averages().table( sort_by="cpu_time_total", max_name_column_width=50, top_level_events_only=True, row_limit=50, ) ) torch.cuda.synchronize() print("Generating...") num = 0 time_1 = time.time() for i in range(5): generated_ids = model.generate(input_ids, max_new_tokens=MAX_NEW_TOKENS) num += len(generated_ids[0]) print("=" * 40) print(f"Example:\n{tokenizer.decode(generated_ids[0])}") print("=" * 40) print(f"Speed: {num / (time.time() - time_1)}token/s") bitsandbytes/benchmarking/int8/training_benchmark.py000066400000000000000000000153301506752665600233060ustar00rootroot00000000000000""" Extracted from tests/test_functional.py Usage: pytest benchmarking/int8/training_benchmark.py """ import time import pytest import torch from bitsandbytes import functional as F k = 20 torch.set_printoptions(precision=5, sci_mode=False, linewidth=120, edgeitems=20, threshold=10000) @pytest.mark.parametrize( ("batch", "seq", "model", "hidden"), [ pytest.param(2, 512, 4 * 1024, 3 * 4 * 1024, id="batch=2, seq=512, model=4k, hidden=12k"), pytest.param(2, 512, 5120, 3 * 5120, id="batch=2, seq=512, model=5k, hidden=15k"), pytest.param(2, 512, 12 * 1024, 4 * 12 * 1024, id="batch=2, seq=512, model=12k, hidden=48k"), ], ) @pytest.mark.benchmark def test_bench_8bit_training(batch, seq, model, hidden): formatB = F.get_special_format_str() A = torch.randn(batch, seq, model, device="cuda").half() grad = torch.randn(batch, seq, model, device="cuda").half() w1 = torch.randint(-128, 127, size=(hidden, model), device="cuda").half() w2 = torch.randint(-128, 127, size=(model, hidden), device="cuda").half() print("") # torch.cuda.synchronize() ## warmup # for i in range(100): # torch.matmul(A, w1.t()) # torch.cuda.synchronize() dtype = torch.int8 A = A.view(-1, A.shape[-1]).contiguous() grad = grad.view(-1, grad.shape[-1]).contiguous() torch.cuda.synchronize() t0 = time.time() for i in range(k): out1 = torch.matmul(A, w1.t()) # fc1 # out2 = torch.matmul(out1, w2.t())# fc2 # d1 = torch.matmul(grad, w2) # delta1 # d2 = torch.matmul(d1, w1) # delta2 # grad1 = torch.einsum('bo,bh->oh', out1, grad) # grad w2 # grad2 = torch.einsum('bh,bo->ho', A, d2) # grad w1 torch.cuda.synchronize() t16 = time.time() - t0 print(t16) # torch.cuda.empty_cache() # Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1) # Cw2, Cw2t, statsw2, statsw2t, coo_tensor = F.double_quant(w2) # CTw1, Sw1 = F.transform2(Cw1, formatB) # CTw2, Sw2 = F.transform2(Cw2, formatB) # CTw2t, Sw2t = F.transform2(Cw2t, formatB, transpose=True) # CTw1t, Sw1t = F.transform2(Cw1t, formatB, transpose=True) # CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A) # C32A, SA = F.transform2(CA, 'col32') ## fc1 # out1_32, Sout1_32 = F.igemmlt(C32A, CTw1, SA, Sw1, dtype=dtype) ##out1 = F.mm_dequant(out1_32, Sout1_32, statsAt, statsw1t) ## fc2 # Cout1, Cout1t, statsout1, statsout1t, coo_tensor = F.double_quant(out1) # C32out1, Sout1 = F.transform2(Cout1, 'col32') # out2_32, Sout2_32 = F.igemmlt(C32out1, CTw2, Sout1, Sw2, dtype=dtype) ##out2 = F.mm_dequant(out2_32, Sout2_32, statsout1t, statsw2t) ## delta1 # Cgrad, Cgradt, statsgrad, statsgradt, coo_tensor = F.double_quant(grad) # C32grad, Sgrad = F.transform2(Cgrad, 'col32') ##d1_32, Sd1_32 = F.igemmlt(C32grad, CTw2t, Sgrad, Sw2t, dtype=dtype) ##d1 = F.mm_dequant(d1_32, Sd1_32, statsgradt, statsw2) ## delta2 # Cd1, Cd1t, statsd1, statsd1t, coo_tensor = F.double_quant(d1) # C32d1, Sd1 = F.transform2(Cd1, 'col32') ##d2_32, Sd2_32 = F.igemmlt(C32d1, CTw1t, Sd1, Sw1t, dtype=dtype) ##d2 = F.mm_dequant(d2_32, Sd2_32, statsd1t, statsw1) ## grad1 # C32out1t, Sout1t = F.transform2(Cout1t, 'col32', transpose=True) # CTgradt, Sgradt = F.transform2(Cgradt, formatB, transpose=True) ##grad1_32, Sgrad1_32 = F.igemmlt(C32out1t, CTgradt, Sout1t, Sgradt, dtype=dtype) ##grad1 = F.mm_dequant(grad1_32, Sgrad1_32, statsout1, statsgrad) ## grad2 # C32At, SAt = F.transform2(CAt, 'col32', transpose=True) # CTd1t, Sd1t = F.transform2(Cd1t, formatB, transpose=True) ##grad2_32, Sgrad2_32 = F.igemmlt(C32At, CTd1t, SAt, Sd1t, dtype=dtype) ##grad2 = F.mm_dequant(grad2_32, Sgrad2_32, statsA, statsd1) # Cw2, Cw2t, statsw2, statsw2t, coo_tensor = F.double_quant(w2) # Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1) # Cw2, Cw2t, statsw2, statsw2t, coo_tensor = F.double_quant(w2) # CTw1, Sw1 = F.transform2(Cw1, formatB) # CTw1t, Sw1t = F.transform2(Cw1t, formatB, transpose=True) # CTw2, Sw2 = F.transform2(Cw2, formatB) # CTw2t, Sw2t = F.transform2(Cw2t, formatB, transpose=True) # torch.cuda.synchronize() # t0 = time.time() # for i in range(k): # #Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1) # #CTw1, Sw1 = F.transform2(Cw1, formatB) # #Cw1, Cw1t, statsw1, statsw1t, coo_tensor = F.double_quant(w1) # #CTw1, Sw1 = F.transform2(Cw1, formatB) # #CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A, threshold=3.5) # CA, CAt, statsA, statsAt, coo_tensor = F.double_quant(A) # #CTw1t, Sw1t = F.transform2(Cw1t, formatB, transpose=True) # #CTw2, Sw2 = F.transform2(Cw2, formatB) # #CTw2t, Sw2t = F.transform2(Cw2t, formatB, transpose=True) # C32A, SA = F.transform2(CA, 'col32') # # fc1 # out1_32, Sout1_32 = F.igemmlt(C32A, CTw1, SA, Sw1, dtype=dtype) # #out1dn = F.mm_dequant(out1_32, Sout1_32, statsA, statsw1) # #print(coo_tensor.nnz) # #out1sp = F.spmm_coo(coo_tensor, w1.t()) # #print(w1.t().shape) # #out1 = out1dn + out1sp # # fc2 # Cout1, Cout1t, statsout1, statsout1t, coo_tensor = F.double_quant(out1) # C32out1, Sout1 = F.transform2(Cout1, 'col32') # out2_32, Sout2_32 = F.igemmlt(C32out1, CTw2, Sout1, Sw2, dtype=dtype) # #out2 = F.mm_dequant(out2_32, Sout2_32, statsout1, statsw2) # # delta1 # Cgrad, Cgradt, statsgrad, statsgradt, coo_tensor = F.double_quant(grad) # C32grad, Sgrad = F.transform2(Cgrad, 'col32') # d1_32, Sd1_32 = F.igemmlt(C32grad, CTw2t, Sgrad, Sw2t, dtype=dtype) # #d1 = F.mm_dequant(d1_32, Sd1_32, statsgrad, statsw2t) # # delta2 # Cd1, Cd1t, statsd1, statsd1t, coo_tensor = F.double_quant(d1) # C32d1, Sd1 = F.transform2(Cd1, 'col32') # d2_32, Sd2_32 = F.igemmlt(C32d1, CTw1t, Sd1, Sw1t, dtype=dtype) # #d2 = F.mm_dequant(d2_32, Sd2_32, statsd1, statsw1t) # # grad1 # #C32out1t, Sout1t = F.transform2(Cout1t, 'col32', transpose=True) # #CTgradt, Sgradt = F.transform2(Cgradt, formatB, transpose=True) # #grad1_32, Sgrad1_32 = F.igemmlt(C32out1t, CTgradt, Sout1t, Sgradt, dtype=dtype) # #grad1 = F.mm_dequant(grad1_32, Sgrad1_32, statsout1t, statsgradt) # ## grad2 # #C32At, SAt = F.transform2(CAt, 'col32', transpose=True) # #CTd1t, Sd1t = F.transform2(Cd1t, formatB, transpose=True) # #grad2_32, Sgrad2_32 = F.igemmlt(C32At, CTd1t, SAt, Sd1t, dtype=dtype) # #grad2 = F.mm_dequant(grad2_32, Sgrad2_32, statsAt, statsd1t) # torch.cuda.synchronize() # t8 = time.time() - t0 # print(t8) bitsandbytes/benchmarking/matmul_benchmark.py000066400000000000000000000176121506752665600221150ustar00rootroot00000000000000""" Extracted from tests/test_functional.py Usage: pytest benchmarking/matmul_benchmark.py """ import time import pytest import torch import bitsandbytes as bnb from bitsandbytes import functional as F k = 20 torch.set_printoptions(precision=5, sci_mode=False, linewidth=120, edgeitems=20, threshold=10000) @pytest.mark.parametrize( ("batch", "seq", "model", "hidden"), [ # pytest.param(1, 128, 6656, 4 * 6656, id="batch=1, seq=128, model=6656, hidden=26k"), pytest.param(1, 1, 3584, 512, id="batch=1, seq=128, model=3584, hidden=19k"), # pytest.param(4, 128, 6656, 4 * 6656, id="batch=4, seq=128, model=6656, hidden=26k"), # pytest.param(16, 256, 6656, 4 * 6656, id="batch=16, seq=256, model=6656, hidden=26k") ], ) @pytest.mark.benchmark def test_bench_matmul(batch, seq, model, hidden): iters = 1000 formatB = F.get_special_format_str() A = torch.randn(batch, seq, model, device="cuda").half() B = torch.empty(hidden, model, dtype=torch.float16, device="cuda") torch.nn.init.xavier_uniform_(B) B_fp4, state = F.quantize_fp4(B) B_fp4_c, state_c = F.quantize_fp4(B, compress_statistics=True) B_nf4, state_nf4 = F.quantize_nf4(B) B_nf4_c, state_nf4_c = F.quantize_nf4(B, compress_statistics=True) linear8bit = bnb.nn.Linear8bitLt(model, hidden, False, False).cuda().half() linear8bit.eval() outliers = torch.randint(0, model, size=(5,)).cuda() A[:, :, outliers] = 8.0 linearMixedBit = bnb.nn.Linear8bitLt(model, hidden, False, False, threshold=6.0).cuda().half() # linearMixedBit.eval() linear8bit_train = bnb.nn.Linear8bitLt(model, hidden, False).cuda().half() linear8bit_train_thresh = bnb.nn.Linear8bitLt(model, hidden, False, threshold=6.0).cuda().half() bnb.matmul_4bit(A, B_nf4.t(), quant_state=state_nf4) # warmup for i in range(iters): torch.matmul(A, B.t()) torch.cuda.synchronize() print("") torch.cuda.synchronize() t0 = time.time() for i in range(iters): torch.matmul(A, B.t()) torch.cuda.synchronize() print( f"pytorch fp16: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s", ) # torch.cuda.synchronize() # t0 = time.time() # for i in range(iters): # bnb.matmul_4bit(A, B_fp4.t(), quant_state=state) # torch.cuda.synchronize() # print( f"bnb fp4: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s" ) # torch.cuda.synchronize() # t0 = time.time() # for i in range(iters): # bnb.matmul_4bit(A, B_fp4.t(), quant_state=state_c) # torch.cuda.synchronize() # print( f"bnb fp4 + compressed stats: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s" ) torch.cuda.synchronize() t0 = time.time() for i in range(iters): bnb.matmul_4bit(A, B_nf4.t(), quant_state=state_nf4) torch.cuda.synchronize() print(f"bnb nf4: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s") torch.cuda.synchronize() t0 = time.time() for i in range(iters): bnb.matmul_4bit(A, B_nf4_c.t(), quant_state=state_nf4_c) torch.cuda.synchronize() print( f"bnb nf4+DQ: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s" ) torch.cuda.synchronize() t0 = time.time() for i in range(iters): bnb.matmul(A, B) torch.cuda.synchronize() print( f"B -> CB (each iteration): [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s" ) torch.cuda.synchronize() t0 = time.time() for i in range(iters): bnb.matmul(A, B, threshold=6.0) torch.cuda.synchronize() print( f"B -> CB + threshold: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s" ) CA, SCA, _ = F.int8_vectorwise_quant(A, threshold=0.0) CB, SCB, _ = F.int8_vectorwise_quant(B) torch.cuda.synchronize() t0 = time.time() for i in range(iters): # CA, CAt, SCA, SCAt, coo_tensorA = F.double_quant(A, threshold=0.0) out32 = F.int8_linear_matmul(CA, CB) torch.cuda.synchronize() print( f"no overhead int8 [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s" ) # C32A, SA = F.transform(CA, "col32") # CxB, SB = F.transform(CB, to_order=formatB) # torch.cuda.synchronize() # t0 = time.time() # for i in range(iters): # out32, Sout32 = F.igemmlt(C32A, CxB, SA, SB) # torch.cuda.synchronize() # print(f"no overhead matmul-lt: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s") # CA, CAt, SCA, SCAt, coo_tensorA = F.double_quant(A, threshold=0.0) # C32A, SA = F.transform(CA, "col32") # CB, CBt, SCB, SCBt, coo_tensorB = F.double_quant(B) # CxB, SB = F.transform(CB, to_order=formatB) # torch.cuda.synchronize() # t0 = time.time() # for i in range(iters): # out32, Sout32 = F.igemmlt(C32A, CxB, SA, SB) # torch.cuda.synchronize() # print(f"no overhead matmul-lt: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s") # BA, statsB = F.vectorwise_quant(B, dim=1) # CxB, SB = F.nvidia_transform(CB, to_order=formatB) # torch.cuda.synchronize() # t0 = time.time() # for i in range(iters): # A2 = A.view(-1, A.shape[-1]).contiguous() # CA, statsA = F.vectorwise_quant(A2, dim=1) # C32A, SA = F.nvidia_transform(CA, "col32") # out32, Sout32 = F.igemmlt(C32A, CxB, SA, SB) # Cout, Sout = F.nvidia_transform(out32, "row", state=Sout32) # F.vectorwise_mm_dequant(Cout, statsA, statsB.t()) # torch.cuda.synchronize() # print(f"vector pytorch + nvidia: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s") # BA, statsB = F.vectorwise_quant(B, dim=1, quant_type="linear") # CxB, SB = F.nvidia_transform(CB, to_order=formatB) # torch.cuda.synchronize() # t0 = time.time() # for i in range(iters): # A2 = A.view(-1, A.shape[-1]).contiguous() # CA, statsA = F.vectorwise_quant(A2, dim=1, quant_type="linear") # C32A, SA = F.nvidia_transform(CA, "col32") # out32, Sout32 = F.igemmlt(C32A, CxB, SA, SB) # Cout, Sout = F.nvidia_transform(out32, "row", state=Sout32) # out = Cout * statsB * statsA * (1.0 / (127 * 127)) # torch.cuda.synchronize() # print(f"linear pytorch + nvidia: [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s") linear8bit(A) torch.cuda.synchronize() t0 = time.time() for i in range(iters): linear8bit(A) torch.cuda.synchronize() print( f"bnb linear8bitlt (eval): [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s" ) linearMixedBit(A) torch.cuda.synchronize() t0 = time.time() for i in range(iters): linearMixedBit(A) torch.cuda.synchronize() print( f"bnb linear8bitlt with threshold (eval): [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time() - t0:.4f}s" ) # linear8bit_train(A) # torch.cuda.synchronize() # t0 = time.time() # for i in range(iters): # linear8bit_train(A) # torch.cuda.synchronize() # print( f"bnb linear8bitlt (training): [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s") # linear8bit_train_thresh(A) # torch.cuda.synchronize() # t0 = time.time() # for i in range(iters): # linear8bit_train(A) # torch.cuda.synchronize() # print( f"bnb linear8bitlt with threshold (training): [{batch},{seq},{model}], [{model},{hidden}]->[{batch},{seq},{hidden}]: {time.time()-t0:.4f}s") bitsandbytes/benchmarking/optimizer_benchmark.py000066400000000000000000000032451506752665600226350ustar00rootroot00000000000000""" Extracted from tests/test_optim.py Usage: pytest benchmarking/optimizer_benchmark.py """ import time import pytest from tests.helpers import describe_dtype, id_formatter import torch import bitsandbytes as bnb str2optimizers = {"paged_adamw": (torch.optim.AdamW, bnb.optim.PagedAdamW)} @pytest.mark.parametrize("dim1", [2 * 1024], ids=id_formatter("dim1")) @pytest.mark.parametrize("gtype", [torch.float16], ids=describe_dtype) @pytest.mark.parametrize("optim_name", ["paged_adamw"], ids=id_formatter("optim_name")) @pytest.mark.parametrize("mode", ["bnb"], ids=id_formatter("mode")) @pytest.mark.benchmark def test_stream_optimizer_bench(dim1, gtype, optim_name, mode): layers1 = torch.nn.Sequential(*torch.nn.ModuleList([torch.nn.Linear(dim1, dim1) for i in range(10)])) layers1 = layers1.to(gtype) layers1 = layers1.cuda() large_tensor = None if mode == "torch": optim = str2optimizers[optim_name][0](layers1.parameters()) else: optim = str2optimizers[optim_name][1](layers1.parameters()) # 12 GB large_tensor = torch.empty((int(4.5e9),), device="cuda") torch.cuda.synchronize() time.sleep(5) num_batches = 5 batches = torch.randn(num_batches, 128, dim1, device="cuda").to(gtype) lbls = torch.randint(0, 10, size=(num_batches, 128)).cuda() for i in range(num_batches): print(i) b = batches[i] if i == 2: torch.cuda.synchronize() t0 = time.time() out1 = layers1(b) loss1 = torch.nn.functional.cross_entropy(out1, lbls[i]).mean() loss1.backward() optim.step() torch.cuda.synchronize() print(mode, time.time() - t0) bitsandbytes/benchmarking/switchback/000077500000000000000000000000001506752665600203455ustar00rootroot00000000000000bitsandbytes/benchmarking/switchback/README.md000066400000000000000000000006111506752665600216220ustar00rootroot00000000000000Steps: 1. Run `python speed_benchmark/speed_benchmark.py` which times operations and writes their time to `speed_benchmark/info_a100_py2.jsonl` (change the name of the jsonl to a different name for your profiling). 2. Run `python speed_benchmark/make_plot_with_jsonl.py`, which produces the `speed_benchmark/plot_with_info.pdf`. Again make sure you change the jsonl which is being processed. bitsandbytes/benchmarking/switchback/info_a100_py2.jsonl000066400000000000000000001237601506752665600236730ustar00rootroot00000000000000{"repeat": 64, "batch_size": 8192, "dim_out": 4096, "dim_in": 1024, "wm": 4, "switch": false, "standard_fwd": 0.28139352798461914, "standard_gw": 0.2811811864376068, "standard_gx": 0.30258670449256897, "rowwise_fwd": 0.1994594931602478, "rowwise_bwd": 0.16159191727638245, "global_fwd": 0.19502267241477966, "global_bwd": 0.16080215573310852, "x_quantize_rowwise": 0.03306940197944641, "g_quantize_rowwise": 0.08210167288780212, "w_quantize_rowwise": 0.03385916352272034, "w_quantize_colwise_transpose": 0.08635595440864563, "w_quantize_global": 0.09237229824066162, "w_quantize_global_transpose": 0.10007619857788086, "time_standard": 0.8651614189147949, "time_rowwise": 0.8776187896728516, "time_global": 0.944625586271286} {"repeat": 64, "batch_size": 8192, "dim_out": 1024, "dim_in": 4096, "wm": 4, "switch": true, "standard_fwd": 0.262625515460968, "standard_gw": 0.2806223928928375, "standard_gx": 0.31118839979171753, "rowwise_fwd": 0.1828707754611969, "rowwise_bwd": 0.21236762404441833, "global_fwd": 0.16665831208229065, "global_bwd": 0.19929558038711548, "x_quantize_rowwise": 0.08227676153182983, "g_quantize_rowwise": 0.03310292959213257, "w_quantize_rowwise": 0.032648444175720215, "w_quantize_colwise_transpose": 0.09015202522277832, "w_quantize_global": 0.0988692045211792, "w_quantize_global_transpose": 0.10057538747787476, "time_standard": 0.8544363081455231, "time_rowwise": 0.9140409529209137, "time_global": 0.96140056848526} {"repeat": 64, "batch_size": 16384, "dim_out": 4096, "dim_in": 1024, "wm": 4, "switch": false, "standard_fwd": 0.5731917917728424, "standard_gw": 0.5709454417228699, "standard_gx": 0.5963630974292755, "rowwise_fwd": 0.37662312388420105, "rowwise_bwd": 0.281747430562973, "global_fwd": 0.36768242716789246, "global_bwd": 0.28043612837791443, "x_quantize_rowwise": 0.046547502279281616, "g_quantize_rowwise": 0.15532970428466797, "w_quantize_rowwise": 0.032436102628707886, "w_quantize_colwise_transpose": 0.08635222911834717, "w_quantize_global": 0.0947415828704834, "w_quantize_global_transpose": 0.10129809379577637, "time_standard": 1.7405003309249878, "time_rowwise": 1.5499815344810486, "time_global": 1.616980880498886} {"repeat": 64, "batch_size": 16384, "dim_out": 1024, "dim_in": 4096, "wm": 4, "switch": true, "standard_fwd": 0.5341619253158569, "standard_gw": 0.5690865218639374, "standard_gx": 0.599835067987442, "rowwise_fwd": 0.3233291208744049, "rowwise_bwd": 0.41359663009643555, "global_fwd": 0.2831108868122101, "global_bwd": 0.37280842661857605, "x_quantize_rowwise": 0.15563145279884338, "g_quantize_rowwise": 0.046741217374801636, "w_quantize_rowwise": 0.03306940197944641, "w_quantize_colwise_transpose": 0.09020790457725525, "w_quantize_global": 0.0925213098526001, "w_quantize_global_transpose": 0.09945780038833618, "time_standard": 1.7030835151672363, "time_rowwise": 1.6316622495651245, "time_global": 1.6193576157093048} {"repeat": 64, "batch_size": 32768, "dim_out": 4096, "dim_in": 1024, "wm": 4, "switch": false, "standard_fwd": 1.2199915945529938, "standard_gw": 1.1069811880588531, "standard_gx": 1.09761580824852, "rowwise_fwd": 0.738043338060379, "rowwise_bwd": 0.5549229681491852, "global_fwd": 0.7219798862934113, "global_bwd": 0.5512163043022156, "x_quantize_rowwise": 0.08748471736907959, "g_quantize_rowwise": 0.3023110330104828, "w_quantize_rowwise": 0.03182142972946167, "w_quantize_colwise_transpose": 0.08632615208625793, "w_quantize_global": 0.09445473551750183, "w_quantize_global_transpose": 0.10032951831817627, "time_standard": 3.424588590860367, "time_rowwise": 2.9078908264636993, "time_global": 2.9647573828697205} {"repeat": 64, "batch_size": 32768, "dim_out": 1024, "dim_in": 4096, "wm": 4, "switch": true, "standard_fwd": 1.1040829122066498, "standard_gw": 1.1221766471862793, "standard_gx": 1.1548101902008057, "rowwise_fwd": 0.581938773393631, "rowwise_bwd": 0.7480122148990631, "global_fwd": 0.5537159740924835, "global_bwd": 0.7232688367366791, "x_quantize_rowwise": 0.30193477869033813, "g_quantize_rowwise": 0.08745118975639343, "w_quantize_rowwise": 0.03374740481376648, "w_quantize_colwise_transpose": 0.09068101644515991, "w_quantize_global": 0.09645149111747742, "w_quantize_global_transpose": 0.10189786553382874, "time_standard": 3.3810697495937347, "time_rowwise": 2.9659420251846313, "time_global": 2.9868967831134796} {"repeat": 64, "batch_size": 65536, "dim_out": 4096, "dim_in": 1024, "wm": 4, "switch": false, "standard_fwd": 2.4533793330192566, "standard_gw": 2.1938569843769073, "standard_gx": 2.179361879825592, "rowwise_fwd": 1.4615543186664581, "rowwise_bwd": 1.0522231459617615, "global_fwd": 1.4288239181041718, "global_bwd": 1.0450035333633423, "x_quantize_rowwise": 0.1691766083240509, "g_quantize_rowwise": 0.5951300263404846, "w_quantize_rowwise": 0.03337860107421875, "w_quantize_colwise_transpose": 0.08653849363327026, "w_quantize_global": 0.0940859317779541, "w_quantize_global_transpose": 0.09976327419281006, "time_standard": 6.826598197221756, "time_rowwise": 5.5918581783771515, "time_global": 5.625840276479721} {"repeat": 64, "batch_size": 65536, "dim_out": 1024, "dim_in": 4096, "wm": 4, "switch": true, "standard_fwd": 2.1698065102100372, "standard_gw": 2.1875128149986267, "standard_gx": 2.2887587547302246, "rowwise_fwd": 1.0762326419353485, "rowwise_bwd": 1.4638006687164307, "global_fwd": 1.0450668632984161, "global_bwd": 1.4308765530586243, "x_quantize_rowwise": 0.5953535437583923, "g_quantize_rowwise": 0.16899779438972473, "w_quantize_rowwise": 0.03240257501602173, "w_quantize_colwise_transpose": 0.09106099605560303, "w_quantize_global": 0.09546056389808655, "w_quantize_global_transpose": 0.09852275252342224, "time_standard": 6.6460780799388885, "time_rowwise": 5.615361034870148, "time_global": 5.621790885925293} {"repeat": 64, "batch_size": 131072, "dim_out": 4096, "dim_in": 1024, "wm": 4, "switch": false, "standard_fwd": 4.858218133449554, "standard_gw": 4.3631307780742645, "standard_gx": 4.404045641422272, "rowwise_fwd": 2.9063820838928223, "rowwise_bwd": 2.094462513923645, "global_fwd": 2.8426870703697205, "global_bwd": 2.0792782306671143, "x_quantize_rowwise": 0.33241137862205505, "g_quantize_rowwise": 1.1817105114459991, "w_quantize_rowwise": 0.03374367952346802, "w_quantize_colwise_transpose": 0.08633732795715332, "w_quantize_global": 0.09231641888618469, "w_quantize_global_transpose": 0.100012868642807, "time_standard": 13.62539455294609, "time_rowwise": 10.998178273439407, "time_global": 10.991547256708145} {"repeat": 64, "batch_size": 131072, "dim_out": 1024, "dim_in": 4096, "wm": 4, "switch": true, "standard_fwd": 4.246581345796585, "standard_gw": 4.42587211728096, "standard_gx": 4.581417888402939, "rowwise_fwd": 2.1114833652973175, "rowwise_bwd": 2.9050447046756744, "global_fwd": 2.0806826651096344, "global_bwd": 2.85966694355011, "x_quantize_rowwise": 1.1816024780273438, "g_quantize_rowwise": 0.33330172300338745, "w_quantize_rowwise": 0.033445656299591064, "w_quantize_colwise_transpose": 0.09065866470336914, "w_quantize_global": 0.09239837527275085, "w_quantize_global_transpose": 0.09984523057937622, "time_standard": 13.253871351480484, "time_rowwise": 11.081408709287643, "time_global": 11.073369532823563} {"repeat": 64, "batch_size": 8192, "dim_out": 5120, "dim_in": 1280, "wm": 4, "switch": false, "standard_fwd": 0.4859529435634613, "standard_gw": 0.46338513493537903, "standard_gx": 0.42321905493736267, "rowwise_fwd": 0.2761557698249817, "rowwise_bwd": 0.20775198936462402, "global_fwd": 0.2713911235332489, "global_bwd": 0.20639970898628235, "x_quantize_rowwise": 0.033095479011535645, "g_quantize_rowwise": 0.11894106864929199, "w_quantize_rowwise": 0.03125518560409546, "w_quantize_colwise_transpose": 0.1424551010131836, "w_quantize_global": 0.07288157939910889, "w_quantize_global_transpose": 0.08071959018707275, "time_standard": 1.372557133436203, "time_rowwise": 1.2730397284030914, "time_global": 1.2468136847019196} {"repeat": 64, "batch_size": 8192, "dim_out": 1280, "dim_in": 5120, "wm": 4, "switch": true, "standard_fwd": 0.3920421004295349, "standard_gw": 0.44424086809158325, "standard_gx": 0.4759356379508972, "rowwise_fwd": 0.23231282830238342, "rowwise_bwd": 0.28430670499801636, "global_fwd": 0.20883232355117798, "global_bwd": 0.2741999924182892, "x_quantize_rowwise": 0.12018159031867981, "g_quantize_rowwise": 0.03195926547050476, "w_quantize_rowwise": 0.026017427444458008, "w_quantize_colwise_transpose": 0.14733895659446716, "w_quantize_global": 0.07734447717666626, "w_quantize_global_transpose": 0.0788569450378418, "time_standard": 1.3122186064720154, "time_rowwise": 1.2863576412200928, "time_global": 1.235615462064743} {"repeat": 64, "batch_size": 16384, "dim_out": 5120, "dim_in": 1280, "wm": 4, "switch": false, "standard_fwd": 1.0111741721630096, "standard_gw": 0.9267590939998627, "standard_gx": 0.8254274725914001, "rowwise_fwd": 0.5434826016426086, "rowwise_bwd": 0.4077926278114319, "global_fwd": 0.5318708717823029, "global_bwd": 0.40537863969802856, "x_quantize_rowwise": 0.059738755226135254, "g_quantize_rowwise": 0.2299174666404724, "w_quantize_rowwise": 0.02545863389968872, "w_quantize_colwise_transpose": 0.14269724488258362, "w_quantize_global": 0.07300823926925659, "w_quantize_global_transpose": 0.07878988981246948, "time_standard": 2.7633607387542725, "time_rowwise": 2.335846424102783, "time_global": 2.305462956428528} {"repeat": 64, "batch_size": 16384, "dim_out": 1280, "dim_in": 5120, "wm": 4, "switch": true, "standard_fwd": 0.8095316588878632, "standard_gw": 0.8607134222984314, "standard_gx": 0.9204968810081482, "rowwise_fwd": 0.4275888204574585, "rowwise_bwd": 0.5485899746417999, "global_fwd": 0.41000545024871826, "global_bwd": 0.5317628383636475, "x_quantize_rowwise": 0.2301819622516632, "g_quantize_rowwise": 0.059254467487335205, "w_quantize_rowwise": 0.02466142177581787, "w_quantize_colwise_transpose": 0.14865398406982422, "w_quantize_global": 0.07582828402519226, "w_quantize_global_transpose": 0.08231401443481445, "time_standard": 2.5907419621944427, "time_rowwise": 2.2996440529823303, "time_global": 2.2500604391098022} {"repeat": 64, "batch_size": 32768, "dim_out": 5120, "dim_in": 1280, "wm": 4, "switch": false, "standard_fwd": 2.0658522844314575, "standard_gw": 1.718364655971527, "standard_gx": 1.6660578548908234, "rowwise_fwd": 1.066897064447403, "rowwise_bwd": 0.8070804178714752, "global_fwd": 1.0473169386386871, "global_bwd": 0.8021742105484009, "x_quantize_rowwise": 0.11274218559265137, "g_quantize_rowwise": 0.4518181085586548, "w_quantize_rowwise": 0.026501715183258057, "w_quantize_colwise_transpose": 0.14259666204452515, "w_quantize_global": 0.07484853267669678, "w_quantize_global_transpose": 0.07976219058036804, "time_standard": 5.450274795293808, "time_rowwise": 4.326000809669495, "time_global": 4.287026822566986} {"repeat": 64, "batch_size": 32768, "dim_out": 1280, "dim_in": 5120, "wm": 4, "switch": true, "standard_fwd": 2.7549192309379578, "standard_gw": 1.6954988241195679, "standard_gx": 1.8179528415203094, "rowwise_fwd": 0.8649080991744995, "rowwise_bwd": 1.0746456682682037, "global_fwd": 0.8023083209991455, "global_bwd": 1.0471977293491364, "x_quantize_rowwise": 0.45225024223327637, "g_quantize_rowwise": 0.11286512017250061, "w_quantize_rowwise": 0.0252649188041687, "w_quantize_colwise_transpose": 0.14732033014297485, "w_quantize_global": 0.07537379860877991, "w_quantize_global_transpose": 0.0807642936706543, "time_standard": 6.268370896577835, "time_rowwise": 4.372753202915192, "time_global": 4.266258329153061} {"repeat": 64, "batch_size": 65536, "dim_out": 5120, "dim_in": 1280, "wm": 4, "switch": false, "standard_fwd": 4.098430275917053, "standard_gw": 3.3501461148262024, "standard_gx": 5.560480058193207, "rowwise_fwd": 2.112947404384613, "rowwise_bwd": 1.605246216058731, "global_fwd": 2.0697638392448425, "global_bwd": 1.5953518450260162, "x_quantize_rowwise": 0.21921470761299133, "g_quantize_rowwise": 0.8956789970397949, "w_quantize_rowwise": 0.02710893750190735, "w_quantize_colwise_transpose": 0.14268234372138977, "w_quantize_global": 0.07259473204612732, "w_quantize_global_transpose": 0.07899105548858643, "time_standard": 13.009056448936462, "time_rowwise": 8.35302472114563, "time_global": 8.281741291284561} {"repeat": 64, "batch_size": 65536, "dim_out": 1280, "dim_in": 5120, "wm": 4, "switch": true, "standard_fwd": 5.586959421634674, "standard_gw": 3.358360379934311, "standard_gx": 3.6434978246688843, "rowwise_fwd": 1.6269534826278687, "rowwise_bwd": 2.128206193447113, "global_fwd": 1.5950687229633331, "global_bwd": 2.0831897854804993, "x_quantize_rowwise": 0.8954145014286041, "g_quantize_rowwise": 0.21914392709732056, "w_quantize_rowwise": 0.026203691959381104, "w_quantize_colwise_transpose": 0.14658644795417786, "w_quantize_global": 0.07478520274162292, "w_quantize_global_transpose": 0.07964670658111572, "time_standard": 12.58881762623787, "time_rowwise": 8.400868624448776, "time_global": 8.305609226226807} {"repeat": 64, "batch_size": 131072, "dim_out": 5120, "dim_in": 1280, "wm": 4, "switch": false, "standard_fwd": 8.229725062847137, "standard_gw": 6.791356950998306, "standard_gx": 6.806455552577972, "rowwise_fwd": 4.252471029758453, "rowwise_bwd": 3.2062679529190063, "global_fwd": 4.175614565610886, "global_bwd": 3.1837262213230133, "x_quantize_rowwise": 0.4321373999118805, "g_quantize_rowwise": 1.787092536687851, "w_quantize_rowwise": 0.0270158052444458, "w_quantize_colwise_transpose": 0.1424252986907959, "w_quantize_global": 0.07348507642745972, "w_quantize_global_transpose": 0.07829815149307251, "time_standard": 21.827537566423416, "time_rowwise": 16.63876697421074, "time_global": 16.52171090245247} {"repeat": 64, "batch_size": 131072, "dim_out": 1280, "dim_in": 5120, "wm": 4, "switch": true, "standard_fwd": 11.279478669166565, "standard_gw": 6.7345499992370605, "standard_gx": 7.206875830888748, "rowwise_fwd": 3.209315240383148, "rowwise_bwd": 4.256397485733032, "global_fwd": 3.180190920829773, "global_bwd": 4.177983850240707, "x_quantize_rowwise": 1.7836056649684906, "g_quantize_rowwise": 0.4321075975894928, "w_quantize_rowwise": 0.03205239772796631, "w_quantize_colwise_transpose": 0.14675036072731018, "w_quantize_global": 0.09316205978393555, "w_quantize_global_transpose": 0.10086596012115479, "time_standard": 25.220904499292374, "time_rowwise": 16.5947787463665, "time_global": 16.502466052770615} {"repeat": 64, "batch_size": 8192, "dim_out": 5632, "dim_in": 1408, "wm": 4, "switch": false, "standard_fwd": 0.5776733160018921, "standard_gw": 0.5300231277942657, "standard_gx": 0.6005913019180298, "rowwise_fwd": 0.33330172300338745, "rowwise_bwd": 0.2957060933113098, "global_fwd": 0.32876431941986084, "global_bwd": 0.29108673334121704, "x_quantize_rowwise": 0.03466755151748657, "g_quantize_rowwise": 0.12264400720596313, "w_quantize_rowwise": 0.033874064683914185, "w_quantize_colwise_transpose": 0.1775398850440979, "w_quantize_global": 0.09503215551376343, "w_quantize_global_transpose": 0.10617449879646301, "time_standard": 1.7082877457141876, "time_rowwise": 1.5277564525604248, "time_global": 1.5083923935890198} {"repeat": 64, "batch_size": 8192, "dim_out": 1408, "dim_in": 5632, "wm": 4, "switch": true, "standard_fwd": 0.5164109170436859, "standard_gw": 0.5367249250411987, "standard_gx": 0.5876161158084869, "rowwise_fwd": 0.3132447600364685, "rowwise_bwd": 0.3396235406398773, "global_fwd": 0.2943649888038635, "global_bwd": 0.33209100365638733, "x_quantize_rowwise": 0.12357160449028015, "g_quantize_rowwise": 0.035997480154037476, "w_quantize_rowwise": 0.03213062882423401, "w_quantize_colwise_transpose": 0.17676874995231628, "w_quantize_global": 0.09861215949058533, "w_quantize_global_transpose": 0.0998862087726593, "time_standard": 1.6407519578933716, "time_rowwise": 1.5580616891384125, "time_global": 1.5212483704090118} {"repeat": 64, "batch_size": 16384, "dim_out": 5632, "dim_in": 1408, "wm": 4, "switch": false, "standard_fwd": 1.2096501886844635, "standard_gw": 1.0663382709026337, "standard_gx": 1.0961703956127167, "rowwise_fwd": 0.6396733224391937, "rowwise_bwd": 0.5173943936824799, "global_fwd": 0.6296299397945404, "global_bwd": 0.5130060017108917, "x_quantize_rowwise": 0.06211921572685242, "g_quantize_rowwise": 0.2361498773097992, "w_quantize_rowwise": 0.03260001540184021, "w_quantize_colwise_transpose": 0.17679482698440552, "w_quantize_global": 0.09361281991004944, "w_quantize_global_transpose": 0.09913742542266846, "time_standard": 3.372158855199814, "time_rowwise": 2.7310699224472046, "time_global": 2.6999935507774353} {"repeat": 64, "batch_size": 16384, "dim_out": 1408, "dim_in": 5632, "wm": 4, "switch": true, "standard_fwd": 1.1065565049648285, "standard_gw": 1.0664314031600952, "standard_gx": 1.1266544461250305, "rowwise_fwd": 0.5352050065994263, "rowwise_bwd": 0.6464086472988129, "global_fwd": 0.513765960931778, "global_bwd": 0.6284862756729126, "x_quantize_rowwise": 0.23620948195457458, "g_quantize_rowwise": 0.062271952629089355, "w_quantize_rowwise": 0.031460076570510864, "w_quantize_colwise_transpose": 0.17675384879112244, "w_quantize_global": 0.09486451745033264, "w_quantize_global_transpose": 0.09898096323013306, "time_standard": 3.2996423542499542, "time_rowwise": 2.7547404170036316, "time_global": 2.7010105550289154} {"repeat": 64, "batch_size": 32768, "dim_out": 5632, "dim_in": 1408, "wm": 4, "switch": false, "standard_fwd": 2.4367496371269226, "standard_gw": 2.0806193351745605, "standard_gx": 2.19624862074852, "rowwise_fwd": 1.2554042041301727, "rowwise_bwd": 1.0227933526039124, "global_fwd": 1.2322552502155304, "global_bwd": 1.0152235627174377, "x_quantize_rowwise": 0.11792033910751343, "g_quantize_rowwise": 0.4639364778995514, "w_quantize_rowwise": 0.03241002559661865, "w_quantize_colwise_transpose": 0.17657503485679626, "w_quantize_global": 0.09655207395553589, "w_quantize_global_transpose": 0.09958073496818542, "time_standard": 6.713617593050003, "time_rowwise": 5.149658769369125, "time_global": 5.106087774038315} {"repeat": 64, "batch_size": 32768, "dim_out": 1408, "dim_in": 5632, "wm": 4, "switch": true, "standard_fwd": 2.1935217082500458, "standard_gw": 2.0055584609508514, "standard_gx": 2.1882541477680206, "rowwise_fwd": 1.0396353900432587, "rowwise_bwd": 1.2542344629764557, "global_fwd": 1.0161921381950378, "global_bwd": 1.233428716659546, "x_quantize_rowwise": 0.4642195999622345, "g_quantize_rowwise": 0.11782720685005188, "w_quantize_rowwise": 0.033117830753326416, "w_quantize_colwise_transpose": 0.17696991562843323, "w_quantize_global": 0.09416043758392334, "w_quantize_global_transpose": 0.10101497173309326, "time_standard": 6.387334316968918, "time_rowwise": 5.091562867164612, "time_global": 5.032401531934738} {"repeat": 64, "batch_size": 65536, "dim_out": 5632, "dim_in": 1408, "wm": 4, "switch": false, "standard_fwd": 4.804681986570358, "standard_gw": 4.763372242450714, "standard_gx": 4.064023494720459, "rowwise_fwd": 2.484843134880066, "rowwise_bwd": 1.9691288471221924, "global_fwd": 2.441786229610443, "global_bwd": 1.9574686884880066, "x_quantize_rowwise": 0.2294592559337616, "g_quantize_rowwise": 0.9196549654006958, "w_quantize_rowwise": 0.0313781201839447, "w_quantize_colwise_transpose": 0.1768544316291809, "w_quantize_global": 0.09644776582717896, "w_quantize_global_transpose": 0.09847059845924377, "time_standard": 13.632077723741531, "time_rowwise": 10.574690997600555, "time_global": 10.506659746170044} {"repeat": 64, "batch_size": 65536, "dim_out": 1408, "dim_in": 5632, "wm": 4, "switch": true, "standard_fwd": 4.0907710790634155, "standard_gw": 3.9793066680431366, "standard_gx": 4.302978515625, "rowwise_fwd": 1.992940902709961, "rowwise_bwd": 2.4996213614940643, "global_fwd": 1.9551962614059448, "global_bwd": 2.457551658153534, "x_quantize_rowwise": 0.9200014173984528, "g_quantize_rowwise": 0.2293996512889862, "w_quantize_rowwise": 0.0313781201839447, "w_quantize_colwise_transpose": 0.17882883548736572, "w_quantize_global": 0.09540095925331116, "w_quantize_global_transpose": 0.09880587458610535, "time_standard": 12.373056262731552, "time_rowwise": 9.831476956605911, "time_global": 9.73566249012947} {"repeat": 64, "batch_size": 131072, "dim_out": 5632, "dim_in": 1408, "wm": 4, "switch": false, "standard_fwd": 9.655728936195374, "standard_gw": 8.261296898126602, "standard_gx": 8.064884692430496, "rowwise_fwd": 5.007706582546234, "rowwise_bwd": 3.8615092635154724, "global_fwd": 4.920527338981628, "global_bwd": 3.8330331444740295, "x_quantize_rowwise": 0.45276060700416565, "g_quantize_rowwise": 1.8306002020835876, "w_quantize_rowwise": 0.031366944313049316, "w_quantize_colwise_transpose": 0.1766495406627655, "w_quantize_global": 0.09412690997123718, "w_quantize_global_transpose": 0.09780004620552063, "time_standard": 25.981910526752472, "time_rowwise": 19.621890038251877, "time_global": 19.49014514684677} {"repeat": 64, "batch_size": 131072, "dim_out": 1408, "dim_in": 5632, "wm": 4, "switch": true, "standard_fwd": 8.033104240894318, "standard_gw": 8.2889124751091, "standard_gx": 8.622754365205765, "rowwise_fwd": 3.8747042417526245, "rowwise_bwd": 5.003921687602997, "global_fwd": 3.8315393030643463, "global_bwd": 4.9162134528160095, "x_quantize_rowwise": 1.8304847180843353, "g_quantize_rowwise": 0.4522763192653656, "w_quantize_rowwise": 0.03413110971450806, "w_quantize_colwise_transpose": 0.1771189272403717, "w_quantize_global": 0.09519979357719421, "w_quantize_global_transpose": 0.09930506348609924, "time_standard": 24.944771081209183, "time_rowwise": 19.661549478769302, "time_global": 19.51393112540245} {"repeat": 64, "batch_size": 8192, "dim_out": 6656, "dim_in": 1664, "wm": 4, "switch": false, "standard_fwd": 0.7954612374305725, "standard_gw": 0.7456131279468536, "standard_gx": 0.8799619972705841, "rowwise_fwd": 0.43267011642456055, "rowwise_bwd": 0.34622475504875183, "global_fwd": 0.42615458369255066, "global_bwd": 0.344250351190567, "x_quantize_rowwise": 0.03748014569282532, "g_quantize_rowwise": 0.13304129242897034, "w_quantize_rowwise": 0.03294646739959717, "w_quantize_colwise_transpose": 0.2407953143119812, "w_quantize_global": 0.094633549451828, "w_quantize_global_transpose": 0.10305643081665039, "time_standard": 2.4210363626480103, "time_rowwise": 1.96877121925354, "time_global": 1.8842294812202454} {"repeat": 64, "batch_size": 8192, "dim_out": 1664, "dim_in": 6656, "wm": 4, "switch": true, "standard_fwd": 0.7120333611965179, "standard_gw": 0.7622130215167999, "standard_gx": 0.8262209594249725, "rowwise_fwd": 0.3702230751514435, "rowwise_bwd": 0.4419572651386261, "global_fwd": 0.3479123115539551, "global_bwd": 0.4306286573410034, "x_quantize_rowwise": 0.13308599591255188, "g_quantize_rowwise": 0.037495046854019165, "w_quantize_rowwise": 0.03398209810256958, "w_quantize_colwise_transpose": 0.23782625794410706, "w_quantize_global": 0.09853765368461609, "w_quantize_global_transpose": 0.10247156023979187, "time_standard": 2.3004673421382904, "time_rowwise": 2.016782760620117, "time_global": 1.9123442471027374} {"repeat": 64, "batch_size": 16384, "dim_out": 6656, "dim_in": 1664, "wm": 4, "switch": false, "standard_fwd": 1.6292817890644073, "standard_gw": 1.5109702944755554, "standard_gx": 1.482747495174408, "rowwise_fwd": 0.8386112749576569, "rowwise_bwd": 0.6844550371170044, "global_fwd": 0.8220970630645752, "global_bwd": 0.6802082061767578, "x_quantize_rowwise": 0.06883963942527771, "g_quantize_rowwise": 0.25641173124313354, "w_quantize_rowwise": 0.033054500818252563, "w_quantize_colwise_transpose": 0.24027004837989807, "w_quantize_global": 0.0967271625995636, "w_quantize_global_transpose": 0.102948397397995, "time_standard": 4.622999578714371, "time_rowwise": 3.6326125264167786, "time_global": 3.5382024943828583} {"repeat": 64, "batch_size": 16384, "dim_out": 1664, "dim_in": 6656, "wm": 4, "switch": true, "standard_fwd": 1.4877021312713623, "standard_gw": 1.5015341341495514, "standard_gx": 1.529306173324585, "rowwise_fwd": 0.715944916009903, "rowwise_bwd": 0.8529908955097198, "global_fwd": 0.680088996887207, "global_bwd": 0.8224695920944214, "x_quantize_rowwise": 0.2568177878856659, "g_quantize_rowwise": 0.06864592432975769, "w_quantize_rowwise": 0.03343448042869568, "w_quantize_colwise_transpose": 0.23645907640457153, "w_quantize_global": 0.09399279952049255, "w_quantize_global_transpose": 0.10286271572113037, "time_standard": 4.518542438745499, "time_rowwise": 3.665827214717865, "time_global": 3.5264119505882263} {"repeat": 64, "batch_size": 32768, "dim_out": 6656, "dim_in": 1664, "wm": 4, "switch": false, "standard_fwd": 3.261040896177292, "standard_gw": 2.8816498816013336, "standard_gx": 2.8357282280921936, "rowwise_fwd": 1.6594752669334412, "rowwise_bwd": 1.359265297651291, "global_fwd": 1.6287527978420258, "global_bwd": 1.3503879308700562, "x_quantize_rowwise": 0.13146549463272095, "g_quantize_rowwise": 0.5035959184169769, "w_quantize_rowwise": 0.03438442945480347, "w_quantize_colwise_transpose": 0.24086236953735352, "w_quantize_global": 0.0945068895816803, "w_quantize_global_transpose": 0.10332837700843811, "time_standard": 8.978419005870819, "time_rowwise": 6.8106986582279205, "time_global": 6.693687289953232} {"repeat": 64, "batch_size": 32768, "dim_out": 1664, "dim_in": 6656, "wm": 4, "switch": true, "standard_fwd": 2.848360687494278, "standard_gw": 2.8955675661563873, "standard_gx": 3.0499882996082306, "rowwise_fwd": 1.3900883495807648, "rowwise_bwd": 1.6595833003520966, "global_fwd": 1.3514049351215363, "global_bwd": 1.629263162612915, "x_quantize_rowwise": 0.5036592483520508, "g_quantize_rowwise": 0.13118237257003784, "w_quantize_rowwise": 0.03438442945480347, "w_quantize_colwise_transpose": 0.23709610104560852, "w_quantize_global": 0.0951625406742096, "w_quantize_global_transpose": 0.10216236114501953, "time_standard": 8.793916553258896, "time_rowwise": 6.851561367511749, "time_global": 6.708402186632156} {"repeat": 64, "batch_size": 65536, "dim_out": 6656, "dim_in": 1664, "wm": 4, "switch": false, "standard_fwd": 6.4978525042533875, "standard_gw": 6.462603807449341, "standard_gx": 5.5987648665905, "rowwise_fwd": 3.2996535301208496, "rowwise_bwd": 2.6320070028305054, "global_fwd": 3.2426007091999054, "global_bwd": 2.612769603729248, "x_quantize_rowwise": 0.2561397850513458, "g_quantize_rowwise": 0.9984448552131653, "w_quantize_rowwise": 0.033076852560043335, "w_quantize_colwise_transpose": 0.24232640862464905, "w_quantize_global": 0.09618699550628662, "w_quantize_global_transpose": 0.10257214307785034, "time_standard": 18.559221178293228, "time_rowwise": 13.9242522418499, "time_global": 13.771317899227142} {"repeat": 64, "batch_size": 65536, "dim_out": 1664, "dim_in": 6656, "wm": 4, "switch": true, "standard_fwd": 5.5702440440654755, "standard_gw": 5.717620253562927, "standard_gx": 6.08203187584877, "rowwise_fwd": 2.649586647748947, "rowwise_bwd": 3.315173089504242, "global_fwd": 2.6132799685001373, "global_bwd": 3.257807344198227, "x_quantize_rowwise": 0.9980201721191406, "g_quantize_rowwise": 0.256560742855072, "w_quantize_rowwise": 0.03356859087944031, "w_quantize_colwise_transpose": 0.23729726672172546, "w_quantize_global": 0.09495764970779419, "w_quantize_global_transpose": 0.103779137134552, "time_standard": 17.369896173477173, "time_rowwise": 13.207826763391495, "time_global": 13.04202526807785} {"repeat": 64, "batch_size": 131072, "dim_out": 6656, "dim_in": 1664, "wm": 4, "switch": false, "standard_fwd": 13.058379292488098, "standard_gw": 11.480242013931274, "standard_gx": 11.092845350503922, "rowwise_fwd": 6.637874990701675, "rowwise_bwd": 5.24790957570076, "global_fwd": 6.521012634038925, "global_bwd": 5.214303731918335, "x_quantize_rowwise": 0.5057565867900848, "g_quantize_rowwise": 1.989319920539856, "w_quantize_rowwise": 0.03439188003540039, "w_quantize_colwise_transpose": 0.24280324578285217, "w_quantize_global": 0.09520724415779114, "w_quantize_global_transpose": 0.10240450501441956, "time_standard": 35.631466656923294, "time_rowwise": 26.138298213481903, "time_global": 25.908246636390686} {"repeat": 64, "batch_size": 131072, "dim_out": 1664, "dim_in": 6656, "wm": 4, "switch": true, "standard_fwd": 11.13397628068924, "standard_gw": 11.371888220310211, "standard_gx": 12.12756335735321, "rowwise_fwd": 5.2495077252388, "rowwise_bwd": 6.638709455728531, "global_fwd": 5.215313285589218, "global_bwd": 6.5222084522247314, "x_quantize_rowwise": 1.9870512187480927, "g_quantize_rowwise": 0.5058236420154572, "w_quantize_rowwise": 0.034634023904800415, "w_quantize_colwise_transpose": 0.23674964904785156, "w_quantize_global": 0.09457767009735107, "w_quantize_global_transpose": 0.10183081030845642, "time_standard": 34.63342785835266, "time_rowwise": 26.024363934993744, "time_global": 25.798693299293518} {"repeat": 64, "batch_size": 8192, "dim_out": 8192, "dim_in": 2048, "wm": 4, "switch": false, "standard_fwd": 1.2125298380851746, "standard_gw": 1.1111274361610413, "standard_gx": 1.0840706527233124, "rowwise_fwd": 0.6057210266590118, "rowwise_bwd": 0.51865354180336, "global_fwd": 0.5952082574367523, "global_bwd": 0.5167685449123383, "x_quantize_rowwise": 0.045686960220336914, "g_quantize_rowwise": 0.15827640891075134, "w_quantize_rowwise": 0.04361197352409363, "w_quantize_colwise_transpose": 0.34067779779434204, "w_quantize_global": 0.13644620776176453, "w_quantize_global_transpose": 0.14925003051757812, "time_standard": 3.407727926969528, "time_rowwise": 2.823755145072937, "time_global": 2.7127638459205627} {"repeat": 64, "batch_size": 8192, "dim_out": 2048, "dim_in": 8192, "wm": 4, "switch": true, "standard_fwd": 1.0731369256973267, "standard_gw": 1.1365897953510284, "standard_gx": 1.1498592793941498, "rowwise_fwd": 0.5573518574237823, "rowwise_bwd": 0.615488737821579, "global_fwd": 0.5220361053943634, "global_bwd": 0.5939789116382599, "x_quantize_rowwise": 0.15765801072120667, "g_quantize_rowwise": 0.04369020462036133, "w_quantize_rowwise": 0.047359615564346313, "w_quantize_colwise_transpose": 0.5526281893253326, "w_quantize_global": 0.13606995344161987, "w_quantize_global_transpose": 0.15017390251159668, "time_standard": 3.359586000442505, "time_rowwise": 3.1107664108276367, "time_global": 2.7401968836784363} {"repeat": 64, "batch_size": 16384, "dim_out": 8192, "dim_in": 2048, "wm": 4, "switch": false, "standard_fwd": 2.4274885654449463, "standard_gw": 2.1799951791763306, "standard_gx": 2.1426528692245483, "rowwise_fwd": 1.195710152387619, "rowwise_bwd": 1.027170568704605, "global_fwd": 1.1747106909751892, "global_bwd": 1.0251589119434357, "x_quantize_rowwise": 0.08098781108856201, "g_quantize_rowwise": 0.3052949905395508, "w_quantize_rowwise": 0.043764710426330566, "w_quantize_colwise_transpose": 0.33987686038017273, "w_quantize_global": 0.13646483421325684, "w_quantize_global_transpose": 0.14739856123924255, "time_standard": 6.750136613845825, "time_rowwise": 5.172800272703171, "time_global": 5.050010979175568} {"repeat": 64, "batch_size": 16384, "dim_out": 2048, "dim_in": 8192, "wm": 4, "switch": true, "standard_fwd": 2.1661892533302307, "standard_gw": 2.0948275923728943, "standard_gx": 2.306375652551651, "rowwise_fwd": 1.0587647557258606, "rowwise_bwd": 1.1999905109405518, "global_fwd": 1.0296404361724854, "global_bwd": 1.1749230325222015, "x_quantize_rowwise": 0.3054030239582062, "g_quantize_rowwise": 0.08077546954154968, "w_quantize_rowwise": 0.047225505113601685, "w_quantize_colwise_transpose": 0.600133091211319, "w_quantize_global": 0.13613328337669373, "w_quantize_global_transpose": 0.1484006643295288, "time_standard": 6.567392498254776, "time_rowwise": 5.387119948863983, "time_global": 4.97010350227356} {"repeat": 64, "batch_size": 32768, "dim_out": 8192, "dim_in": 2048, "wm": 4, "switch": false, "standard_fwd": 4.807606339454651, "standard_gw": 4.170913249254227, "standard_gx": 4.117622971534729, "rowwise_fwd": 2.370934933423996, "rowwise_bwd": 1.9481778144836426, "global_fwd": 2.3383721709251404, "global_bwd": 1.9443817436695099, "x_quantize_rowwise": 0.1547597348690033, "g_quantize_rowwise": 0.6000511348247528, "w_quantize_rowwise": 0.04361942410469055, "w_quantize_colwise_transpose": 0.3403201699256897, "w_quantize_global": 0.13600289821624756, "w_quantize_global_transpose": 0.1474134624004364, "time_standard": 13.096142560243607, "time_rowwise": 9.628776460886002, "time_global": 9.491894394159317} {"repeat": 64, "batch_size": 32768, "dim_out": 2048, "dim_in": 8192, "wm": 4, "switch": true, "standard_fwd": 4.1619837284088135, "standard_gw": 4.181284457445145, "standard_gx": 4.635505378246307, "rowwise_fwd": 1.9684135913848877, "rowwise_bwd": 2.3750364780426025, "global_fwd": 1.9445866346359253, "global_bwd": 2.3551955819129944, "x_quantize_rowwise": 0.6004162132740021, "g_quantize_rowwise": 0.15468522906303406, "w_quantize_rowwise": 0.04730746150016785, "w_quantize_colwise_transpose": 0.5999617278575897, "w_quantize_global": 0.1364201307296753, "w_quantize_global_transpose": 0.14847144484519958, "time_standard": 12.978773564100266, "time_rowwise": 9.927105158567429, "time_global": 9.521059691905975} {"repeat": 64, "batch_size": 65536, "dim_out": 8192, "dim_in": 2048, "wm": 4, "switch": false, "standard_fwd": 9.52371209859848, "standard_gw": 8.354485034942627, "standard_gx": 8.69860127568245, "rowwise_fwd": 4.717472940683365, "rowwise_bwd": 3.8843750953674316, "global_fwd": 4.645414650440216, "global_bwd": 3.8761012256145477, "x_quantize_rowwise": 0.3024861216545105, "g_quantize_rowwise": 1.1897757649421692, "w_quantize_rowwise": 0.04366785287857056, "w_quantize_colwise_transpose": 0.33988431096076965, "w_quantize_global": 0.1359507441520691, "w_quantize_global_transpose": 0.14724582433700562, "time_standard": 26.576798409223557, "time_rowwise": 18.832147121429443, "time_global": 18.651459366083145} {"repeat": 64, "batch_size": 65536, "dim_out": 2048, "dim_in": 8192, "wm": 4, "switch": true, "standard_fwd": 8.307881653308868, "standard_gw": 8.214320987462997, "standard_gx": 9.21182706952095, "rowwise_fwd": 3.8919784128665924, "rowwise_bwd": 4.72346693277359, "global_fwd": 3.8761794567108154, "global_bwd": 4.673641175031662, "x_quantize_rowwise": 1.1893920600414276, "g_quantize_rowwise": 0.3024972975254059, "w_quantize_rowwise": 0.04708021879196167, "w_quantize_colwise_transpose": 0.6039328873157501, "w_quantize_global": 0.13624504208564758, "w_quantize_global_transpose": 0.14867261052131653, "time_standard": 25.734029710292816, "time_rowwise": 18.972668796777725, "time_global": 18.540948629379272} {"repeat": 64, "batch_size": 131072, "dim_out": 8192, "dim_in": 2048, "wm": 4, "switch": false, "standard_fwd": 19.30372044444084, "standard_gw": 16.480475664138794, "standard_gx": 17.61433482170105, "rowwise_fwd": 9.49602946639061, "rowwise_bwd": 7.768530398607254, "global_fwd": 9.3533955514431, "global_bwd": 7.749464362859726, "x_quantize_rowwise": 0.5977451801300049, "g_quantize_rowwise": 2.3684948682785034, "w_quantize_rowwise": 0.04375725984573364, "w_quantize_colwise_transpose": 0.34042075276374817, "w_quantize_global": 0.13628974556922913, "w_quantize_global_transpose": 0.14671683311462402, "time_standard": 53.398530930280685, "time_rowwise": 37.09545359015465, "time_global": 36.83258220553398} {"repeat": 64, "batch_size": 131072, "dim_out": 2048, "dim_in": 8192, "wm": 4, "switch": true, "standard_fwd": 18.041003495454788, "standard_gw": 17.770148813724518, "standard_gx": 17.70009845495224, "rowwise_fwd": 7.756810635328293, "rowwise_bwd": 9.502101689577103, "global_fwd": 7.7384114265441895, "global_bwd": 9.36170294880867, "x_quantize_rowwise": 2.3686252534389496, "g_quantize_rowwise": 0.5980581045150757, "w_quantize_rowwise": 0.04723668098449707, "w_quantize_colwise_transpose": 0.6035342812538147, "w_quantize_global": 0.13603642582893372, "w_quantize_global_transpose": 0.1485198736190796, "time_standard": 53.511250764131546, "time_rowwise": 38.64651545882225, "time_global": 38.121502846479416} {"repeat": 64, "batch_size": 8192, "dim_out": 16384, "dim_in": 4096, "wm": 4, "switch": false, "standard_fwd": 4.598241299390793, "standard_gw": 4.294309765100479, "standard_gx": 4.261095076799393, "rowwise_fwd": 2.0976848900318146, "rowwise_bwd": 1.9718967378139496, "global_fwd": 2.0763762295246124, "global_bwd": 1.9703581929206848, "x_quantize_rowwise": 0.08216872811317444, "g_quantize_rowwise": 0.4405900835990906, "w_quantize_rowwise": 0.1553371548652649, "w_quantize_colwise_transpose": 1.6110725700855255, "w_quantize_global": 0.481240451335907, "w_quantize_global_transpose": 0.5061514675617218, "time_standard": 13.153646141290665, "time_rowwise": 10.653059929609299, "time_global": 9.85119491815567} {"repeat": 64, "batch_size": 8192, "dim_out": 4096, "dim_in": 16384, "wm": 4, "switch": true, "standard_fwd": 4.35885414481163, "standard_gw": 4.29583340883255, "standard_gx": 4.5370906591415405, "rowwise_fwd": 2.0015686750411987, "rowwise_bwd": 2.097565680742264, "global_fwd": 1.969795674085617, "global_bwd": 2.075403928756714, "x_quantize_rowwise": 0.43984130024909973, "g_quantize_rowwise": 0.08216127753257751, "w_quantize_rowwise": 0.22544339299201965, "w_quantize_colwise_transpose": 2.4342015385627747, "w_quantize_global": 0.48087164759635925, "w_quantize_global_transpose": 0.5099289119243622, "time_standard": 13.19177821278572, "time_rowwise": 11.576615273952484, "time_global": 9.85383614897728} {"repeat": 64, "batch_size": 16384, "dim_out": 16384, "dim_in": 4096, "wm": 4, "switch": false, "standard_fwd": 9.09888744354248, "standard_gw": 8.230950683355331, "standard_gx": 8.465446531772614, "rowwise_fwd": 4.182614386081696, "rowwise_bwd": 3.747660666704178, "global_fwd": 4.138719290494919, "global_bwd": 3.74777615070343, "x_quantize_rowwise": 0.15515834093093872, "g_quantize_rowwise": 0.8699297904968262, "w_quantize_rowwise": 0.15544891357421875, "w_quantize_colwise_transpose": 1.6132444143295288, "w_quantize_global": 0.48100948333740234, "w_quantize_global_transpose": 0.5051903426647186, "time_standard": 25.795284658670425, "time_rowwise": 18.955007195472717, "time_global": 18.128734081983566} {"repeat": 64, "batch_size": 16384, "dim_out": 4096, "dim_in": 16384, "wm": 4, "switch": true, "standard_fwd": 8.378107100725174, "standard_gw": 8.923027664422989, "standard_gx": 9.049762040376663, "rowwise_fwd": 3.765825182199478, "rowwise_bwd": 4.183519631624222, "global_fwd": 3.744799643754959, "global_bwd": 4.1590481996536255, "x_quantize_rowwise": 0.8693933486938477, "g_quantize_rowwise": 0.1553073525428772, "w_quantize_rowwise": 0.2258792519569397, "w_quantize_colwise_transpose": 2.4386271834373474, "w_quantize_global": 0.4811100661754608, "w_quantize_global_transpose": 0.5102269351482391, "time_standard": 26.350896805524826, "time_rowwise": 20.5615796148777, "time_global": 18.842913210392} {"repeat": 64, "batch_size": 32768, "dim_out": 16384, "dim_in": 4096, "wm": 4, "switch": false, "standard_fwd": 18.266115337610245, "standard_gw": 17.671160399913788, "standard_gx": 17.10302010178566, "rowwise_fwd": 8.347474038600922, "rowwise_bwd": 7.514089345932007, "global_fwd": 8.263226598501205, "global_bwd": 7.487393915653229, "x_quantize_rowwise": 0.3021806478500366, "g_quantize_rowwise": 1.7319358885288239, "w_quantize_rowwise": 0.15519559383392334, "w_quantize_colwise_transpose": 1.6133114695549011, "w_quantize_global": 0.48247724771499634, "w_quantize_global_transpose": 0.506427139043808, "time_standard": 53.04029583930969, "time_rowwise": 37.3353473842144, "time_global": 36.44480183720589} {"repeat": 64, "batch_size": 32768, "dim_out": 4096, "dim_in": 16384, "wm": 4, "switch": true, "standard_fwd": 17.73649826645851, "standard_gw": 16.359902918338776, "standard_gx": 18.0993489921093, "rowwise_fwd": 7.493957877159119, "rowwise_bwd": 8.352488279342651, "global_fwd": 7.486194372177124, "global_bwd": 8.28903540968895, "x_quantize_rowwise": 1.7313472926616669, "g_quantize_rowwise": 0.30205026268959045, "w_quantize_rowwise": 0.2255477011203766, "w_quantize_colwise_transpose": 2.4363920092582703, "w_quantize_global": 0.4815347492694855, "w_quantize_global_transpose": 0.5103759467601776, "time_standard": 52.195750176906586, "time_rowwise": 36.90168634057045, "time_global": 35.16044095158577} {"repeat": 64, "batch_size": 65536, "dim_out": 16384, "dim_in": 4096, "wm": 4, "switch": false, "standard_fwd": 36.309611052274704, "standard_gw": 32.85098075866699, "standard_gx": 34.34552624821663, "rowwise_fwd": 16.74525812268257, "rowwise_bwd": 15.026237815618515, "global_fwd": 16.574162989854813, "global_bwd": 14.977734535932541, "x_quantize_rowwise": 0.5954466760158539, "g_quantize_rowwise": 3.4569576382637024, "w_quantize_rowwise": 0.15521422028541565, "w_quantize_colwise_transpose": 1.6133897006511688, "w_quantize_global": 0.4822872579097748, "w_quantize_global_transpose": 0.5065612494945526, "time_standard": 103.50611805915833, "time_rowwise": 70.44348493218422, "time_global": 69.44413110613823} {"repeat": 64, "batch_size": 65536, "dim_out": 4096, "dim_in": 16384, "wm": 4, "switch": true, "standard_fwd": 35.40017828345299, "standard_gw": 33.037226647138596, "standard_gx": 36.30436211824417, "rowwise_fwd": 15.043705701828003, "rowwise_bwd": 16.756191849708557, "global_fwd": 15.011314302682877, "global_bwd": 16.580048948526382, "x_quantize_rowwise": 3.4548528492450714, "g_quantize_rowwise": 0.5951337516307831, "w_quantize_rowwise": 0.22584572434425354, "w_quantize_colwise_transpose": 2.4329908192157745, "w_quantize_global": 0.4813261330127716, "w_quantize_global_transpose": 0.5101598799228668, "time_standard": 104.74176704883575, "time_rowwise": 71.54594734311104, "time_global": 69.67006251215935} {"repeat": 64, "batch_size": 131072, "dim_out": 16384, "dim_in": 4096, "wm": 4, "switch": false, "standard_fwd": 73.40333238244057, "standard_gw": 73.76311346888542, "standard_gx": 70.41774317622185, "rowwise_fwd": 33.37597846984863, "rowwise_bwd": 30.345775187015533, "global_fwd": 33.00366923213005, "global_bwd": 30.218638479709625, "x_quantize_rowwise": 1.1825822293758392, "g_quantize_rowwise": 6.902601569890976, "w_quantize_rowwise": 0.15529245138168335, "w_quantize_colwise_transpose": 1.6109198331832886, "w_quantize_global": 0.48149004578590393, "w_quantize_global_transpose": 0.5066059529781342, "time_standard": 217.58418902754784, "time_rowwise": 147.33626320958138, "time_global": 146.05870097875595} {"repeat": 64, "batch_size": 131072, "dim_out": 4096, "dim_in": 16384, "wm": 4, "switch": true, "standard_fwd": 71.5160183608532, "standard_gw": 73.76786693930626, "standard_gx": 72.98104092478752, "rowwise_fwd": 30.291248112916946, "rowwise_bwd": 33.36654230952263, "global_fwd": 30.181586742401123, "global_bwd": 33.082425594329834, "x_quantize_rowwise": 6.902430206537247, "g_quantize_rowwise": 1.1815279722213745, "w_quantize_rowwise": 0.2262219786643982, "w_quantize_colwise_transpose": 2.4421699345111847, "w_quantize_global": 0.4816502332687378, "w_quantize_global_transpose": 0.5105249583721161, "time_standard": 218.26492622494698, "time_rowwise": 148.17800745368004, "time_global": 146.1080126464367} bitsandbytes/benchmarking/switchback/make_plot_with_jsonl.py000066400000000000000000000126361506752665600251420ustar00rootroot00000000000000import matplotlib.gridspec as gridspec import matplotlib.pyplot as plt import pandas as pd cmap = plt.get_cmap("cool") if __name__ == "__main__": fig = plt.figure(tight_layout=True, figsize=(12, 3.5)) gs = gridspec.GridSpec(1, 2) dims_to_consider = [1024, 1280, 1408, 1664, 2048, 4096] batch_size_for_plot1 = 32768 batch_sizes_for_plot2 = [2**14, 2**15, 2**16, 2**17] dims_to_xtick = [1024, 2048, 4096] logscale_plot1 = True ax = fig.add_subplot(gs[0, 0]) # TODO: change this to what you want. rdf = pd.read_json("speed_benchmark/info_a100_py2.jsonl", lines=True) df = rdf[rdf.batch_size == batch_size_for_plot1] # first plot the time occupied by different operations for k, marker, ls, color, name in [ ("standard_gx+standard_gw+standard_fwd", "s", "-", "C2", "Standard fp16 (sum of parts)"), ( "x_quantize_rowwise+g_quantize_rowwise+w_quantize_global+w_quantize_global_transpose+standard_gw+global_fwd+global_bwd", "o", "-", "C4", "SwitchBack int8 (sum of parts)", ), ("standard_fwd", "^", "--", "C2", "Matmul XW (standard)"), ("standard_gw", "^", "-.", "C2", "Matmul GW (standard)"), ("standard_gx", "^", ":", "gray", "Matmul GX (both)"), ("global_fwd", "^", "--", "C4", "Int8 Matmul XW (switchback)"), ("global_bwd", "^", "-.", "C4", "Int8 Matmul GW (switchback)"), ("x_quantize_rowwise", "P", "--", "C4", "Quantize rowwise X (switchback)"), ("g_quantize_rowwise", "P", "-.", "C4", "Quantize rowwise G (switchback)"), ("w_quantize_global", ".", "--", "C4", "Quantize global W (switchback)"), ("w_quantize_global_transpose", ".", "-.", "C4", "Quantize global and\ntranspose W (switchback)"), ]: xs = [] ys = [] for embed_dim in dims_to_consider: # average over dim -> 4*dim and 4*dim -> dim df_ = df[df.dim_in == embed_dim] df_ = df_[df_.dim_out == embed_dim * 4] xs.append(embed_dim) y_ = 0 for k_ in k.split("+"): y_ += df_[k_].values[0] df_ = df[df.dim_in == embed_dim * 4] df_ = df_[df_.dim_out == embed_dim] for k_ in k.split("+"): y_ += df_[k_].values[0] ys.append(y_ * 0.5) ax.plot( xs, ys, color=color, label=name, marker=marker, markersize=5 if marker == "s" else 5, linestyle=ls, linewidth=2 if "+" in k else 1.0, ) ax.set_xlabel("dim", fontsize=13) ax.set_ylabel("time (ms)", fontsize=13) ax.grid() ax.set_xscale("log") if logscale_plot1: ax.set_yscale("log") ax.tick_params(axis="x", labelsize=11) ax.tick_params(axis="y", labelsize=11) ax.set_xticks(dims_to_xtick) ax.set_xticklabels(dims_to_xtick) ax.set_xticks([], minor=True) leg = ax.legend(loc="upper center", bbox_to_anchor=(-0.64, 1.0), ncol=1, fontsize=10) leg.get_texts()[0].set_fontweight("bold") leg.get_texts()[1].set_fontweight("bold") plt.subplots_adjust(left=0.1) ax.set_title(" Linear layer, batch * sequence length = 32k", fontsize=10, loc="left", y=1.05, pad=-20) ax = fig.add_subplot(gs[0, 1]) # now plot the % speedup for different batch sizes for j, batch_size in enumerate(batch_sizes_for_plot2): all_xs, all_ys = [], [] for k, marker, ls, color, name in [ ("standard_gx+standard_gw+standard_fwd", "s", "-", "C2", "Standard fp16 (total time)"), ( "x_quantize_rowwise+g_quantize_rowwise+w_quantize_global+w_quantize_global_transpose+standard_gw+global_fwd+global_bwd", "o", "-", "C4", "SwitchBack int8 (total time)", ), ]: xs, ys = [], [] df = rdf[rdf.batch_size == batch_size] for embed_dim in dims_to_consider: df_ = df[df.dim_in == embed_dim] df_ = df_[df_.dim_out == embed_dim * 4] xs.append(embed_dim) y_ = 0 for k_ in k.split("+"): y_ += df_[k_].values[0] df_ = df[df.dim_in == embed_dim * 4] df_ = df_[df_.dim_out == embed_dim] for k_ in k.split("+"): y_ += df_[k_].values[0] ys.append(y_ * 0.5) all_xs.append(xs) all_ys.append(ys) color = cmap(j * 0.25) real_ys = [-((all_ys[1][i] - all_ys[0][i]) / all_ys[0][i]) * 100 for i in range(len(all_ys[0]))] markers = ["^", "v", "P", "o"] ax.plot( all_xs[0], real_ys, color=color, label=f"batch * sequence length = {batch_size}", marker=markers[j], markersize=5 if marker == "s" else 5, ) ax.legend() ax.set_xlabel("dim", fontsize=13) ax.set_xscale("log") ax.grid() ax.set_ylabel(r"% speedup", fontsize=13) ax.tick_params(axis="x", labelsize=11) ax.tick_params(axis="y", labelsize=11) ax.set_xticks(dims_to_xtick) ax.set_xticklabels(dims_to_xtick) ax.set_xticks([], minor=True) ax.set_title(" Linear layer summary, varying dimensions", fontsize=10, loc="left", y=1.05, pad=-20) plt.savefig("speed_benchmark/plot_with_info.pdf", bbox_inches="tight") bitsandbytes/benchmarking/switchback/plot_with_info.pdf000066400000000000000000001040741506752665600240720ustar00rootroot00000000000000%PDF-1.4 % 1 0 obj << /Type /Catalog /Pages 2 0 R >> endobj 8 0 obj << /Font 3 0 R /XObject 7 0 R /ExtGState 4 0 R /Pattern 5 0 R /Shading 6 0 R /ProcSet [ /PDF /Text /ImageB /ImageC /ImageI ] >> endobj 11 0 obj << /Type /Page /Parent 2 0 R /Resources 8 0 R /MediaBox [ 0 0 814.7277579804 244.8 ] /Contents 9 0 R /Annots 10 0 R >> endobj 9 0 obj << /Length 12 0 R /Filter /FlateDecode >> stream x]IsǑ_ї&fU~8ac|+! 2.Ͱ|Y]Uݯ{$R NԒWYp??{珻'NwNuNwNz">'YGTݓQvco;{cphiꍵz%dBA'R]7FofȚ>&ѱ:jr/M{bcb+SO EUڡzgQ gwϻwS լС4 V{ |ܽÿ{y7rɠݓ֝݋_^{ﺿu/SI7)XI+`c3#sHtɡTz6;P.:ÝUKU6&z- kJ4f^]^=;86-}jɫ^8)3}PrQ]C* H׽m3^-?R =i =6 Z1ޡBX㏈:r=Xgvh߇nn__|~x=6Մ&_t&lS٢iҦ,,־9 -` 1raUA%m4[Sm!۔ƨ8 P[6GE۟Bi&09C!D ֩͟)'!2)Ʒc(3F6I4{yk"ORhF$ b_AnG4TXA#y?,*FV18\D,!gRH Kq>uVs i4UB| b ~0<j2 .ծ5'4r@:(doD,Qv$f>sdsd:۠#bmHdb +kb";Ya(#k|4E2l=C4CӸi,[Z' YL^Q @bFi$Lްr 0"3"0-6].9PPBH+k(y!Mfp*bhUECá(c=d {oS^|0CCNk/ 15fM90$T`vdv sI|8^3mIcDRF'BVnccѬ$1 Fxiתм$E(f|V :8B-PCl#1&&:zEjBW6B ~p&RzPbë ~5.|KC |&eFN!Nod=q#Γ_DصueĬ=xCq4yfyhp&FHBRo_:RXcfm26֐G\!=dLY 0z2$Ef=g~4>wُ)eLN&&a"H`: rMk|4B3liZBO#W/ÄbZ⒗M9z4}!R*y0aks3coBPX|'dpdYD?{n.Ƅ.4JkçʵY.*iꧻzwj5gWx;ʘCD2j"B̄rlD3DR3whmaw(_J3/澶Zx=i# TY6Չы|fCOgwŠw\"^M -#LE$aL,OȘF<;&%sᤍ؛AoތImp<+'(囻_yψ||DG'ΔIkZ#L qLʗoTc}wL_ ~`>8*1~N(f,lyd(itm>DI3eR<93ìH#1=<+IFSHTBAŴ7/ ʪ{Vt4,@B|txn.u!9G7rN]& LYXiBl`' V(\4M=4#xZ& `ʇ!uZ3a$|֑>:TVM@O4ɎXLa ݎEx/|"sL IŐkLm}E';C)-B^ߚه $>9w'E!(OTC_u 耙l'R],I s#ΤBgfKI;1Q׮7/ϮαY'L$lEz6i_D.k{m /ȯUzQ(XyP|^͈r3:}zgNPOףlQj=@[xoh-)gz-J=Vm#/(F- +oFcruyL;t7G_X_Y_Y4vqj}]x\恫C{y{uy}qs{{\=?9Y/VG(/Rn=R˅^W )wƧ xtR!I|&Aهt>s=@D|g_J~͋&GbJЊׯblN\މL)xBmKK _4Dw):(؍m+ےzm 9i]OHgТ$kv;wK}y-ģTS4ϯ}?![,k_hBvzqy"ԅ>[g\K^UT\GDn~  #U4Q@J˧sBW\R ($> _c1VQG fEB!C N:`쨼|r돀  "̪3.` P:)O$/mNTAHT7Sxh& %qԐlOhrzaSh,%7[ReC'I6žARBL] endstream endobj 12 0 obj 5363 endobj 10 0 obj [ ] endobj 32 0 obj << /Length 273 /Filter /FlateDecode >> stream x=K0D:G#qRpn¢KFv.R<^%/v\ܷ|r ϔk][c,*΋<&\ל*xuf+?{ )pRy!DcΊz1 jV9dU \LMKX!zXQ(2:3rB# -ZtO\e=Z8NV=)tq>e~ ?1T7aXpNŵrѰ^o\P8I cx?9h endstream endobj 33 0 obj << /Length 343 /Filter /FlateDecode >> stream x5R;n%1 @gyb%5H%BHvB;,E?OQ/g8:>Ofb/CúS?OwCLΓ:'gu)XK'BعțJAv V!ٍN:Q.\DNmyغ%gM?ViPH/olšxz*2$g_n2nm80R2!Z&ݦ1I۾0 ,m؜"մJ~QRP7243Ӕq.// REF'C ]ws!>2!`$㧵ȡ[oTxtg}$fGZNJd=(\57]h(̼h/<?:o endstream endobj 34 0 obj << /Length 316 /Filter /FlateDecode >> stream x=K0 C9.ЀT1oљE@Aև3W۰Jz+c[ԴW8Q̲'uE||WEQ;u<7Kȃ+s>8dV>*{AME|޿|ןK29}vPr*mQp/yY_V[k;n `'9ԜFFckF({p|T(zIh3b| ݚL.p“r[ ګ-gu[t`pn)&h*ޗCӛ%n]-7+*ju ܫfJ>pyt endstream endobj 35 0 obj << /Length 237 /Filter /FlateDecode >> stream x5QKn1'sW=u[CTi$<`!8x#%5lC#Ah"2xPV$_&w8'L`EX1? o7nAhͩPEOUU9@p2{cP4+sP\KxD^ՍS#mtvp!rlȎw{:7\t}Sfhr{}UU endstream endobj 36 0 obj << /Length 238 /Filter /FlateDecode >> stream x5Kn1 Cs ] LPtR%tʐHytCzENYߊͿy>oOM(b.^SܗJq1Yj?2'SҵX mG^òt.V$#DE/ ěYqjD\ 3]t 2[YeOXw<츲VHIsBZ|Y}\ endstream endobj 37 0 obj << /Length 400 /Filter /FlateDecode >> stream x5SIn1 +j=Sג T%Ԓ%CJG]Cfĭ{Z{R\ak)eȱoc5ߐF:BC{_Ka멉FF=<1 *aF%*տex_M+|J,cN_&p :0^֦z&S#$-x;#jrgF=4큷F\JJRy ^A>Cr6^N!_BLkNςb{PGO> stream xMA!＂'Tj׌5{q/zpw*R3zpy٤"YOBkMJ ܍ ȃl$Hs5pܰA-~Tñ5֨0vp`> stream xEAr0 z@ %ǝN&=؋!dКæ^*-#G?L 0s#H-'@^^lj&AE,=svXux5r&:xi#gr:qZHwEԟYg#UIQ;ҺaM58J[x_=~gFc endstream endobj 40 0 obj << /Length 74 /Filter /FlateDecode >> stream x361V0P0F f )\@B.H  4XiLEW Q bAC 4$ endstream endobj 41 0 obj << /Length 82 /Filter /FlateDecode >> stream x5ͻ 0)m0" :@" }QiZ)5,n,tbʅ=Q:^z>® endstream endobj 42 0 obj << /Length 281 /Filter /FlateDecode >> stream xERI1> stream xEK1C>AsNf162-Bn50ðkpȈ1חɷ̷ s0>t\r1 Fˉ *@>tqL ~Y1GjUr:zkGxjrNۋ$ĩG> tGPTD xΎv&]&>+_?i Fh endstream endobj 44 0 obj << /Length 213 /Filter /FlateDecode >> stream x5P 1 =`o 4ٿ('Ł> stream xMA ; ,ڵiz_^`eaiLzU/Pɴ)B4 Ě#(7O;_oY4^# ^> endstream endobj 46 0 obj << /Length 241 /Filter /FlateDecode >> stream xE;n1 C{X>5"I䳗$&s! rK:[kNl ~lLfp ˍOzXȚ!g«镫l٪K9)'{"o2V!b`櫔-kh0VT ͽV3եwHY.Ϭ8N6dKHS(`z&mGfnx}{v٩Cfn77SHԦ÷n~^ endstream endobj 47 0 obj << /Length 153 /Filter /FlateDecode >> stream x=;0 Cw`:_+%hCM?扁8XS<"%}!ƚش7G{(k-uļU>[ q n5H: l+IB7!.gM):X-Hum6A0 endstream endobj 48 0 obj << /Length 151 /Filter /FlateDecode >> stream x=O91 )#{@b'h&GNtIxލf:Fpv޸v3m #rIF XʓlW3ܲDZq\x8I/Dnbp٫؀TJ9~W{m0 endstream endobj 49 0 obj << /Length 158 /Filter /FlateDecode >> stream xEPA! tzӃ&L@vt{4am2Au #U-Utఱi:cu ,J #;bnWM (UŎesBɓFS^mQ C%'3筬涿_ګ=?om8, endstream endobj 50 0 obj << /Length 346 /Filter /FlateDecode >> stream x5InD1D>hLVo lpQU8oʔ4yKƑ\StOψj%,$&g7y#Sܓ3:>b{ ת(=bΎ~*mzU,_Ʃ$oC.3{,aKd'ϩbzdڸ&mI^H4WcW;#3<8Y;3nbfޢ&1ZmDHWYꝹ\YUaj)~@؎Rَ6ٳaGcs#$&$biarú%Jm©k-[%XַbIk kU3ҁR[# &Ng|/֫~ endstream endobj 51 0 obj << /Length 330 /Filter /FlateDecode >> stream x5RKnD1\R'jE{mm^0Iݒ-e.&*O^[~۝w;G܎: 2=EfόzBU U=ܜ,lOGAAe|L <܈BDz^k i6f)+8R{KFl""Az'11sl7nl'0S1q1G+*u .qPOʞ,'MTJ,p8GF L.nh 6h0 #dz1s<|,8}9>Oq{R`9#d`|^H7as)c6#4  ]} endstream endobj 52 0 obj << /Length 18 /Filter /FlateDecode >> stream x36P0C4U endstream endobj 53 0 obj << /Length 132 /Filter /FlateDecode >> stream xEAC!D#b4]o;߄8όN΅؋[H=񷩙*(C2M/c{yR- \5qm֦̒2.U~TL\\Xv`I,ֳ/}Y 2L endstream endobj 54 0 obj << /Length 204 /Filter /FlateDecode >> stream xMQAn! `%!jCkmfF&& ix/k8+"*4EÖ@<8vǻ6Jn^7) V( NG♷+;`CN2zwmRI?,ǙOtMEIr0&vb(,&FCBQxM(VL! Z2\CU`>xoݾyL endstream endobj 55 0 obj << /Length 93 /Filter /FlateDecode >> stream x=0TA B?! +61FAw ^HޒM0iJ5i﹖?kkgJc}c_x endstream endobj 30 0 obj << /Type /Font /BaseFont /EVICAO+DejaVuSans-Bold /FirstChar 0 /LastChar 255 /FontDescriptor 29 0 R /Subtype /Type3 /Name /EVICAO+DejaVuSans-Bold /FontBBox [ -1070 -416 1976 1175 ] /FontMatrix [ 0.001 0 0 0.001 0 0 ] /CharProcs 31 0 R /Encoding << /Type /Encoding /Differences [ 32 /space 40 /parenleft /parenright 49 /one 54 /six 56 /eight 66 /B 83 /S 97 /a 99 /c /d 102 /f 104 /h /i 107 /k 109 /m /n /o /p 114 /r /s /t /u 119 /w ] >> /Widths 28 0 R >> endobj 29 0 obj << /Type /FontDescriptor /FontName /EVICAO+DejaVuSans-Bold /Flags 32 /FontBBox [ -1070 -416 1976 1175 ] /Ascent 929 /Descent -236 /CapHeight 0 /XHeight 0 /ItalicAngle 0 /StemV 0 /MaxWidth 1440 >> endobj 28 0 obj [ 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 348 456 521 838 696 1002 872 306 457 457 523 838 380 415 380 365 696 696 696 696 696 696 696 696 696 696 400 400 838 838 838 580 1000 774 762 734 830 683 683 821 837 372 372 775 637 995 837 850 733 850 770 720 682 812 774 1103 771 724 725 457 365 457 838 500 500 675 716 593 716 678 435 716 712 343 343 665 343 1042 712 687 716 716 493 595 478 712 652 924 645 652 582 712 365 712 838 600 696 600 380 435 657 1000 500 500 500 1440 720 412 1167 600 725 600 600 380 380 657 657 639 500 1000 500 1000 595 412 1094 600 582 724 348 456 696 696 636 696 365 500 500 1000 564 646 838 415 1000 500 500 838 438 438 500 736 636 380 500 438 564 646 1035 1035 1035 580 774 774 774 774 774 774 1085 734 683 683 683 683 372 372 372 372 838 837 850 850 850 850 850 838 850 812 812 812 812 724 738 719 675 675 675 675 675 675 1048 593 678 678 678 678 343 343 343 343 687 712 687 687 687 687 687 838 687 712 712 712 712 652 716 652 ] endobj 31 0 obj << /B 32 0 R /S 33 0 R /a 34 0 R /c 35 0 R /d 36 0 R /eight 37 0 R /f 38 0 R /h 39 0 R /i 40 0 R /k 41 0 R /m 42 0 R /n 43 0 R /o 44 0 R /one 45 0 R /p 46 0 R /parenleft 47 0 R /parenright 48 0 R /r 49 0 R /s 50 0 R /six 51 0 R /space 52 0 R /t 53 0 R /u 54 0 R /w 55 0 R >> endobj 60 0 obj << /Length 247 /Filter /FlateDecode >> stream xMQIn0 yOA%#K\D^PBF^ ֜?F?T[ 1Q$tQ7H7 ~WXw+[:v*B<HZD=sCt >7!Di^um46'GI)flm*V2 7TFZ622ZOv&'q.;;b >|iqA"4ťgxO\& endstream endobj 61 0 obj << /Length 51 /Filter /FlateDecode >> stream x34U0P F )\@> 4XiLEWW V endstream endobj 62 0 obj << /Length 61 /Filter /FlateDecode >> stream x355W0PF F )\@>ehifY@Has`zr2 endstream endobj 63 0 obj << /Length 90 /Filter /FlateDecode >> stream x= 0CLSTU׆|z[ȸ  څo'u`]^Bd;Jf&$qD;MJ endstream endobj 64 0 obj << /Length 255 /Filter /FlateDecode >> stream x5QKnD1\`@$jE{mmҷ c%S"奶Jɗu\2M~2kd+lIlЭӜ`\YCD,݂> stream x= 0LTU&`a57EVEPX]GLO;Rtp8M_J endstream endobj 66 0 obj << /Length 90 /Filter /FlateDecode >> stream x- 0{Om0vßNÈ MiI4Gg`Ljd9␲F\Ҩ  WcϜ? endstream endobj 67 0 obj << /Length 307 /Filter /FlateDecode >> stream x=Kn1 C>.ٞ('%EjLYS4p;lffZb|F]Y'f:Q96M,.x&[?Р5. 7tWe)4c{2jL]lR{<^DUG"cA JhZwEA{ jȲ;)ftR.n!kB+%H؆3r'hލKh!n.`S:EbdNriUNPӹý lEC)s_ klxJ&KVۨN<Zt endstream endobj 68 0 obj << /Length 108 /Filter /FlateDecode >> stream xMA "Ot:=_+R-( XNc wy'e I"? Uk;Chx9y(Tw+c). endstream endobj 69 0 obj << /Length 244 /Filter /FlateDecode >> stream xEMr!/W=ϤRYL ͼ$+ZGC쉕2|/2|2;:ٙpOn\e3a[;*;^CwR?9.R5tedz ̐Yt+u@ӔRl[J'4wk|uOT!ZͭlxA$ 's!Q:.`%ܠf|zRI0|h6ГgƚKEW4;I5ⶾ6qw+}>K] endstream endobj 70 0 obj << /Length 232 /Filter /FlateDecode >> stream x5QIn0 'ŠR@%%bc#/195&~grW,OG­I^ve?3VΤƓ E6=<2%:0ufډ!#RF*Sms mNv{I疬eX=QK1}l'(l>jD(g(gU!h]ݧmcFJӿ MMD g4ϽN!T endstream endobj 71 0 obj << /Length 68 /Filter /FlateDecode >> stream x36P0P07W544U022P042QH12443s`9`a$r`Zs: P9\\iqg endstream endobj 72 0 obj << /Length 231 /Filter /FlateDecode >> stream x5O9! y>0U@6N!x##fZd fSLſ" ap֬nvX6Y^LWg.ci9n]uSXG0 tÔ sTɎ28',v9~ 6!*z6yrA]E%  Qb_vt(sBA.!*PRQp>謟_\- endstream endobj 73 0 obj << /Length 249 /Filter /FlateDecode >> stream x=P;D! 9/$#pF-f߮)PL~3$G 1%Bn CR zt6:3?a7cE1t=&9 seVH'"3)*{x,6['= RRɥ?mʔ :f,dM8˻IR2v"}<ȣ:dMά#Oۙp_dNt΃96Mi*/RZ~ugdZ" endstream endobj 74 0 obj << /Length 395 /Filter /FlateDecode >> stream x=RKn@\yRUݼokCR*10}ʐK]qɷ^[<|.[Zy%>ܗ!]t.8G 2*DͪѡBN }9/װ =2A$)BnQAaPYQ2jo c BmH@DTg$gb`ѲTD{LjΨD>OM(L-VnS_|t*4U Xy9Hl!:n32`K9`GYutpL~OtZur@MAF2>)z,F3ar4k"X"bDls=L9l֡33*!ںj@vp?3m endstream endobj 75 0 obj << /Length 74 /Filter /FlateDecode >> stream x0P0P040S047R076R015QH1 \00,,`Ad3l#SS fpQ endstream endobj 76 0 obj << /Length 249 /Filter /FlateDecode >> stream xMQI0 @!^CC9 X 1,=!s7~ٻYz"SQR.bB]ϡ=kY9,s3c}I2!N uZ¸kbZvBzM;" 2;-+{n?R\ endstream endobj 77 0 obj << /Length 94 /Filter /FlateDecode >> stream xE TA O&2|` yF&^XJcP@r-p'TwiIU)ܤyg&yc c#< endstream endobj 78 0 obj << /Length 341 /Filter /FlateDecode >> stream xERKnD1ۿSpHᗐ󴪺[L `lxSdʰ%#$הO}"C4r$"I vkK쑨FZѥ͢Q9{z\H!P,er %A[w+ak.e2\}b.6)hUI£znvlU30r%[}]E+@QY@v1Hpp<řpK8@z%/W ud/žBw^dí> stream xEq1 C%0 AzB bO°\WqatÏ-3 G2cX ' Ѥvb\/:"̒@#|:Ǔ3t^!**na.@RԏQꚡ*+kjW]J>.2NݽZὕ=?v: endstream endobj 80 0 obj << /Length 72 /Filter /FlateDecode >> stream x32P0P4& f )\@B.H  %g m DDJ% endstream endobj 81 0 obj << /Length 83 /Filter /FlateDecode >> stream x=90ОS#"qxV0z!8`ʺtSƜ xJM zԴԾhKՉ? endstream endobj 82 0 obj << /Length 47 /Filter /FlateDecode >> stream x32P0P4& f )\V.L,іp "g ' endstream endobj 83 0 obj << /Length 258 /Filter /FlateDecode >> stream xEKr D#> stream x240S065U276r,#s# $`Ad3 | endstream endobj 85 0 obj << /Length 163 /Filter /FlateDecode >> stream xE;! C{N##>f2)6oclRX ڂ肇%gc6n5uVhR}ith6s+ fz:rCp_b9S6;ܬ~+UaiYK f@ S9ګ=? : endstream endobj 86 0 obj << /Length 322 /Filter /FlateDecode >> stream x5Qm0 50 ~% H;کHӼU^2\.u*YaCu|f!S֖{yJ,p 6s +QDnY@蚫Xb V)Xv$ r1YSE)|,ّZrY n}a8쳋DyL$dl>[E6v8Ѐ$AS8>aj('89֯1 NKRWϠA3.IwcJd)6 SbVofQlܠbod[7}_{d endstream endobj 87 0 obj << /Length 218 /Filter /FlateDecode >> stream x=P1 ]XzgKH"ER*5)/u%YS:|ynd6%*E/% }ՖC4h9~ 3*K6p*3 mtV[ Ф`׶ r " JMr^>C5Xs S'no^#]X.i5AMZ-^# @qRkp/'S endstream endobj 88 0 obj << /Length 83 /Filter /FlateDecode >> stream xE 0D{`~&>Q޿ %n'=\ )n3C ,]q0 \GPaվRBf endstream endobj 89 0 obj << /Length 239 /Filter /FlateDecode >> stream xMPm1 50<C hK[e#9oɺCNinͯ vPT%J\n4I/8k{Yx)=2f90=NƯ;$B %+ۘRN-{|=ލF-j80SѕqCgGW4 zB59p`˳Re]:4J(^7 endstream endobj 90 0 obj << /Length 150 /Filter /FlateDecode >> stream x=O90 ~ uXޓ"xȲc1I4@,#zB0,p6IAVIW>[L!^H<MdLjR[uU$] cp=a5z0cGUKnv}`0] endstream endobj 91 0 obj << /Length 151 /Filter /FlateDecode >> stream x5 0 C ϲ$#I9P7OhL 3Ke> stream x5RK#1)T1>O^Eۑ{$f!C?R{ ,[ɽJ%sՔϓb"BG5$ L(aϳ`bvm?Յvx2@f)KP$m|3T9Zbh!7م>IdȠbJ Tك;9Yֱ,AqU\U;vFj_ovRy8[f=1p@.zwCDj yB9 x%a䦴 YE4`Wdih4U endstream endobj 93 0 obj << /Length 243 /Filter /FlateDecode >> stream xMQ1 =8پy.xxe6DCISda}JFH.mKl<7BSkꔫS,MY{U7^QUa;`k&lQ "1P3i s1g* N*W/e*x0"?h_x)Bܕz!Mc,\wLtX ! ^p֍ #^ӰwEu[Ɣ}Q}]'{]^ endstream endobj 94 0 obj << /Length 160 /Filter /FlateDecode >> stream xE91s'H\˵G4Mzh>Fk,JDPc;(HLHcH#X]T*Ri¿ZDpTO\M1T;I^EZ1D0 oE&پ` 8 L7 endstream endobj 95 0 obj << /Length 334 /Filter /FlateDecode >> stream x-RKr s ]3s--|FЀ_䟯imӮzL;lOs^&_Zg:R3s(VFv6Hj\loXJЅn\F#"<& ⏜r*aꚮtrUq53IfHy]8 Sfy'rԩgK9YhtNONW3!9 bZWib _SQSz;~,s|s endstream endobj 96 0 obj << /Length 70 /Filter /FlateDecode >> stream x336S0P0 F )\@>,́,# .C c0mbl`fbdY 1 2 endstream endobj 97 0 obj << /Length 320 /Filter /FlateDecode >> stream x5RKn1)@O󪪋okL0`)/YK%ۥ[%BޏR-G- Q=X/":pJWIj3 P@ƻb2R$siqr&'=YߓKMNI]IJD}B<>G1ɀ6m&%:fU`)W l[00Sbbe5ˆQ]pwhcG i8^8^9c, ch&zОaZT| endstream endobj 98 0 obj << /Length 18 /Filter /FlateDecode >> stream x36P0C4R endstream endobj 99 0 obj << /Length 133 /Filter /FlateDecode >> stream xEK!Dq2s퀝nc=T=.k!񷩖EY>TzRV!ܠ,I*!P-en6 S0l_["u?*/ȥֻޢ1 endstream endobj 100 0 obj << /Length 340 /Filter /FlateDecode >> stream x5R9n1 } n H qtR;ZTvˇTt7we ɫRHV!25gRqɂ?å X#D2f[~ix+-XR8y- V~AG}TX IpPvH 9{,vQ3(1EA%9źtб،7:*[ZYc>,2zXtǃ 4 4e9QVxC 1! J<% _6?m^JkL endstream endobj 101 0 obj << /Length 251 /Filter /FlateDecode >> stream x-QIrA+fc+ @tZ⠌ ,W呫h31~+mX@T IcP5vzq ְ̓ f)1ge]=Aɬ\%La#g?2.s=7 FXanV34H!AF6r/tU1s-m_M8kX]q.'/6Qi@HnsxK{T`śn@6߼_`[ endstream endobj 102 0 obj << /Length 174 /Filter /FlateDecode >> stream xMIC! C31Aݢ-Xvs21F<\d "]VYϷ;ݏf_=ŰWǁK^*%c"R6;Trݛ+pC#< endstream endobj 103 0 obj << /Length 75 /Filter /FlateDecode >> stream x34R0P06fF & )\@>ehd fpY)XX&fP!.cSs@EƦ`?++  endstream endobj 104 0 obj << /Length 89 /Filter /FlateDecode >> stream x5 0 D{O! ؿIpcݓ^R@:'A^ VEQIMESϦM'3uzZ(mp} endstream endobj 105 0 obj << /Length 141 /Filter /FlateDecode >> stream x=0C @)|Oj#K zcÆ`%Tk@%7ș`Gzb8\f}B%h{Siܦq5)ꜣg44s {Sw. endstream endobj 106 0 obj << /Length 76 /Filter /FlateDecode >> stream x=;0 CGh~$B)$`Z>r$;|!WZҚ M 8 endstream endobj 107 0 obj << /Length 215 /Filter /FlateDecode >> stream x5Q9! @$/xO(͌V\!R!S>z.j4疴-s/fi<CѨQCԥS(>nc! b(,AͩEm'5^ޗSƖV8=O&t2ԫ5 E"<Z42j;0NqdJoWGT endstream endobj 58 0 obj << /Type /Font /BaseFont /BMQQDV+DejaVuSans /FirstChar 0 /LastChar 255 /FontDescriptor 57 0 R /Subtype /Type3 /Name /BMQQDV+DejaVuSans /FontBBox [ -1021 -463 1794 1233 ] /FontMatrix [ 0.001 0 0 0.001 0 0 ] /CharProcs 59 0 R /Encoding << /Type /Encoding /Differences [ 32 /space 37 /percent 40 /parenleft /parenright /asterisk 44 /comma 48 /zero /one /two /three /four /five /six /seven /eight /nine 61 /equal 71 /G 73 /I 76 /L /M 81 /Q 87 /W /X 97 /a /b /c /d /e 103 /g /h /i 107 /k /l /m /n /o /p /q /r /s /t /u /v /w 121 /y /z ] >> /Widths 56 0 R >> endobj 57 0 obj << /Type /FontDescriptor /FontName /BMQQDV+DejaVuSans /Flags 32 /FontBBox [ -1021 -463 1794 1233 ] /Ascent 929 /Descent -236 /CapHeight 0 /XHeight 0 /ItalicAngle 0 /StemV 0 /MaxWidth 1342 >> endobj 56 0 obj [ 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 600 318 401 460 838 636 950 780 275 390 390 500 838 318 361 318 337 636 636 636 636 636 636 636 636 636 636 337 337 838 838 838 531 1000 684 686 698 770 632 575 775 752 295 295 656 557 863 748 787 603 787 695 635 611 732 684 989 685 611 685 390 337 390 838 500 500 613 635 550 635 615 352 635 634 278 278 579 278 974 634 612 635 635 411 521 392 634 592 818 592 592 525 636 337 636 838 600 636 600 318 352 518 1000 500 500 500 1342 635 400 1070 600 685 600 600 318 318 518 518 590 500 1000 500 1000 521 400 1023 600 525 611 318 401 636 636 636 636 337 500 500 1000 471 612 838 361 1000 500 500 838 401 401 500 636 636 318 500 401 471 612 969 969 969 531 684 684 684 684 684 684 974 698 632 632 632 632 295 295 295 295 775 748 787 787 787 787 787 838 787 732 732 732 732 611 605 630 613 613 613 613 613 613 982 550 615 615 615 615 278 278 278 278 612 634 612 612 612 612 612 838 612 634 634 634 634 592 635 592 ] endobj 59 0 obj << /G 60 0 R /I 61 0 R /L 62 0 R /M 63 0 R /Q 64 0 R /W 65 0 R /X 66 0 R /a 67 0 R /asterisk 68 0 R /b 69 0 R /c 70 0 R /comma 71 0 R /d 72 0 R /e 73 0 R /eight 74 0 R /equal 75 0 R /five 76 0 R /four 77 0 R /g 78 0 R /h 79 0 R /i 80 0 R /k 81 0 R /l 82 0 R /m 83 0 R /n 85 0 R /nine 86 0 R /o 87 0 R /one 88 0 R /p 89 0 R /parenleft 90 0 R /parenright 91 0 R /percent 92 0 R /q 93 0 R /r 94 0 R /s 95 0 R /seven 96 0 R /six 97 0 R /space 98 0 R /t 99 0 R /three 100 0 R /two 101 0 R /u 102 0 R /v 103 0 R /w 104 0 R /y 105 0 R /z 106 0 R /zero 107 0 R >> endobj 3 0 obj << /F2 30 0 R /F1 58 0 R >> endobj 4 0 obj << /A1 << /Type /ExtGState /CA 0 /ca 1 >> /A2 << /Type /ExtGState /CA 1 /ca 1 >> /A3 << /Type /ExtGState /CA 0.8 /ca 0.8 >> >> endobj 5 0 obj << >> endobj 6 0 obj << >> endobj 7 0 obj << /M0 13 0 R /M1 14 0 R /M2 15 0 R /M3 16 0 R /M4 17 0 R /M5 18 0 R /M6 19 0 R /M7 20 0 R /M8 21 0 R /M9 22 0 R /M10 23 0 R /M11 24 0 R /M12 25 0 R /M13 26 0 R /M14 27 0 R /F1-DejaVuSans-minus 84 0 R >> endobj 13 0 obj << /Type /XObject /Subtype /Form /BBox [ -7.5 -7.5 7.5 7.5 ] /Length 38 /Filter /FlateDecode >> stream x3P2P53U\pf a"\\N\/ f endstream endobj 14 0 obj << /Type /XObject /Subtype /Form /BBox [ -7.5 -7.5 7.5 7.5 ] /Length 134 /Filter /FlateDecode >> stream xm= > stream x3P2Pb#=S\.]&r\\N\ endstream endobj 16 0 obj << /Type /XObject /Subtype /Form /BBox [ -7.5 -7.5 7.5 7.5 ] /Length 35 /Filter /FlateDecode >> stream x3P2Pb#=S\.]&r\\N\ endstream endobj 17 0 obj << /Type /XObject /Subtype /Form /BBox [ -7.5 -7.5 7.5 7.5 ] /Length 35 /Filter /FlateDecode >> stream x3P2Pb#=S\.]&r\\N\ endstream endobj 18 0 obj << /Type /XObject /Subtype /Form /BBox [ -7.5 -7.5 7.5 7.5 ] /Length 35 /Filter /FlateDecode >> stream x3P2Pb#=S\.]&r\\N\ endstream endobj 19 0 obj << /Type /XObject /Subtype /Form /BBox [ -7.5 -7.5 7.5 7.5 ] /Length 35 /Filter /FlateDecode >> stream x3P2Pb#=S\.]&r\\N\ endstream endobj 20 0 obj << /Type /XObject /Subtype /Form /BBox [ -7.5 -7.5 7.5 7.5 ] /Length 71 /Filter /FlateDecode >> stream x3P2P5г0]#=S\.T~Ip&618HztEW2* endstream endobj 21 0 obj << /Type /XObject /Subtype /Form /BBox [ -7.5 -7.5 7.5 7.5 ] /Length 71 /Filter /FlateDecode >> stream x3P2P5г0]#=S\.T~Ip&618HztEW2* endstream endobj 22 0 obj << /Type /XObject /Subtype /Form /BBox [ -6.25 -6.25 6.25 6.25 ] /Length 136 /Filter /FlateDecode >> stream xm; D{ |2נV!ϳj~&uZ_ A=IS5PΡ.ފT4=Ȼp-q&HoP?(Qcx 1Lc:$L]+pN endstream endobj 23 0 obj << /Type /XObject /Subtype /Form /BBox [ -6.25 -6.25 6.25 6.25 ] /Length 136 /Filter /FlateDecode >> stream xm; D{ |2נV!ϳj~&uZ_ A=IS5PΡ.ފT4=Ȼp-q&HoP?(Qcx 1Lc:$L]+pN endstream endobj 24 0 obj << /Type /XObject /Subtype /Form /BBox [ -7.5 -7.5 7.5 7.5 ] /Length 35 /Filter /FlateDecode >> stream x3P2Pb#=S\.]&r\\N\ endstream endobj 25 0 obj << /Type /XObject /Subtype /Form /BBox [ -7.5 -7.5 7.5 7.5 ] /Length 39 /Filter /FlateDecode >> stream x3P2P5P53U å`fpq9q endstream endobj 26 0 obj << /Type /XObject /Subtype /Form /BBox [ -7.5 -7.5 7.5 7.5 ] /Length 71 /Filter /FlateDecode >> stream x3P2P5г0]#=S\.T~Ip&618HztEW2* endstream endobj 27 0 obj << /Type /XObject /Subtype /Form /BBox [ -7.5 -7.5 7.5 7.5 ] /Length 134 /Filter /FlateDecode >> stream xm= > endobj 108 0 obj << /Creator (Matplotlib v3.7.1, https://matplotlib.org) /Producer (Matplotlib pdf backend v3.7.1) /CreationDate (D:20230401183100Z) >> endobj xref 0 109 0000000000 65535 f 0000000016 00000 n 0000032401 00000 n 0000028953 00000 n 0000028996 00000 n 0000029138 00000 n 0000029159 00000 n 0000029180 00000 n 0000000065 00000 n 0000000343 00000 n 0000005802 00000 n 0000000208 00000 n 0000005781 00000 n 0000029399 00000 n 0000029567 00000 n 0000029832 00000 n 0000029997 00000 n 0000030162 00000 n 0000030327 00000 n 0000030492 00000 n 0000030657 00000 n 0000030858 00000 n 0000031059 00000 n 0000031330 00000 n 0000031601 00000 n 0000031766 00000 n 0000031935 00000 n 0000032136 00000 n 0000013134 00000 n 0000012922 00000 n 0000012453 00000 n 0000014195 00000 n 0000005822 00000 n 0000006168 00000 n 0000006584 00000 n 0000006973 00000 n 0000007283 00000 n 0000007594 00000 n 0000008067 00000 n 0000008276 00000 n 0000008540 00000 n 0000008686 00000 n 0000008840 00000 n 0000009194 00000 n 0000009457 00000 n 0000009743 00000 n 0000009899 00000 n 0000010213 00000 n 0000010439 00000 n 0000010663 00000 n 0000010894 00000 n 0000011313 00000 n 0000011716 00000 n 0000011806 00000 n 0000012011 00000 n 0000012288 00000 n 0000027328 00000 n 0000027121 00000 n 0000026553 00000 n 0000028381 00000 n 0000014486 00000 n 0000014806 00000 n 0000014929 00000 n 0000015062 00000 n 0000015224 00000 n 0000015552 00000 n 0000015716 00000 n 0000015878 00000 n 0000016258 00000 n 0000016439 00000 n 0000016756 00000 n 0000017061 00000 n 0000017201 00000 n 0000017505 00000 n 0000017827 00000 n 0000018295 00000 n 0000018441 00000 n 0000018763 00000 n 0000018929 00000 n 0000019343 00000 n 0000019580 00000 n 0000019724 00000 n 0000019879 00000 n 0000019998 00000 n 0000020329 00000 n 0000020501 00000 n 0000020737 00000 n 0000021132 00000 n 0000021423 00000 n 0000021578 00000 n 0000021890 00000 n 0000022113 00000 n 0000022337 00000 n 0000022816 00000 n 0000023132 00000 n 0000023365 00000 n 0000023772 00000 n 0000023914 00000 n 0000024307 00000 n 0000024397 00000 n 0000024603 00000 n 0000025017 00000 n 0000025342 00000 n 0000025590 00000 n 0000025738 00000 n 0000025900 00000 n 0000026115 00000 n 0000026264 00000 n 0000032461 00000 n trailer << /Size 109 /Root 1 0 R /Info 108 0 R >> startxref 32613 %%EOF bitsandbytes/benchmarking/switchback/speed_benchmark.py000066400000000000000000000136751506752665600240450ustar00rootroot00000000000000import json import time import torch from bitsandbytes.triton.int8_matmul_mixed_dequantize import ( int8_matmul_mixed_dequantize, ) from bitsandbytes.triton.int8_matmul_rowwise_dequantize import ( int8_matmul_rowwise_dequantize, ) from bitsandbytes.triton.quantize_columnwise_and_transpose import ( quantize_columnwise_and_transpose, ) from bitsandbytes.triton.quantize_global import ( quantize_global, quantize_global_transpose, ) from bitsandbytes.triton.quantize_rowwise import quantize_rowwise # KNOW ISSUE: need to optimize "w_quantize_colwise_transpose" when embeddim is too large. def get_time(k, fn, info_dict): for _ in range(repeat // 2): fn() torch.cuda.synchronize() start = time.time() for _ in range(repeat): fn() torch.cuda.synchronize() end = time.time() ms = (end - start) / repeat * 1000 print(f"time {k}: {ms:.3f} ms") info_dict[k] = ms if __name__ == "__main__": torch.manual_seed(0) wm = 4 for dim in [1024, 1280, 1408, 1664, 2048, 4096]: # note "batch_size" is actually "batch_size * embed_dim", which is why it's large for batch_size in [256 * 32, 256 * 64, 256 * 128, 256 * 256, 256 * 512]: # switch switches dim_in and dim_out for switch in [False, True]: # hparams repeat = 64 batch_size = batch_size dim_out = dim * wm dim_in = dim if switch: dim_out = dim dim_in = wm * dim dim_in = round(dim_in) dim_out = round(dim_out) # simulate forward pass x = torch.randn(batch_size, dim_in, dtype=torch.float16).cuda() g = torch.randn(batch_size, dim_out, dtype=torch.float16).cuda() w = torch.randn(dim_out, dim_in, dtype=torch.float16).cuda() x_int8 = x.clone().to(torch.int8) g_int8 = g.clone().to(torch.int8) w_int8 = w.clone().to(torch.int8) wt_int8 = w.t().contiguous().clone().to(torch.int8) state_x_rowwise = x.max(dim=1)[0] state_g_rowwise = g.max(dim=1)[0] state_w_columnwise = w.max(dim=0)[0] state_w_rowwise = w.max(dim=1)[0] state_w_global = w.max() info = { "repeat": repeat, "batch_size": batch_size, "dim_out": dim_out, "dim_in": dim_in, "wm": wm, "switch": switch, } get_time("standard_fwd", lambda: x.matmul(w.t()), info) get_time("standard_gw", lambda: g.t().matmul(x), info) get_time("standard_gx", lambda: g.matmul(w), info) get_time( "rowwise_fwd", lambda: int8_matmul_rowwise_dequantize( x_int8, w_int8.t(), state_x_rowwise, state_w_columnwise, None, ), info, ) get_time( "rowwise_bwd", lambda: int8_matmul_rowwise_dequantize( g_int8, wt_int8.t(), state_x_rowwise, state_w_rowwise, None, ), info, ) get_time( "global_fwd", lambda: int8_matmul_mixed_dequantize(x_int8, w_int8.t(), state_x_rowwise, state_w_global, None), info, ) get_time( "global_bwd", lambda: int8_matmul_mixed_dequantize(g_int8, wt_int8.t(), state_x_rowwise, state_w_global, None), info, ) get_time("x_quantize_rowwise", lambda: quantize_rowwise(x), info) get_time("g_quantize_rowwise", lambda: quantize_rowwise(g), info) get_time("w_quantize_rowwise", lambda: quantize_rowwise(w), info) get_time("w_quantize_colwise_transpose", lambda: quantize_columnwise_and_transpose(w), info) get_time("w_quantize_global", lambda: quantize_global(w), info) get_time("w_quantize_global_transpose", lambda: quantize_global_transpose(w), info) time_standard = info["standard_fwd"] + info["standard_gx"] + info["standard_gw"] time_rowwise = ( info["x_quantize_rowwise"] + info["g_quantize_rowwise"] + info["w_quantize_colwise_transpose"] + info["w_quantize_rowwise"] + info["standard_gw"] + info["rowwise_fwd"] + info["rowwise_bwd"] ) time_global = ( info["x_quantize_rowwise"] + info["g_quantize_rowwise"] + info["w_quantize_global"] + info["w_quantize_global_transpose"] + info["standard_gw"] + info["global_fwd"] + info["global_bwd"] ) print("TOTAL STANDARD", time_standard) print("TOTAL ROWWISE", time_rowwise) print("TOTAL GLOBAL", time_global) print("speedup", -100 * (time_global - time_standard) / time_standard) info["time_standard"] = time_standard info["time_rowwise"] = time_rowwise info["time_global"] = time_global info_json = json.dumps(info) # TODO: change this to what you want. with open("speed_benchmark/info.jsonl", "a") as file: file.write(info_json + "\n") bitsandbytes/benchmarking/xpu/000077500000000000000000000000001506752665600170375ustar00rootroot00000000000000bitsandbytes/benchmarking/xpu/inference_benchmark.py000066400000000000000000000166161506752665600233730ustar00rootroot00000000000000import argparse import time # import intel_extension_for_pytorch as ipex import numpy as np import torch from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig MAX_NEW_TOKENS = 256 get_time = time.time system_prompt = "You are a helpful assistant" user_prompt = """Summarize this text please: ```Tell me, O muse, of that ingenious hero who travelled far and wide after he had sacked the famous town of Troy. Many cities did he visit, and many were the nations with whose manners and customs he was acquainted; moreover he suffered much by sea while trying to save his own life and bring his men safely home; but do what he might he could not save his men, for they perished through their own sheer folly in eating the cattle of the Sun-god Hyperion; so the god prevented them from ever reaching home. Tell me, too, about all these things, O daughter of Jove, from whatsoever source you may know them. So now all who escaped death in battle or by shipwreck had got safely home except Ulysses, and he, though he was longing to return to his wife and country, was detained by the goddess Calypso, who had got him into a large cave and wanted to marry him. But as years went by, there came a time when the gods settled that he should go back to Ithaca; even then, however, when he was among his own people, his troubles were not yet over; nevertheless all the gods had now begun to pity him except Neptune, who still persecuted him without ceasing and would not let him get home. Now Neptune had gone off to the Ethiopians, who are at the world's end, and lie in two halves, the one looking West and the other East. He had gone there to accept a hecatomb of sheep and oxen, and was enjoying himself at his festival; but the other gods met in the house of Olympian Jove, and the sire of gods and men spoke first. At that moment he was thinking of Aegisthus, who had been killed by Agamemnon's son Orestes; so he said to the other gods: "See now, how men lay blame upon us gods for what is after all nothing but their own folly. Look at Aegisthus; he must needs make love to Agamemnon's wife unrighteously and then kill Agamemnon, though he knew it would be the death of him; for I sent Mercury to warn him not to do either of these things, inasmuch as Orestes would be sure to take his revenge when he grew up and wanted to return home. Mercury told him this in all good will but he would not listen, and now he has paid for everything in full." Then Minerva said, "Father, son of Saturn, King of kings, it served Aegisthus right, and so it would any one else who does as he did; but Aegisthus is neither here nor there; it is for Ulysses that my heart bleeds, when I think of his sufferings in that lonely sea-girt island, far away, poor man, from all his friends. It is an island covered with forest, in the very middle of the sea, and a goddess lives there, daughter of the magician Atlas, who looks after the bottom of the ocean, and carries the great columns that keep heaven and earth asunder. This daughter of Atlas has got hold of poor unhappy Ulysses, and keeps trying by every kind of blandishment to make him forget his home, so that he is tired of life, and thinks of nothing but how he may once more see the smoke of his own chimneys. You, sir, take no heed of this, and yet when Ulysses was before Troy did he not propitiate you with many a burnt sacrifice? Why then should you keep on being so angry with him?" And Jove said, "My child, what are you talking about? How can I forget Ulysses than whom there is no more capable man on earth, nor more liberal in his offerings to the immortal gods that live in heaven? Bear in mind, however, that Neptune is still furious with Ulysses for having blinded an eye of Polyphemus king of the Cyclopes. Polyphemus is son to Neptune by the nymph Thoosa, daughter to the sea-king Phorcys; therefore though he will not kill Ulysses outright, he torments him by preventing him from getting home. Still, let us lay our heads together and see how we can help him to return; Neptune will then be pacified, for if we are all of a mind he can hardly stand out against us."```""" prompt = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}, ] def get_inputs(tokenizer): inputs = tokenizer.apply_chat_template( prompt, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True, ) return inputs def get_streamer(tokenizer): streamer = Streamer(tokenizer) # streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) return streamer class Streamer: def __init__(self, tokenizer, print_median=False): self.times = [] self.print_median = print_median self.tokenizer = tokenizer def put(self, t): self.times.append(get_time()) if len(self.times) > 1: print(f"Token latency: {1000 * (self.times[-1] - self.times[-2]):.1f} ms") if len(self.times) % 10 == 3 and self.print_median: ts = np.array(self.times) diff = ts[1:] - ts[:-1] # print("Token latency:", 1000 * diff, "ms") print("Token latency median:", np.median(1000 * diff), "ms") def print_report(self): times = np.array(self.times) diff = times[1:] - times[:-1] print(f"Median latency: {round(np.median(diff) * 1000, 2)}ms") percentiles = [10, 25, 50, 75, 90] print( "Latency percentiles", {p: round(1000 * float(np.percentile(diff, p)), 1) for p in percentiles}, ) def end(self, *args): pass def parse_arguments(): parser = argparse.ArgumentParser(description="Run inference benchmark for LLM models") parser.add_argument( "--device", type=str, default="xpu", help="Device to run inference on (e.g., xpu, cuda, cpu)", ) parser.add_argument( "--model-id", type=str, default="unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit", help="Model ID from Hugging Face or local path", ) parser.add_argument( "--attn", type=str, default="eager", choices=["eager", "flash_attention", "sdpa"], help="Attention implementation to use", ) return parser.parse_args() if __name__ == "__main__": args = parse_arguments() device = args.device model_id = args.model_id print(f"Running inference on {device} with model {model_id}") print(f"Using attention implementation: {args.attn}") tokenizer = AutoTokenizer.from_pretrained(model_id) model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation=args.attn) inputs = get_inputs(tokenizer) streamer = get_streamer(tokenizer) inputs = inputs.to(device) model = model.to(device) generation_config = GenerationConfig( use_cache=True, forced_eos_token_id=1, eos_token_id=1, max_new_tokens=MAX_NEW_TOKENS, do_sample=False, ) outputs = model.generate( **inputs, streamer=streamer, generation_config=generation_config, ) # Print the final outputs (including the input prompt) output_text = tokenizer.decode(outputs[0], skip_special_tokens=True) print(r"\Output (including prompt):") print("-" * 40) print(output_text) print("-" * 40) print(f"Peak memory usage: {torch.xpu.max_memory_allocated() / 1024**2:.0f}MB") streamer.print_report() bitsandbytes/bitsandbytes/000077500000000000000000000000001506752665600162665ustar00rootroot00000000000000bitsandbytes/bitsandbytes/__init__.py000066400000000000000000000042671506752665600204100ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. import importlib import sys import torch from . import _ops, research, utils from .autograd._functions import ( MatmulLtState, matmul, matmul_4bit, ) from .backends.cpu import ops as cpu_ops from .backends.default import ops as default_ops from .nn import modules from .optim import adam # This is a signal for integrations with transformers/diffusers. # Eventually we may remove this but it is currently required for compatibility. features = {"multi_backend"} supported_torch_devices = { "cpu", "cuda", # NVIDIA/AMD GPU "xpu", # Intel GPU "hpu", # Intel Gaudi "npu", # Ascend NPU "mps", # Apple Silicon } if torch.cuda.is_available(): from .backends.cuda import ops as cuda_ops if hasattr(torch, "xpu") and torch.xpu.is_available(): from .backends.xpu import ops as xpu_ops if importlib.util.find_spec("habana_frameworks") and importlib.util.find_spec("habana_frameworks.torch"): # In case not automatically imported import habana_frameworks.torch if hasattr(torch, "hpu") and torch.hpu.is_available(): from .backends.hpu import ops as hpu_ops def _import_backends(): """ Discover and autoload all available backends installed as separate packages. Packages with an entrypoint for "bitsandbytes.backends" will be loaded. Inspired by PyTorch implementation: https://pytorch.org/tutorials/prototype/python_extension_autoload.html """ from importlib.metadata import entry_points if sys.version_info < (3, 10): extensions = entry_points().get("bitsandbytes.backends", []) else: extensions = entry_points(group="bitsandbytes.backends") for ext in extensions: try: entry = ext.load() entry() except Exception as e: raise RuntimeError(f"bitsandbytes: failed to load backend {ext.name}: {e}") from e _import_backends() __pdoc__ = { "libbitsandbytes": False, "optim.optimizer.Optimizer8bit": False, "optim.optimizer.MockArgs": False, } __version__ = "0.48.1" bitsandbytes/bitsandbytes/__main__.py000066400000000000000000000001321506752665600203540ustar00rootroot00000000000000if __name__ == "__main__": from bitsandbytes.diagnostics.main import main main() bitsandbytes/bitsandbytes/_ops.py000066400000000000000000000352461506752665600176120ustar00rootroot00000000000000from collections.abc import Sequence from math import prod from typing import Optional import torch _IS_TORCH_GTE_24 = False if hasattr(torch.library, "register_fake"): _IS_TORCH_GTE_24 = True register_fake = torch.library.register_fake register_kernel = torch.library.register_kernel else: # PyTorch <= 2.3 register_fake = torch.library.impl_abstract register_kernel = torch.library.impl # Int8 mixed precision matmul + dequant + bias torch.library.define( "bitsandbytes::int8_mixed_scaled_mm", "(Tensor A, Tensor CA, Tensor CB, Tensor SCA, Tensor SCB, Tensor? outlier_cols=None, Tensor? bias=None) -> (Tensor, Tensor?)", ) @register_fake("bitsandbytes::int8_mixed_scaled_mm") def _( A: torch.Tensor, CA: torch.Tensor, CB: torch.Tensor, SCA: torch.Tensor, SCB: torch.Tensor, outlier_cols: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: shapeC = (*CA.shape[:-1], CB.shape[0]) out = torch.empty(shapeC, device=A.device, dtype=A.dtype) outlier_cols = torch.library.get_ctx().new_dynamic_size() subA = A.new_empty(outlier_cols, dtype=torch.int64) return out, subA # Higher level op: int8 matmul + dequant + bias torch.library.define( "bitsandbytes::int8_scaled_mm", "(Tensor A, Tensor B, Tensor row_stats, Tensor col_stats, Tensor? bias=None, ScalarType? dtype=None) -> Tensor", ) @register_fake("bitsandbytes::int8_scaled_mm") def _( A: torch.Tensor, B: torch.Tensor, row_stats: torch.Tensor, col_stats: torch.Tensor, bias: Optional[torch.Tensor] = None, dtype: Optional[torch.dtype] = None, ) -> torch.Tensor: shapeC = (*A.shape[:-1], B.shape[0]) return torch.empty(shapeC, device=A.device, dtype=dtype or torch.float16) torch.library.define( "bitsandbytes::int8_linear_matmul", "(Tensor A, Tensor B) -> Tensor", ) @register_fake("bitsandbytes::int8_linear_matmul") def _(A: torch.Tensor, B: torch.Tensor): torch._check(A.dtype == torch.int8, lambda: "A must be int8") torch._check(B.dtype == torch.int8, lambda: "B must be int8") shapeC = (*A.shape[:-1], B.shape[0]) return torch.empty(shapeC, device=A.device, dtype=torch.int32) # More info on `out` overloads: # https://github.com/pytorch/pytorch/issues/125044 torch.library.define( "bitsandbytes::int8_linear_matmul.out", "(Tensor A, Tensor B, Tensor! out) -> ()", ) @register_fake("bitsandbytes::int8_linear_matmul.out") def _(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor): shapeC = (*A.shape[:-1], B.shape[0]) torch._check(A.dtype == torch.int8, lambda: "A must be int8") torch._check(B.dtype == torch.int8, lambda: "B must be int8") torch._check(out.shape == shapeC, lambda: f"Expected out.shape == {shapeC}, got {out.shape}") torch._check(out.device == A.device, lambda: f"Expected out.device == {A.device}, got {out.device}") torch._check(out.dtype == torch.int32, lambda: f"Expected out.dtype == int32, got {out.dtype}") torch.library.define( "bitsandbytes::int8_vectorwise_quant", "(Tensor A, float threshold=0.0) -> (Tensor, Tensor, Tensor?)", ) @register_fake("bitsandbytes::int8_vectorwise_quant") def _(A: torch.Tensor, threshold=0.0): out_row = torch.empty(A.shape, device=A.device, dtype=torch.int8) row_stats = torch.empty(prod(A.shape[:-1]), device=A.device, dtype=torch.float32) if threshold == 0.0: return out_row, row_stats, None outlier_cols = torch.library.get_ctx().new_dynamic_size() return out_row, row_stats, A.new_empty(outlier_cols, dtype=torch.int64) torch.library.define("bitsandbytes::int8_vectorwise_dequant", "(Tensor A, Tensor stats) -> Tensor") @register_fake("bitsandbytes::int8_vectorwise_dequant") def _(A: torch.Tensor, stats: torch.Tensor) -> torch.Tensor: torch._check(A.dtype == torch.int8, lambda: "A must be int8") return torch.empty_like(A, dtype=torch.float32) # Default PyTorch-native implementation @register_kernel("bitsandbytes::int8_vectorwise_dequant", "default") def _(A: torch.Tensor, stats: torch.Tensor): # To dequantize we divide by 127, or multiply by the reciprocal. return A * stats.view(-1, 1) * 7.874015718698502e-3 torch.library.define( "bitsandbytes::int8_mm_dequant", "(Tensor A, Tensor row_stats, Tensor col_stats, ScalarType? dtype=None, Tensor? bias=None) -> Tensor", ) @register_fake("bitsandbytes::int8_mm_dequant") def _( A: torch.Tensor, row_stats: torch.Tensor, col_stats: torch.Tensor, dtype: Optional[torch.dtype] = None, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: torch._check(A.dtype == torch.int32, lambda: "A must be int32") return torch.empty_like(A, dtype=dtype or torch.float16) torch.library.define( "bitsandbytes::int8_double_quant", "(Tensor A, float threshold=0.0) -> (Tensor, Tensor, Tensor, Tensor, Tensor?)", ) @register_fake("bitsandbytes::int8_double_quant") def _( A: torch.Tensor, threshold=0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: out_row = torch.empty_like(A, dtype=torch.int8) out_col = torch.empty_like(A, dtype=torch.int8) row_stats = torch.empty(prod(A.shape[:-1]), device=A.device, dtype=torch.float32) col_stats = torch.empty(A.shape[-1], device=A.device, dtype=torch.float32) outlier_n = torch.library.get_ctx().new_dynamic_size() outlier_cols = A.new_empty(outlier_n, dtype=torch.int64) return out_row, out_col, row_stats, col_stats, outlier_cols torch.library.define( "bitsandbytes::dequantize_4bit", "(Tensor A, Tensor absmax, int blocksize, str quant_type, int[] shape, ScalarType dtype) -> Tensor", ) @register_fake("bitsandbytes::dequantize_4bit") def _( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, shape: Sequence[int], dtype: torch.dtype, ) -> torch.Tensor: torch._check_is_size(blocksize) return torch.empty(shape, dtype=dtype, device=A.device) torch.library.define( "bitsandbytes::dequantize_4bit.out", "(Tensor A, Tensor absmax, int blocksize, str quant_type, int[] shape, ScalarType dtype, Tensor! out) -> ()", ) @register_fake("bitsandbytes::dequantize_4bit.out") def _( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, shape: Sequence[int], dtype: torch.dtype, out: torch.Tensor, ) -> None: torch._check_is_size(blocksize) torch._check(out.shape == shape, lambda: f"Expected out.shape == {shape}, got {out.shape}") torch._check(out.device == A.device, lambda: f"Expected out.device == {A.device}, got {out.device}") torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}") torch.library.define( "bitsandbytes::quantize_4bit", "(Tensor A, int blocksize, str quant_type, ScalarType quant_storage) -> (Tensor, Tensor)", ) @register_fake("bitsandbytes::quantize_4bit") def _( A: torch.Tensor, blocksize: int, quant_type: str, quant_storage: torch.dtype ) -> tuple[torch.Tensor, torch.Tensor]: torch._check_is_size(blocksize) n = A.numel() blocks = -(n // -blocksize) absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32) out = torch.empty(((n + 1) // (quant_storage.itemsize * 2), 1), device=A.device, dtype=quant_storage) return out, absmax torch.library.define( "bitsandbytes::dequantize_blockwise", "(Tensor A, Tensor absmax, Tensor code, int blocksize, ScalarType dtype) -> Tensor", ) @register_fake("bitsandbytes::dequantize_blockwise") def _(A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype) -> torch.Tensor: torch._check_is_size(blocksize) torch._check(A.dtype == torch.uint8, lambda: f"A must be uint8, got {A.dtype}") return torch.empty_like(A, dtype=dtype) torch.library.define( "bitsandbytes::dequantize_blockwise.out", "(Tensor A, Tensor absmax, Tensor code, int blocksize, ScalarType dtype, Tensor! out) -> ()", ) @register_fake("bitsandbytes::dequantize_blockwise.out") def _( A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype, out: torch.Tensor ): torch._check_is_size(blocksize) torch._check(A.dtype == torch.uint8, lambda: f"A must be uint8, got {A.dtype}") torch._check(out.shape == A.shape, lambda: f"Expected out.shape == {A.shape}, got {out.shape}") torch._check(out.device == A.device, lambda: f"Expected out.device == {A.device}, got {out.device}") torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}") torch.library.define("bitsandbytes::quantize_blockwise", "(Tensor A, Tensor code, int blocksize) -> (Tensor, Tensor)") @register_fake("bitsandbytes::quantize_blockwise") def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]: torch._check_is_size(blocksize) n = A.numel() blocks = -(n // -blocksize) absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32) out = torch.empty_like(A, dtype=torch.uint8) return out, absmax torch.library.define( "bitsandbytes::gemv_4bit", "(Tensor A, Tensor B, int[] shapeB, Tensor absmax, Tensor code, int blocksize) -> Tensor", ) @register_fake("bitsandbytes::gemv_4bit") def _( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int ) -> torch.Tensor: torch._check_is_size(blocksize) torch._check(A.numel() == A.size(-1), lambda: f"A must be a vector with leading dimensions of 1, got {A.shape}") torch._check( A.dtype in [torch.float16, torch.bfloat16, torch.float32], lambda: f"A must be float16, bfloat16, or float32, got {A.dtype}", ) torch._check( B.dtype in [torch.uint8, torch.bfloat16, torch.float16, torch.float32], lambda: f"B must be backed by storage of type uint8, bfloat16, float16, or float32, got {B.dtype}", ) shape = (*A.shape[:-1], shapeB[0]) return torch.empty(shape, device=A.device, dtype=A.dtype) torch.library.define( "bitsandbytes::gemv_4bit.out", "(Tensor A, Tensor B, int[] shapeB, Tensor absmax, Tensor code, int blocksize, Tensor! out) -> ()", ) @register_fake("bitsandbytes::gemv_4bit.out") def _( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int, out: torch.Tensor, ) -> None: torch._check_is_size(blocksize) torch._check(A.numel() == A.size(-1), lambda: f"A must be a vector with leading dimensions of 1, got {A.shape}") torch._check( A.dtype in [torch.float16, torch.bfloat16, torch.float32], lambda: f"A must be float16, bfloat16, or float32, got {A.dtype}", ) torch._check( B.dtype in [torch.uint8, torch.bfloat16, torch.float16, torch.float32], lambda: f"B must be backed by storage of type uint8, bfloat16, float16, or float32, got {B.dtype}", ) torch._check( out.shape == (*A.shape[:-1], shapeB[0]), lambda: f"Expected out.shape == {(*A.shape[:-1], shapeB[0])}, got {out.shape}", ) torch._check(out.device == A.device, lambda: f"Expected out.device == {A.device}, got {out.device}") torch._check(out.dtype == A.dtype, lambda: f"Expected out.dtype == {A.dtype}, got {out.dtype}") torch.library.define( "bitsandbytes::optimizer_update_32bit", "(str optimizer_name, Tensor(a0!) g, Tensor(a1!) p, Tensor(a2!) state1, Tensor(a3!)? state2, Tensor(a4!)? unorm_vec, float max_unorm, float param_norm, float beta1, float beta2, float beta3, float alpha, float eps, float weight_decay, int step, float lr, float gnorm_scale, bool skip_zeros=False) -> ()", ) @register_fake("bitsandbytes::optimizer_update_32bit") def _( optimizer_name: str, g: torch.Tensor, p: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], unorm_vec: Optional[torch.Tensor], max_unorm: float, param_norm: float, beta1: float, beta2: float, beta3: float, alpha: float, eps: float, weight_decay: float, step: int, lr: float, gnorm_scale: float, skip_zeros=False, ) -> None: torch._check( g.numel() == p.numel(), lambda: f"g and p must have the same number of elements, got {g.numel()} and {p.numel()}", ) compute_dtypes = [torch.float16, torch.bfloat16, torch.float32] torch._check( g.dtype in compute_dtypes, lambda: f"g must be bfloat16, float16, or float32, got {g.dtype}", ) torch._check( g.dtype == p.dtype, lambda: f"Expected all tensors to have the same dtype, got g.dtype={g.dtype}, p.dtype={p.dtype}", ) torch.library.define( "bitsandbytes::optimizer_update_8bit_blockwise", "(str optimizer_name, Tensor(a0!) g, Tensor(a1!) p, Tensor(a2!) state1, Tensor(a3!)? state2, float beta1, float beta2, float beta3, float alpha, float eps, int step, float lr, Tensor(a4!) qmap1, Tensor(a5!)? qmap2, Tensor(a6!) absmax1, Tensor(a7!)? absmax2, float weight_decay, float gnorm_scale, bool skip_zeros=False) -> ()", ) @register_fake("bitsandbytes::optimizer_update_8bit_blockwise") def _( optimizer_name: str, g: torch.Tensor, p: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], beta1: float, beta2: float, beta3: float, alpha: float, eps: float, step: int, lr: float, qmap1: torch.Tensor, qmap2: Optional[torch.Tensor], absmax1: torch.Tensor, absmax2: Optional[torch.Tensor], weight_decay: float, gnorm_scale: float, skip_zeros=False, ) -> None: torch._check( g.numel() == p.numel(), lambda: f"g and p must have the same number of elements, got {g.numel()} and {p.numel()}", ) compute_dtypes = [torch.float16, torch.bfloat16, torch.float32] torch._check( g.dtype in compute_dtypes, lambda: f"g must be bfloat16, float16, or float32, got {g.dtype}", ) torch._check( g.dtype == p.dtype, lambda: f"Expected all tensors to have the same dtype, got g.dtype={g.dtype}, p.dtype={p.dtype}", ) torch._check( state1.dtype == torch.uint8, lambda: f"state1 must be uint8, got {state1.dtype}", ) torch._check( qmap1.dtype == absmax1.dtype == torch.float32, lambda: f"Expected qmap1 and absmax1 to be float32, got qmap1.dtype={qmap1.dtype}, absmax1.dtype={absmax1.dtype}", ) if state2 is not None: torch._check( state2.dtype == torch.uint8, lambda: f"state2 must be uint8, got {state2.dtype}", ) torch._check( qmap2.dtype == absmax2.dtype == torch.float32, lambda: f"Expected qmap2 and absmax2 to be float32, got qmap2.dtype={qmap2.dtype}, absmax2.dtype={absmax2.dtype}", ) bitsandbytes/bitsandbytes/autograd/000077500000000000000000000000001506752665600200745ustar00rootroot00000000000000bitsandbytes/bitsandbytes/autograd/__init__.py000066400000000000000000000001031506752665600221770ustar00rootroot00000000000000from ._functions import get_inverse_transform_indices, undo_layout bitsandbytes/bitsandbytes/autograd/_functions.py000066400000000000000000000402061506752665600226170ustar00rootroot00000000000000from dataclasses import dataclass from math import prod from typing import Callable, Optional import warnings from warnings import warn import torch from typing_extensions import deprecated import bitsandbytes.functional as F # The inverse transformation for the colTuring and colAmpere format were contributed by Alex Borzunov: # https://github.com/bigscience-workshop/petals/blob/main/src/petals/utils/linear8bitlt_patch.py """ This class pools outlier dimensions across layers. This is particularly important for small models where outlier features are less systematic and occur with low frequency. """ class GlobalOutlierPooler: _instance = None def __init__(self): raise RuntimeError("Call get_instance() instead") def initialize(self): self.outliers = set() self.model_dim = None @classmethod def get_instance(cls): if cls._instance is None: cls._instance = cls.__new__(cls) cls._instance.initialize() return cls._instance def add_outliers(self, outlier_idx, feature_dim): if self.model_dim is None: self.model_dim = feature_dim if feature_dim != self.model_dim: return # we do not encode outliers for the 2nd FFN layer self.outliers.update(outlier_idx.tolist()) def get_current_outlier_idx(self): return torch.Tensor(list(self.outliers)).to(torch.int64) @deprecated( "This function is deprecated and will be removed in a future release.", category=FutureWarning, ) def get_inverse_transform_indices( transform_tile: Callable[[torch.Tensor], torch.Tensor], tile_size: tuple[int, int], ): """ Compute a permutation of indices that invert the specified (tiled) matrix transformation :param transform_tile: a function that applies forward transform to a tensor of shape [dim1, dim2] :param tile_size: higher-level tile dimensions, i.e. (8, 32) for Turing and (32, 32) for Ampere :note: we assume that tile_transform applies to a cpu-based int8 tensor of shape tile_size :example: transform_tile function for the turing layout (bitsandbytes.functional as F) :returns: indices """ d1, d2 = tile_size assert 0 < d1 * d2 < 2**64 tile_indices = torch.arange(d1 * d2, dtype=torch.int64).view(d1, d2) # encode each position in tile as a tuple of <= 8 unique bytes permuted_tile_indices = torch.zeros_like(tile_indices) for i in range(8): # select i-th byte, apply transformation and trace where each index ended up ith_dim_indices = torch.div(tile_indices, 256**i, rounding_mode="trunc") % 256 sample_tile_i = (ith_dim_indices - 128).to(torch.int8).contiguous() assert torch.all(sample_tile_i.int() + 128 == ith_dim_indices), "int overflow" permuted_tile_i = transform_tile(sample_tile_i) ith_permuted_indices = permuted_tile_i.to(tile_indices.dtype) + 128 permuted_tile_indices += ith_permuted_indices * (256**i) if d1 * d2 < 256**i: break # if all indices fit in i bytes, stop early return permuted_tile_indices _is_compiling = torch.compiler.is_compiling @deprecated( "This function is deprecated and will be removed in a future release.", category=FutureWarning, ) def undo_layout(permuted_tensor: torch.Tensor, tile_indices: torch.LongTensor) -> torch.Tensor: """ Undo a tiled permutation such as turing or ampere layout :param permuted_tensor: torch tensor in a permuted layout :param tile_indices: reverse transformation indices, from get_inverse_transform_indices :return: contiguous row-major tensor """ (rows, cols), (tile_rows, tile_cols) = permuted_tensor.shape, tile_indices.shape assert rows % tile_rows == cols % tile_cols == 0, "tensor must contain a whole number of tiles" tensor = permuted_tensor.reshape(-1, tile_indices.numel()).t() outputs = torch.empty_like(tensor) # note: not using .index_copy because it was slower on cuda outputs[tile_indices.flatten()] = tensor outputs = outputs.reshape(tile_rows, tile_cols, cols // tile_cols, rows // tile_rows) outputs = outputs.permute(3, 0, 2, 1) # (rows // tile_rows, tile_rows), (cols // tile_cols, tile_cols) return outputs.reshape(rows, cols).contiguous() @dataclass class MatmulLtState: _tile_indices: Optional[torch.Tensor] = None # TODO: remove force_no_igemmlt: bool = False CB: Optional[torch.Tensor] = None CxB: Optional[torch.Tensor] = None # TODO: Deprecate/remove SB: Optional[torch.Tensor] = None SCB: Optional[torch.Tensor] = None CxBt: Optional[torch.Tensor] = None # TODO: Deprecate/remove SBt: Optional[torch.Tensor] = None CBt: Optional[torch.Tensor] = None subB: Optional[torch.Tensor] = None outlier_pool: Optional[GlobalOutlierPooler] = None has_accumulated_gradients = False threshold = 0.0 idx: Optional[torch.Tensor] = None is_training = True has_fp16_weights = True use_pool = False formatB = "row" # TODO: Deprecate/remove def reset_grads(self): self.CB = None self.CxB = None self.SB = None self.SCB = None self.CxBt = None self.SBt = None self.CBt = None @property def tile_indices(self): raise ValueError("tile_indices is no longer supported.") class MatMul8bitLt(torch.autograd.Function): @staticmethod def forward( ctx: torch.autograd.function.FunctionCtx, A: torch.Tensor, B: torch.Tensor, out: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None, state: Optional[MatmulLtState] = None, ): state = state or MatmulLtState() # default of pytorch behavior if inputs are empty ctx.is_empty = False if prod(A.shape) == 0: ctx.is_empty = True ctx.A = A ctx.B = B ctx.bias = bias if A.shape[-1] == B.shape[0]: return torch.empty(A.shape[:-1] + B.shape[1:], dtype=A.dtype, device=A.device) else: return torch.empty(A.shape[:-1] + B.shape[:1], dtype=A.dtype, device=A.device) input_shape = A.shape # Cast A to fp16 if A.dtype != torch.float16 and not _is_compiling(): warnings.warn(f"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization") if len(A.shape) == 3: A = A.reshape(-1, A.shape[-1]) # 1. Quantize A. Note that as a side-effect, outliers are suppressed in CA/CAt. if ctx.needs_input_grad[1]: # Slower path CA, CAt, SCA, SCAt, outlier_cols = F.int8_double_quant(A.to(torch.float16), threshold=state.threshold) else: # Fast path CA, SCA, outlier_cols = F.int8_vectorwise_quant(A.to(torch.float16), threshold=state.threshold) CAt = SCAt = None has_grad = False if state.has_fp16_weights or state.CB is None: has_grad = getattr(B, "grad", None) is not None is_transposed = not B.is_contiguous() and B.shape[0] == B.stride(1) if is_transposed: B = B.contiguous() if (state.is_training and not has_grad) or state.CB is None or state.SCB is None: state.reset_grads() # 2. Quantize B state.CB, state.SCB, _ = F.int8_vectorwise_quant(B.to(torch.float16)) # Handle sparse decomposition if state.threshold > 0.0: state.idx = outlier_cols # Mixed Int8 Matmul + Dequant + Bias output, subA = torch.ops.bitsandbytes.int8_mixed_scaled_mm( A, CA, state.CB, SCA, state.SCB, outlier_cols, bias, ) else: # Int8 Matmul + Dequant + Bias output = torch.ops.bitsandbytes.int8_scaled_mm.default( CA, state.CB, SCA, state.SCB, bias=bias, dtype=A.dtype ) subA = None # 5. Save state ctx.state = state ctx.grad_shape = input_shape ctx.dtype_A = A.dtype ctx.dtype_bias = None if bias is None else bias.dtype if any(ctx.needs_input_grad[:2]): ctx.tensors = (CAt, subA, A) ctx.tensor_states = (SCAt, state.idx) else: ctx.tensors = [None, None, None] ctx.tensor_states = (None, None) ctx.save_for_backward(None, None) output_shape = (*input_shape[:-1], state.CB.shape[0]) if len(input_shape) == 3: return output.reshape(output_shape) return output @staticmethod def backward(ctx: torch.autograd.function.FunctionCtx, grad_output: torch.Tensor): if ctx.is_empty: bias_grad = None if ctx.bias is None else torch.zeros_like(ctx.bias) return torch.zeros_like(ctx.A), torch.zeros_like(ctx.B), None, bias_grad, None req_gradA, req_gradB, _, req_gradBias, _ = ctx.needs_input_grad CAt, subA, A = ctx.tensors SCAt, idx = ctx.tensor_states state: MatmulLtState = ctx.state grad_A = grad_B = grad_bias = None if req_gradBias: # compute grad_bias first before changing grad_output dtype grad_bias = grad_output.sum(0, dtype=ctx.dtype_bias) # Cast grad_output to fp16 if len(grad_output.shape) == 3: grad_output = grad_output.reshape(-1, grad_output.shape[-1]).contiguous() if req_gradB: Cgrad, _, _, SCgradt, _ = F.int8_double_quant(grad_output.to(torch.float16)) grad_B = torch.ops.bitsandbytes.int8_scaled_mm.default( Cgrad.t().contiguous(), CAt.t(), SCgradt, SCAt, dtype=torch.float16, ) if state.threshold > 0.0 and subA is not None and subA.numel() > 0: grad_B[:, idx] += torch.matmul(grad_output.t(), subA) if req_gradA: if state.CB is not None: CB = state.CB.to(ctx.dtype_A, copy=True).mul_(state.SCB.unsqueeze(1).mul(1.0 / 127.0)) grad_A = torch.matmul(grad_output.to(ctx.dtype_A), CB).view(ctx.grad_shape) else: raise Exception("State must contain CB matrix for backward") return grad_A, grad_B, None, grad_bias, None class MatMul8bitFp(torch.autograd.Function): # For Intel CPU and XPU MatMul8bitFp is much faster (~3x) than MatMul8bitLt in finetune. # Because the MatMul8bitLt has more mechanisms in computing grad. # We don't have fast kernel for quant/dequant 8bit in CPU/XPU, so it's very slow. # We'd like to use dequant + matmul to run finetune with good performance. @staticmethod def forward(ctx, A, B, out=None, bias=None, state=MatmulLtState): if state.has_fp16_weights or state.CB is None: has_grad = getattr(B, "grad", None) is not None is_transposed = not B.is_contiguous() and B.shape[0] == B.stride(1) if is_transposed: B = B.contiguous() if (state.is_training and not has_grad) or state.CB is None or state.SCB is None: state.reset_grads() state.CB, state.SCB, _ = F.int8_vectorwise_quant(B.to(torch.float16)) B = state.CB CB = state.CB.data.to(A.dtype).mul_(state.SCB.unsqueeze(1).mul(1.0 / 127.0)) output = torch.nn.functional.linear(A, CB, bias) ctx.state = state ctx.dtype_A = A.dtype ctx.grad_shape = A.shape ctx.A = A ctx.dtype_bias = None if bias is None else bias.dtype return output @staticmethod def backward(ctx, grad_output): req_gradA, req_gradB, _, req_gradBias, _ = ctx.needs_input_grad A = ctx.A state = ctx.state grad_A = grad_B = grad_bias = None if req_gradBias: # compute grad_bias first before changing grad_output dtype grad_bias = grad_output.sum(0, dtype=ctx.dtype_bias) # Cast grad_output to fp16 if len(grad_output.shape) == 3: grad_output = grad_output.reshape(-1, grad_output.shape[-1]).contiguous() if req_gradB: grad_B = torch.matmul(A.t(), grad_output).t() if req_gradA: if state.CB is not None: CB = state.CB.to(ctx.dtype_A, copy=True).mul_(state.SCB.unsqueeze(1).mul(1.0 / 127.0)) grad_A = torch.matmul(grad_output.to(ctx.dtype_A), CB).view(ctx.grad_shape) else: raise Exception("State must contain CB matrix for backward") return grad_A, grad_B, None, grad_bias, None class MatMul4Bit(torch.autograd.Function): # forward is the same, but we added the fallback for pre-turing GPUs # backward is mostly the same, but adds one extra clause (see "elif state.CxB is not None") @staticmethod def forward(ctx, A, B, out=None, bias=None, quant_state: Optional[F.QuantState] = None): # default of pytorch behavior if inputs are empty ctx.is_empty = False if prod(A.shape) == 0: ctx.is_empty = True ctx.A = A ctx.B = B ctx.bias = bias B_shape = quant_state.shape if A.shape[-1] == B_shape[0]: return torch.empty(A.shape[:-1] + B_shape[1:], dtype=A.dtype, device=A.device) else: return torch.empty(A.shape[:-1] + B_shape[:1], dtype=A.dtype, device=A.device) # 1. Dequantize # 2. MatmulnN output = torch.nn.functional.linear(A, F.dequantize_4bit(B, quant_state).to(A.dtype).t(), bias) # 3. Save state ctx.state = quant_state ctx.dtype_A, ctx.dtype_B, ctx.dtype_bias = A.dtype, B.dtype, None if bias is None else bias.dtype if any(ctx.needs_input_grad[:2]): ctx.tensors = (None, B) else: ctx.tensors = (None, None) return output @staticmethod def backward(ctx, grad_output): if ctx.is_empty: bias_grad = None if ctx.bias is None else torch.zeros_like(ctx.bias) return torch.zeros_like(ctx.A), torch.zeros_like(ctx.B), None, bias_grad, None req_gradA, _, _, req_gradBias, _ = ctx.needs_input_grad _, B = ctx.tensors grad_A, grad_B, grad_bias = None, None, None if req_gradBias: # compute grad_bias first before changing grad_output dtype grad_bias = grad_output.sum(0, dtype=ctx.dtype_bias) # not supported by PyTorch. TODO: create work-around # if req_gradB: grad_B = torch.matmul(grad_output.t(), A) if req_gradA: grad_A = torch.matmul(grad_output, F.dequantize_4bit(B, ctx.state).to(grad_output.dtype).t()) return grad_A, grad_B, None, grad_bias, None def matmul( A: torch.Tensor, B: torch.Tensor, out: Optional[torch.Tensor] = None, state: Optional[MatmulLtState] = None, threshold=0.0, bias: Optional[torch.Tensor] = None, ): state = state or MatmulLtState() if threshold > 0.0: state.threshold = threshold # MatMul8bitLt is slower because no fast kernel for quant/dequant 8bit in CPU/XPU if state.is_training: if A.device.type in ("cpu", "xpu"): return MatMul8bitFp.apply(A, B, out, bias, state) return MatMul8bitLt.apply(A, B, out, bias, state) def matmul_4bit( A: torch.Tensor, B: torch.Tensor, quant_state: F.QuantState, out: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None, ): assert quant_state is not None if A.numel() == A.shape[-1] and A.requires_grad == False and A.device.type != "hpu": if A.shape[-1] % quant_state.blocksize != 0: warn( f"Some matrices hidden dimension is not a multiple of {quant_state.blocksize} and efficient inference kernels are not supported for these (slow). Matrix input size found: {A.shape}", ) return MatMul4Bit.apply(A, B, out, bias, quant_state) else: out = F.gemv_4bit(A, B.t(), out, state=quant_state) if bias is not None: out += bias return out else: return MatMul4Bit.apply(A, B, out, bias, quant_state) bitsandbytes/bitsandbytes/backends/000077500000000000000000000000001506752665600200405ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/__init__.py000066400000000000000000000000001506752665600221370ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/cpu/000077500000000000000000000000001506752665600206275ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/cpu/__init__.py000066400000000000000000000000001506752665600227260ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/cpu/ops.py000066400000000000000000000074311506752665600220070ustar00rootroot00000000000000import ctypes as ct import logging import torch from bitsandbytes.functional import get_ptr from ..._ops import register_kernel from ...cextension import ErrorHandlerMockBNBNativeLibrary, lib logger = logging.getLogger(__name__) # torch._int_mm for s8@s8->s32 is supported on CPU from torch 2.4+. # However, we can overflow if we use this without AVX512_VNNI support. # This is fixed in torch 2.6+, so we set this as the minimum to be safe. # For more information: https://github.com/pytorch/pytorch/pull/136942 # TODO(matthewdouglas): aarch64? if torch.__version__ >= (2, 6): @register_kernel("bitsandbytes::int8_linear_matmul", "cpu") def _(A: torch.Tensor, B: torch.Tensor): return torch._int_mm( A.reshape(-1, A.shape[-1]), B.t(), ).reshape(*A.shape[:-1], B.shape[0]) if not isinstance(lib, ErrorHandlerMockBNBNativeLibrary): @register_kernel("bitsandbytes::quantize_blockwise", "cpu") def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]: torch._check_is_size(blocksize) n = A.numel() # Only FP32 has c++ kernrl if A.dtype == torch.float32: blocks = -(n // -blocksize) absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32) out = torch.empty_like(A, dtype=torch.uint8) lib.cquantize_blockwise_cpu_fp32( get_ptr(code), get_ptr(A), get_ptr(absmax), get_ptr(out), ct.c_longlong(blocksize), ct.c_longlong(n), ) else: rem = n % blocksize has_rem = rem > 0 blocks = n // blocksize + has_rem absmax = torch.zeros((blocks,), device=A.device, dtype=torch.float32) A_reshaped = A.reshape(n) A_com = A_reshaped[: n - rem] A_com_reshaped = A_com.reshape(n // blocksize, blocksize) absmax[: blocks - has_rem] = torch.abs(A_com_reshaped).max(dim=-1)[0] scaled_A = torch.clamp(A_com_reshaped * (1 / absmax[: blocks - has_rem].view(-1, 1)), -1, 1) scaled_A = scaled_A.reshape(-1) if has_rem: absmax[-1] = torch.abs(A_reshaped[n - rem :]).max() scaled_A_rem = torch.clamp(A_reshaped[n - rem :] * (1 / absmax[-1]), -1, 1) scaled_A = torch.cat([scaled_A, scaled_A_rem], dim=0) diff = torch.abs(scaled_A.unsqueeze(-1) - code.to(scaled_A.device)) out = torch.argmin(diff, dim=-1).to(torch.uint8).to(scaled_A.device).reshape(A.shape) return out, absmax @register_kernel("bitsandbytes::dequantize_blockwise", "cpu") def _( A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype ) -> torch.Tensor: torch._check_is_size(blocksize) torch._check(A.dtype == torch.uint8, lambda: f"A must be uint8, got {A.dtype}") # Only FP32 has c++ kernrl if dtype == torch.float32: out = torch.empty_like(A, dtype=dtype) lib.cdequantize_blockwise_cpu_fp32( get_ptr(code), get_ptr(A), get_ptr(absmax), get_ptr(out), ct.c_longlong(blocksize), ct.c_longlong(A.numel()), ) else: out = code[A.reshape(-1).int()] blocks = out.shape[-1] // blocksize res = out.shape[-1] % blocksize if res != 0: out = torch.nn.functional.pad(out, (0, blocksize - res), mode="constant", value=0) out = (out.view(-1, blocksize) * absmax.view(-1, 1)).to(dtype).reshape(-1) out = out[: blocks * blocksize + res] out = out.reshape(A.shape) return out bitsandbytes/bitsandbytes/backends/cuda/000077500000000000000000000000001506752665600207545ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/cuda/__init__.py000066400000000000000000000000001506752665600230530ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/cuda/ops.py000066400000000000000000000607561506752665600221450ustar00rootroot00000000000000from collections.abc import Sequence import ctypes as ct from math import prod from typing import Optional import torch from bitsandbytes.functional import CUBLAS_Context, _cuda_device_of, _get_tensor_stream, get_ptr from ..._ops import register_kernel from ...cextension import HIP_ENVIRONMENT, lib @register_kernel("bitsandbytes::int8_linear_matmul", "cuda") def _(A: torch.Tensor, B: torch.Tensor): out = torch.empty((*A.shape[:-1], B.shape[0]), device=A.device, dtype=torch.int32) return _int8_linear_matmul_impl(A, B, out) @register_kernel("bitsandbytes::int8_linear_matmul.out", "cuda") def _(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor): _int8_linear_matmul_impl(A, B, out) def _int8_linear_matmul_impl(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor): A, B = B, A shapeA = A.shape shapeB = B.shape torch._check(A.dtype == torch.int8, lambda: "B must be int8") torch._check(B.dtype == torch.int8, lambda: "A must be int8") torch._check(A.ndim == 2, lambda: "Only two dimensional matrices are supported for argument B") torch._check(B.ndim in [2, 3], lambda: "Only two or three dimensional matrices are supported for argument A") torch._check(prod(shapeB) > 0, lambda: f"Input tensor dimensions need to be > 0: {shapeB}") torch._check(out.dtype == torch.int32) shapeC = (*shapeB[:-1], shapeA[0]) torch._check(out.shape == shapeC, lambda: f"Output shape {out.shape} does not match expected shape {shapeC}") k, m = shapeA n = prod(shapeB[:-1]) lda = shapeA[-1] # Weights (outputs, inputs) ldb = shapeB[-1] # Activations (batch, tokens, inputs) ldc = shapeC[-1] # Output (batch, tokens, outputs) torch._check( lda == ldb, lambda: f"int8_linear_matmul only supports B^T @ A. Inner dimensions do not match: B @ A = {shapeB} @ {shapeA}", ) # cuBLASLt does not support int8 matmul with inner dimensions that are not divisible by 4. # We'll fall back to a slower fp32 calculation in this circumstance. # Fortunately, this should not be very common. if lda % 4 != 0: result = torch.matmul(B.float(), A.float().t()).to(torch.int32) return out.copy_(result) with _cuda_device_of(A): ctx = CUBLAS_Context.get_instance().get_context(A.device) ptrA = get_ptr(A) ptrB = get_ptr(B) ptrC = get_ptr(out) ptrRowScale = None m = ct.c_int32(m) n = ct.c_int32(n) k = ct.c_int32(k) lda = ct.c_int32(lda) ldb = ct.c_int32(ldb) ldc = ct.c_int32(ldc) stream = _get_tensor_stream(A) has_error = lib.cigemmlt_32(ctx, m, n, k, ptrA, ptrB, ptrC, ptrRowScale, lda, ldb, ldc, stream) if has_error: if has_error == 100: # `ERR_NOT_IMPLEMENTED` is defined as 100 in `ops.cu` # TODO: Warn and implement a fallback to fp32 compute? raise NotImplementedError("int8_linear_matmul not implemented!") else: raise RuntimeError( f"cublasLt ran into an error!\n\t{shapeA=}, {shapeB=}, {shapeC=}\n\t{(lda, ldb, ldc)=}\n\t{(m, n, k)=}" ) return out @register_kernel("bitsandbytes::int8_mm_dequant", "cuda") def _( A: torch.Tensor, row_stats: torch.Tensor, col_stats: torch.Tensor, dtype: Optional[torch.dtype] = None, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: torch._check(A.dtype == torch.int32, lambda: f"A must be int32, got {A.dtype}") torch._check(row_stats.dtype == torch.float32, lambda: f"row_stats must be float32, got {row_stats.dtype}") torch._check(col_stats.dtype == torch.float32, lambda: f"col_stats must be float32, got {col_stats.dtype}") # Note: cuda kernel only currently supports fp16 output. # We'll later cast to desired dtype if needed. out = torch.empty_like(A, dtype=torch.float16) ptrA = get_ptr(A) ptrOut = get_ptr(out) ptrRowStats = get_ptr(row_stats) ptrColStats = get_ptr(col_stats) numRows = ct.c_int32(prod(A.shape[:-1])) numCols = ct.c_int32(A.shape[-1]) # Note: fused bias in the kernel is only supported for fp16 # TODO(matthewdouglas): Consider supporting bf16 fused bias ptrBias = get_ptr(bias) if bias is not None and bias.dtype == torch.float16 else None with _cuda_device_of(A): lib.cdequant_mm_int32_fp16( ptrA, ptrRowStats, ptrColStats, ptrOut, ptrBias, numRows, numCols, _get_tensor_stream(A) ) # Add bias separately if not fused in kernel if bias is not None and bias.dtype != torch.float16: out.add_(bias) return out.to(dtype or torch.float16) @register_kernel("bitsandbytes::int8_vectorwise_quant", "cuda") def _(A: torch.Tensor, threshold=0.0): torch._check(A.dtype == torch.float16, lambda: f"A must be float16, got {A.dtype}") torch._check(threshold >= 0.0, lambda: "threshold must be non-negative") rows = prod(A.shape[:-1]) cols = A.shape[-1] row_stats = torch.empty(rows, device=A.device, dtype=torch.float32) out_row = torch.empty(A.shape, device=A.device, dtype=torch.int8) outlier_cols = None if threshold > 0.0: # TODO we could improve perf of this outliers = A.abs() >= threshold if outliers.any(): outlier_cols = torch.argwhere(outliers.any(dim=0)).view(-1) else: # Needed for torch.compile support. outlier_cols = torch.empty(0, device=A.device, dtype=torch.int64) with _cuda_device_of(A): lib.cint8_vector_quant( get_ptr(A), get_ptr(out_row), get_ptr(row_stats), ct.c_float(threshold), ct.c_int32(rows), ct.c_int32(cols), _get_tensor_stream(A), ) # Zero out values from outlier columns across all rows. # The kernel will handle this for outliers themselves, so we can optimize for rows=1. if rows > 1 and outlier_cols is not None: out_row[:, outlier_cols] = 0 return out_row, row_stats, outlier_cols @register_kernel("bitsandbytes::int8_double_quant", "cuda") def _( A: torch.Tensor, threshold=0.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: # Use CUDA kernel for rowwise and COO tensor quant_row, row_stats, outlier_cols = torch.ops.bitsandbytes.int8_vectorwise_quant.default( A, threshold=threshold, ) # PyTorch impl for colwise col_stats, outlier_mask = _get_col_absmax(A, threshold=threshold) if threshold > 0.0 and outlier_mask is not None: A = A.masked_fill(outlier_mask, 0.0) quant_col = torch.round(A.mul(127.0) / col_stats.unsqueeze(0)).to(torch.int8) return quant_row, quant_col, row_stats, col_stats.flatten().float(), outlier_cols def _get_col_absmax( A: torch.Tensor, threshold=0.0, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: torch._check(A.is_floating_point()) outlier_mask = None absA = A.abs().view(-1, A.shape[-1]) if threshold > 0.0: # Filter outliers from stats when enabled outlier_mask = absA >= threshold absA.masked_fill_(outlier_mask, 0.0) # shape [cols]; unsqueeze(0) gives [1,cols] col_stats = absA.amax(dim=0, keepdim=False).float() return col_stats, outlier_mask @register_kernel("bitsandbytes::quantize_blockwise", "cuda") def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]: torch._check_is_size(blocksize) if HIP_ENVIRONMENT: torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128]) else: torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64]) torch._check(code.dtype == torch.float32, lambda: f"code must be float32, got {code.dtype}") n = A.numel() blocks = -(n // -blocksize) absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32) out = torch.empty_like(A, dtype=torch.uint8) with _cuda_device_of(A): args = ( get_ptr(code), get_ptr(A), get_ptr(absmax), get_ptr(out), ct.c_int32(blocksize), ct.c_int(A.numel()), ) if A.dtype == torch.float16: lib.cquantize_blockwise_fp16(*args) elif A.dtype == torch.bfloat16: lib.cquantize_blockwise_bf16(*args) elif A.dtype == torch.float32: lib.cquantize_blockwise_fp32(*args) else: raise ValueError(f"Blockwise quantization only supports 16/32-bit floats, but got {A.dtype}") return out, absmax @register_kernel("bitsandbytes::dequantize_blockwise", "cuda") def _(A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype) -> torch.Tensor: out = torch.empty_like(A, dtype=dtype) _dequantize_blockwise_impl(A, absmax, code, blocksize, dtype, out=out) return out @register_kernel("bitsandbytes::dequantize_blockwise.out", "cuda") def _( A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype, out: torch.Tensor, ) -> None: torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}") torch._check(out.shape == A.shape, lambda: f"Expected out.shape == {A.shape}, got {out.shape}") _dequantize_blockwise_impl(A, absmax, code, blocksize, dtype, out=out) def _dequantize_blockwise_impl( A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype, out: torch.Tensor ) -> None: if HIP_ENVIRONMENT: torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128]) else: torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64]) torch._check(A.dtype == torch.uint8, lambda: f"A must be uint8, got {A.dtype}") torch._check( dtype in [torch.float16, torch.bfloat16, torch.float32], lambda: f"Blockwise dequantization only supports 16bit/32bit floating types, got {dtype}", ) with _cuda_device_of(A): args = ( get_ptr(code), get_ptr(A), get_ptr(absmax), get_ptr(out), ct.c_int(blocksize), ct.c_int(A.numel()), _get_tensor_stream(A), ) if dtype == torch.float16: lib.cdequantize_blockwise_fp16(*args) elif dtype == torch.bfloat16: lib.cdequantize_blockwise_bf16(*args) elif dtype == torch.float32: lib.cdequantize_blockwise_fp32(*args) @register_kernel("bitsandbytes::quantize_4bit", "cuda") def _( A: torch.Tensor, blocksize: int, quant_type: str, quant_storage: torch.dtype ) -> tuple[torch.Tensor, torch.Tensor]: if HIP_ENVIRONMENT: torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128]) else: torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64]) torch._check(quant_type in ["fp4", "nf4"]) torch._check( A.dtype in [torch.bfloat16, torch.float16, torch.float32], lambda: f"Blockwise 4bit quantization only supports 16/32-bit floats, but got {A.dtype}", ) n = A.numel() blocks = -(n // -blocksize) absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32) out = torch.empty(((n + 1) // (quant_storage.itemsize * 2), 1), device=A.device, dtype=quant_storage) with _cuda_device_of(A): args = ( None, get_ptr(A), get_ptr(absmax), get_ptr(out), ct.c_int32(blocksize), ct.c_int(n), ) if A.dtype == torch.bfloat16: if quant_type == "fp4": lib.cquantize_blockwise_bf16_fp4(*args) else: lib.cquantize_blockwise_bf16_nf4(*args) elif A.dtype == torch.float16: if quant_type == "fp4": lib.cquantize_blockwise_fp16_fp4(*args) else: lib.cquantize_blockwise_fp16_nf4(*args) elif A.dtype == torch.float32: if quant_type == "fp4": lib.cquantize_blockwise_fp32_fp4(*args) else: lib.cquantize_blockwise_fp32_nf4(*args) return out, absmax @register_kernel("bitsandbytes::dequantize_4bit", "cuda") def _( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, shape: Sequence[int], dtype: torch.dtype, ) -> torch.Tensor: out = torch.empty(shape, dtype=dtype, device=A.device) _dequantize_4bit_impl(A, absmax, blocksize, quant_type, dtype, out=out) return out @register_kernel("bitsandbytes::dequantize_4bit.out", "cuda") def _( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, shape: Sequence[int], dtype: torch.dtype, out: torch.Tensor, ) -> None: torch._check(out.shape == shape, lambda: f"Expected out.shape == {shape}, got {out.shape}") torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}") _dequantize_4bit_impl(A, absmax, blocksize, quant_type, dtype, out=out) def _dequantize_4bit_impl( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, dtype: torch.dtype, out: torch.Tensor, ) -> None: if HIP_ENVIRONMENT: torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128]) else: torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64]) torch._check(quant_type in ["fp4", "nf4"]) torch._check( dtype in [torch.bfloat16, torch.float16, torch.float32], lambda: f"Blockwise 4bit dequantization only supports 16/32-bit floats, but got {dtype}", ) with _cuda_device_of(A): args = ( None, get_ptr(A), get_ptr(absmax), get_ptr(out), ct.c_int(blocksize), ct.c_int(out.numel()), _get_tensor_stream(A), ) if out.dtype == torch.bfloat16: if quant_type == "fp4": lib.cdequantize_blockwise_bf16_fp4(*args) else: lib.cdequantize_blockwise_bf16_nf4(*args) elif out.dtype == torch.float16: if quant_type == "fp4": lib.cdequantize_blockwise_fp16_fp4(*args) else: lib.cdequantize_blockwise_fp16_nf4(*args) elif out.dtype == torch.float32: if quant_type == "fp4": lib.cdequantize_blockwise_fp32_fp4(*args) else: lib.cdequantize_blockwise_fp32_nf4(*args) @register_kernel("bitsandbytes::gemv_4bit", "cuda") def _( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int ) -> torch.Tensor: shape = (*A.shape[:-1], shapeB[0]) out = torch.empty(shape, device=A.device, dtype=A.dtype) _gemv_4bit_impl(A, B, shapeB, absmax, code, blocksize, out=out) return out @register_kernel("bitsandbytes::gemv_4bit.out", "cuda") def _( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int, out: torch.Tensor, ) -> None: torch._check( out.shape == (*A.shape[:-1], shapeB[0]), lambda: f"Expected out.shape == {(*A.shape[:-1], shapeB[0])}, got {out.shape}", ) torch._check(out.dtype == A.dtype, lambda: f"Expected out.dtype == {A.dtype}, got {out.dtype}") _gemv_4bit_impl(A, B, shapeB, absmax, code, blocksize, out=out) def _gemv_4bit_impl( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int, out: torch.Tensor, ) -> None: torch._check_is_size(blocksize) # Note: these checks are not strictly necessary, and cost more than they are worth, so they are commented out for now. # torch._check( # A.numel() == A.size(-1), # lambda: f"A must be a vector with leading dimensions of 1, got {A.shape}", # ) # torch._check( # A.dtype in [torch.float16, torch.bfloat16, torch.float32], # lambda: f"A must be float16, bfloat16, or float32, got {A.dtype}", # ) # torch._check( # B.dtype in [torch.uint8, torch.bfloat16, torch.float16, torch.float32], # lambda: f"B must be backed by storage of type uint8, bfloat16, float16, or float32, got {B.dtype}", # ) # torch._check(absmax.dtype == torch.float32, lambda: f"absmax must be float32, got {absmax.dtype}") # torch._check(code.dtype == torch.float32, lambda: f"code must be float32, got {code.dtype}") m = ct.c_int32(shapeB[0]) n = ct.c_int32(1) k = ct.c_int32(shapeB[1]) lda = m ldb = ct.c_int32((A.shape[-1] + 1) // 2) ldc = m stream = _get_tensor_stream(A) with _cuda_device_of(A): if A.dtype == torch.float16: lib.cgemm_4bit_inference_naive_fp16( m, n, k, get_ptr(A), get_ptr(B), get_ptr(absmax), get_ptr(code), get_ptr(out), lda, ldb, ldc, ct.c_int32(blocksize), stream, ) elif A.dtype == torch.bfloat16: lib.cgemm_4bit_inference_naive_bf16( m, n, k, get_ptr(A), get_ptr(B), get_ptr(absmax), get_ptr(code), get_ptr(out), lda, ldb, ldc, ct.c_int32(blocksize), stream, ) elif A.dtype == torch.float32: lib.cgemm_4bit_inference_naive_fp32( m, n, k, get_ptr(A), get_ptr(B), get_ptr(absmax), get_ptr(code), get_ptr(out), lda, ldb, ldc, ct.c_int32(blocksize), stream, ) """C FUNCTIONS FOR OPTIMIZERS""" str2optimizer32bit = { "adam": ( lib.cadam32bit_grad_fp32, lib.cadam32bit_grad_fp16, lib.cadam32bit_grad_bf16, ), "momentum": ( lib.cmomentum32bit_grad_32, lib.cmomentum32bit_grad_16, ), "rmsprop": ( lib.crmsprop32bit_grad_32, lib.crmsprop32bit_grad_16, ), "lion": ( lib.clion32bit_grad_fp32, lib.clion32bit_grad_fp16, lib.clion32bit_grad_bf16, ), "adagrad": ( lib.cadagrad32bit_grad_32, lib.cadagrad32bit_grad_16, ), "lamb": ( lib.cadam32bit_grad_fp32, lib.cadam32bit_grad_fp16, lib.cadam32bit_grad_bf16, ), "ademamix": ( lib.cademamix32bit_grad_fp32, lib.cademamix32bit_grad_fp16, lib.cademamix32bit_grad_bf16, ), } str2optimizer8bit_blockwise = { "adam": ( lib.cadam_8bit_blockwise_grad_fp32, lib.cadam_8bit_blockwise_grad_fp16, lib.cadam_8bit_blockwise_grad_bf16, ), "momentum": ( lib.cmomentum_8bit_blockwise_grad_fp32, lib.cmomentum_8bit_blockwise_grad_fp16, lib.cmomentum_8bit_blockwise_grad_bf16, ), "rmsprop": ( lib.crmsprop_8bit_blockwise_grad_fp32, lib.crmsprop_8bit_blockwise_grad_fp16, lib.crmsprop_8bit_blockwise_grad_bf16, ), "lion": ( lib.clion_8bit_blockwise_grad_fp32, lib.clion_8bit_blockwise_grad_fp16, lib.clion_8bit_blockwise_grad_bf16, ), "adagrad": ( lib.cadagrad_8bit_blockwise_grad_fp32, lib.cadagrad_8bit_blockwise_grad_fp16, lib.cadagrad_8bit_blockwise_grad_bf16, ), "ademamix": ( lib.cademamix_8bit_blockwise_grad_fp32, lib.cademamix_8bit_blockwise_grad_fp16, lib.cademamix_8bit_blockwise_grad_bf16, ), } def _optimizer_update_32bit_impl( optimizer_name: str, g: torch.Tensor, p: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], unorm_vec: Optional[torch.Tensor], max_unorm: float, param_norm: float, beta1: float, beta2: float, beta3: float, alpha: float, eps: float, weight_decay: float, step: int, lr: float, gnorm_scale: float, skip_zeros=False, ) -> None: optim_fns = str2optimizer32bit.get(optimizer_name, None) if optim_fns is None: raise ValueError( f"Unsupported optimizer name: {optimizer_name}. Supported optimizers: {list(str2optimizer8bit_blockwise.keys())}" ) if g.dtype == torch.float32: optim_func = optim_fns[0] elif g.dtype == torch.float16: optim_func = optim_fns[1] elif g.dtype == torch.bfloat16 and len(optim_fns) == 3: optim_func = optim_fns[2] else: raise ValueError( f"Gradient+optimizer bit data type combination not supported: grad {g.dtype}, optimizer {state1.dtype}", ) with _cuda_device_of(g): optim_func( get_ptr(g), get_ptr(p), get_ptr(state1), get_ptr(state2), get_ptr(unorm_vec), ct.c_float(max_unorm), ct.c_float(param_norm), ct.c_float(beta1), ct.c_float(beta2), ct.c_float(beta3), ct.c_float(alpha), ct.c_float(eps), ct.c_float(weight_decay), ct.c_int32(step), ct.c_float(lr), ct.c_float(gnorm_scale), ct.c_bool(skip_zeros), ct.c_int32(g.numel()), ) def _optimizer_update_8bit_blockwise_impl( optimizer_name: str, g: torch.Tensor, p: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], beta1: float, beta2: float, beta3: float, alpha: float, eps: float, step: int, lr: float, qmap1: torch.Tensor, qmap2: Optional[torch.Tensor], absmax1: torch.Tensor, absmax2: Optional[torch.Tensor], weight_decay: float, gnorm_scale: float, skip_zeros=False, ) -> None: # torch._check( # g.numel() == p.numel(), # lambda: f"g and p must have the same number of elements, got {g.numel()} and {p.numel()}", # ) # compute_dtypes = [torch.float16, torch.bfloat16, torch.float32] # torch._check( # g.dtype in compute_dtypes, # lambda: f"g must be bfloat16, float16, or float32, got {g.dtype}", # ) # torch._check( # g.dtype == p.dtype, # lambda: f"Expected all tensors to have the same dtype, got g.dtype={g.dtype}, p.dtype={p.dtype}", # ) # torch._check( # state1.dtype == torch.uint8, # lambda: f"state1 must be uint8, got {state1.dtype}", # ) # torch._check( # qmap1.dtype == absmax1.dtype == torch.float32, # lambda: f"Expected qmap1 and absmax1 to be float32, got qmap1.dtype={qmap1.dtype}, absmax1.dtype={absmax1.dtype}", # ) # if state2 is not None: # torch._check( # state2.dtype == torch.uint8, # lambda: f"state2 must be uint8, got {state2.dtype}", # ) # torch._check( # qmap2.dtype == absmax2.dtype == torch.float32, # lambda: f"Expected qmap2 and absmax2 to be float32, got qmap2.dtype={qmap2.dtype}, absmax2.dtype={absmax2.dtype}", # ) optimizer_fns = str2optimizer8bit_blockwise.get(optimizer_name) if optimizer_fns is None: raise ValueError( f"Unsupported optimizer name: {optimizer_name}. Supported optimizers: {list(str2optimizer8bit_blockwise.keys())}" ) if g.dtype == torch.float32: optimizer_fn = optimizer_fns[0] elif g.dtype == torch.float16: optimizer_fn = optimizer_fns[1] elif g.dtype == torch.bfloat16: optimizer_fn = optimizer_fns[2] else: raise ValueError( f"Unsupported gradient dtype: {g.dtype}. Supported dtypes: torch.float32, torch.float16, torch.bfloat16" ) with _cuda_device_of(g): optimizer_fn( get_ptr(p), get_ptr(g), get_ptr(state1), get_ptr(state2), ct.c_float(beta1), ct.c_float(beta2), ct.c_float(beta3), ct.c_float(alpha), ct.c_float(eps), ct.c_int32(step), ct.c_float(lr), get_ptr(qmap1), get_ptr(qmap2), get_ptr(absmax1), get_ptr(absmax2), ct.c_float(weight_decay), ct.c_float(gnorm_scale), ct.c_bool(skip_zeros), ct.c_int32(g.numel()), ) register_kernel("bitsandbytes::optimizer_update_8bit_blockwise", "cuda")(_optimizer_update_8bit_blockwise_impl) register_kernel("bitsandbytes::optimizer_update_32bit", "cuda")(_optimizer_update_32bit_impl) bitsandbytes/bitsandbytes/backends/default/000077500000000000000000000000001506752665600214645ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/default/__init__.py000066400000000000000000000000001506752665600235630ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/default/ops.py000066400000000000000000000430211506752665600226370ustar00rootroot00000000000000from collections.abc import Sequence from math import prod, sqrt from typing import Optional import torch from ..._ops import register_kernel from ..utils import CODE @register_kernel("bitsandbytes::int8_mm_dequant", "default") def _( A: torch.Tensor, row_stats: torch.Tensor, col_stats: torch.Tensor, dtype: Optional[torch.dtype] = None, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: torch._check(A.dtype == torch.int32, lambda: f"A must be int32, got {A.dtype}") torch._check(row_stats.dtype == torch.float32, lambda: f"row_stats must be float32, got {row_stats.dtype}") torch._check(col_stats.dtype == torch.float32, lambda: f"col_stats must be float32, got {col_stats.dtype}") A_calc = A.view(-1, A.shape[-1]) row_stats = row_stats.reshape(-1).unsqueeze(-1) col_stats = col_stats.reshape(-1).unsqueeze(0) out = A_calc * (row_stats * col_stats) * 6.200124e-05 if bias is not None: out += bias return out.to(dtype or torch.float16) @register_kernel("bitsandbytes::int8_mixed_scaled_mm", "default") def _( A: torch.Tensor, CA: torch.Tensor, CB: torch.Tensor, SCA: torch.Tensor, SCB: torch.Tensor, outlier_cols: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: subB = None if outlier_cols is not None and outlier_cols.numel(): # Extract the inputs with outliers in original precision subA = A[:, outlier_cols].contiguous() # Dequantize the corresponding weight columns subB = ( torch.ops.bitsandbytes.int8_vectorwise_dequant.default(CB[:, outlier_cols].contiguous(), SCB) .to(A.dtype) .t() ) # TODO: if state.has_fp16_weights: subB = B[:, outlier_cols].t() else: # Needed for torch.compile when there are no outliers. subA = torch.empty(0, device=A.device, dtype=A.dtype) # Int8 Matmul + Dequant + Bias output = torch.ops.bitsandbytes.int8_scaled_mm.default(CA, CB, SCA, SCB, bias=bias, dtype=A.dtype) if subB is not None: # Add the outlier columns back to the output output = output.addmm(subA, subB) return output, subA @register_kernel("bitsandbytes::int8_scaled_mm", "default") def _( A: torch.Tensor, B: torch.Tensor, row_stats: torch.Tensor, col_stats: torch.Tensor, bias: Optional[torch.Tensor] = None, dtype: Optional[torch.dtype] = None, ) -> torch.Tensor: out_i32 = torch.ops.bitsandbytes.int8_linear_matmul.default(A, B) return torch.ops.bitsandbytes.int8_mm_dequant.default( out_i32, row_stats, col_stats, dtype=dtype or torch.float16, bias=bias, ) @register_kernel("bitsandbytes::int8_linear_matmul", "default") def _(A: torch.Tensor, B: torch.Tensor): return _int8_linear_matmul_impl(A, B) @register_kernel("bitsandbytes::int8_linear_matmul.out", "default") def _(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor): torch._check(out.dtype == torch.int32) _int8_linear_matmul_impl(A, B, out) def _int8_linear_matmul_impl(A: torch.Tensor, B: torch.Tensor, out: Optional[torch.Tensor] = None): # Naive implementation: perform matmul in fp32 result = torch.matmul(A.float(), B.float().t()).to(torch.int32) if out is not None: result = out.copy_(result) return result @register_kernel("bitsandbytes::int8_vectorwise_quant", "default") def _(A: torch.Tensor, threshold=0.0): rows = prod(A.shape[:-1]) outlier_cols = None outlier_restore = None if threshold > 0.0: outliers = A.abs() >= threshold if outliers.any(): # Determine which columns contain outliers, and zero out the # outliers ahead of quantization. We need to keep a backup of these # outliers to restore them after quantization. outlier_cols = torch.argwhere(outliers.any(dim=0)).view(-1) outlier_restore = A[outliers].clone() A[outliers] = 0 else: # Needed for torch.compile support. outlier_cols = torch.empty(0, device=A.device, dtype=torch.int64) # Get absmax for each row. row_stats = torch.max(A.abs(), dim=1).values.float() # Quantize row-wise to int8. out_row = torch.round(A * (127.0 / row_stats.unsqueeze(-1))).to(torch.int8) # Zero out values from outlier columns across all rows. if rows > 1 and outlier_cols is not None: out_row[:, outlier_cols] = 0 # Restore outliers. if outlier_restore is not None: A[outliers] = outlier_restore return out_row, row_stats, outlier_cols @register_kernel("bitsandbytes::quantize_blockwise", "default") def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]: torch._check_is_size(blocksize) n = A.numel() rem = n % blocksize has_rem = rem > 0 blocks = n // blocksize + has_rem absmax = torch.zeros((blocks,), device=A.device, dtype=torch.float32) A_reshaped = A.reshape(n) A_com = A_reshaped[: n - rem] A_com_reshaped = A_com.reshape(n // blocksize, blocksize) absmax[: blocks - has_rem] = torch.abs(A_com_reshaped).max(dim=-1)[0] scaled_A = torch.clamp(A_com_reshaped * (1 / absmax[: blocks - has_rem].view(-1, 1)), -1, 1) scaled_A = scaled_A.reshape(-1) if has_rem: absmax[-1] = torch.abs(A_reshaped[n - rem :]).max() scaled_A_rem = torch.clamp(A_reshaped[n - rem :] * (1 / absmax[-1]), -1, 1) scaled_A = torch.cat([scaled_A, scaled_A_rem], dim=0) diff = torch.abs(scaled_A.unsqueeze(-1) - code.to(scaled_A.device)) out = torch.argmin(diff, dim=-1).to(torch.uint8).to(scaled_A.device).reshape(A.shape) return out, absmax @register_kernel("bitsandbytes::dequantize_blockwise", "default") def _(A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype) -> torch.Tensor: torch._check_is_size(blocksize) torch._check(A.dtype == torch.uint8, lambda: f"A must be uint8, got {A.dtype}") out = code[A.reshape(-1).int()] blocks = out.shape[-1] // blocksize res = out.shape[-1] % blocksize if res != 0: out = torch.nn.functional.pad(out, (0, blocksize - res), mode="constant", value=0) out = (out.view(-1, blocksize) * absmax.view(-1, 1)).to(dtype).reshape(-1) out = out[: blocks * blocksize + res] out = out.reshape(A.shape) return out @register_kernel("bitsandbytes::quantize_4bit", "default") def _( A: torch.Tensor, blocksize: int, quant_type: str, quant_storage: torch.dtype ) -> tuple[torch.Tensor, torch.Tensor]: torch._check_is_size(blocksize) torch._check(quant_type in ("nf4", "fp4"), lambda: f"quant_type must be nf4 or fp4, got {quant_type}") torch._check( A.dtype in [torch.bfloat16, torch.float16, torch.float32], lambda: f"Blockwise 4bit quantization only supports 16/32-bit floats, but got {A.dtype}", ) n = A.numel() full_blocks = n // blocksize rem = n % blocksize blocks = full_blocks + 1 if rem else full_blocks absmax = torch.zeros((blocks,), device=A.device, dtype=torch.float32) A_flattened = A.reshape(n) # Scale full blocks of the tensor to [-1, 1] A_full_blocks = A_flattened[: n - rem].reshape(n // blocksize, blocksize) absmax[:full_blocks] = torch.abs(A_full_blocks).max(dim=-1)[0] scaled = torch.clamp(A_full_blocks * (1 / absmax[:full_blocks].view(-1, 1)), -1, 1).reshape(-1) # Scale any partial block if rem: A_rem = A_flattened[-rem:] absmax[-1] = torch.abs(A_rem).max() scaled_rem = torch.clamp(A_rem * (1 / absmax[-1]), -1, 1) scaled = torch.cat([scaled, scaled_rem], dim=0) # Quantize with the lookup table code = CODE[quant_type].to(scaled.device).to(scaled.dtype) quantized = torch.argmin(torch.abs(scaled.view(-1, 1) - code), dim=-1, keepdim=True).to(torch.uint8) # Pack two quantized values per byte packed = quantized[::2] << 4 | quantized[1::2] if quant_storage != torch.uint8: packed = packed.squeeze().view(quant_storage).unsqueeze(1) return packed, absmax.float() @register_kernel("bitsandbytes::dequantize_4bit", "default") def _( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, shape: Sequence[int], dtype: torch.dtype, ) -> torch.Tensor: torch._check_is_size(blocksize) torch._check(quant_type in ("nf4", "fp4"), lambda: f"quant_type must be nf4 or fp4, got {quant_type}") torch._check( dtype in [torch.bfloat16, torch.float16, torch.float32], lambda: f"Blockwise 4bit dequantization only supports 16/32-bit floats, but got {dtype}", ) # Enable non uint8 dtype if A.dtype != torch.uint8: A = A.view(torch.uint8) A = A.reshape(-1) # Map nf4 to [-1, 1] out_dq = torch.empty(A.size(0) * 2, dtype=torch.int32, device=A.device) n = out_dq.numel() out_dq[1::2] = A & 0xF out_dq[::2] = A >> 4 # code is fp32, cast to dtype to avoid the mismatch issue code = CODE[quant_type].to(dtype).to(A.device) out_dq = code[out_dq] # Apply scales if out_dq.numel() != n: assert out_dq.numel() == n + 1 out_dq = torch.narrow(out_dq, 0, 0, n) blocks = n // blocksize blocks += 1 if n % blocksize > 0 else 0 rem = n % blocksize has_rem = rem > 0 out = torch.empty(shape, dtype=dtype, device=A.device).reshape(-1) if has_rem: out[: n - rem] = (out_dq[: n - rem].view(-1, blocksize) * absmax[: blocks - has_rem].view(-1, 1)).reshape(-1) out[n - rem :] = out_dq[n - rem :] * absmax[-1] else: out = out_dq.view(-1, blocksize) * absmax.view(-1, 1) out = out.reshape(-1, *shape[1:]).to(dtype) return out @register_kernel("bitsandbytes::gemv_4bit", "default") def _( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int, ) -> torch.Tensor: # Applied from dequantize_4bit quant_type = "fp4" if code[1] > 0 else "nf4" B_dq = torch.ops.bitsandbytes.dequantize_4bit.default(B, absmax, blocksize, quant_type, shapeB, A.dtype) return torch.nn.functional.linear( A, B_dq, bias=None, ) MOMENTUM = 0 RMSPROP = 1 ADAGRAD = 2 ADAM = 3 # LION should be larger than MOMENTUM, RMSPROP, ADAGRAD due to comparison in kernels LION = 4 ADEMAMIX = 5 name2optimizer_id = { "momentum": MOMENTUM, "rmsprop": RMSPROP, "adagrad": ADAGRAD, "adam": ADAM, "lion": LION, "ademamix": ADEMAMIX, } @torch.compile def _optimizer_precondition_32bit( g: torch.Tensor, p: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], unorm_vec: torch.Tensor, beta1: float, beta2: float, eps: float, weight_decay: float, step: int, lr: float, gnorm_scale: float, optimizer_id: int, ): """Preprocessing optimizer, computing update norm""" g_vals = gnorm_scale * g if optimizer_id == 3: # ADAM correction1 = 1.0 / (1.0 - beta1**step) correction2 = 1.0 / (1.0 - beta2**step) s1_vals = state1 * beta1 + (1.0 - beta1) * g_vals s2_vals = state2 * beta2 + (1.0 - beta2) * g_vals * g_vals s1_vals = s1_vals * correction1 s2_vals = s2_vals * correction2 update_vals = s1_vals / (torch.sqrt(s2_vals) + eps) update_norm = update_vals * update_vals elif optimizer_id == 5: # ADEMAMIX update_norm = state1 elif optimizer_id == 0: # MOMENTUM if step == 1: s1_vals = g_vals else: s1_vals = state1 * beta1 + g_vals update_norm = s1_vals * s1_vals elif optimizer_id == 4: # LION s1_vals = state1 * beta2 + (1.0 - beta2) * g_vals update_norm = s1_vals elif optimizer_id == 1: # RMSPROP s1_vals = state1 * beta1 + (1.0 - beta1) * g_vals * g_vals update_vals = g_vals / (torch.sqrt(s1_vals) + eps) update_norm = update_vals * update_vals elif optimizer_id == 2: # ADAGRAD s1_vals = state1 + g_vals * g_vals update_vals = g_vals / (torch.sqrt(s1_vals) + eps) update_norm = update_vals * update_vals total_norm = torch.sum(update_norm) unorm_vec.add_(total_norm) @torch.compile def _optimizer_update_32bit( g: torch.Tensor, p: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], unorm_vec: Optional[torch.Tensor], max_unorm: float, param_norm: float, beta1: float, beta2: float, beta3: float, alpha: float, eps: float, weight_decay: float, step: int, lr: float, gnorm_scale: float, optimizer_id: int, ): """Unified optimizer update kernel""" p_vals = p.float() g_vals = (gnorm_scale * g).float() if optimizer_id in [0, 1, 2, 4] and weight_decay > 0.0: g_vals = g_vals + p_vals * weight_decay update_scale = 1.0 if max_unorm > 0.0: current_unorm = torch.sqrt(unorm_vec) if optimizer_id in [0, 1, 2, 4]: # 1-state optimizers if current_unorm > max_unorm * param_norm + eps: update_scale = (max_unorm * param_norm + eps) / current_unorm else: # 2-state optimizers if current_unorm > max_unorm * param_norm: update_scale = (max_unorm * param_norm) / current_unorm if optimizer_id == 3: # ADAM s1_vals = state1 * beta1 + (1.0 - beta1) * g_vals s2_vals = state2 * beta2 + (1.0 - beta2) * g_vals * g_vals correction1 = 1.0 - beta1**step correction2 = sqrt(1.0 - beta2**step) step_size = -lr * correction2 / correction1 if weight_decay > 0.0: p_vals = p_vals * (1.0 - lr * weight_decay) update_val = update_scale * step_size * (s1_vals / (torch.sqrt(s2_vals) + eps * correction2)) p_vals = p_vals + update_val state1.copy_(s1_vals) state2.copy_(s2_vals) elif optimizer_id == 5: # ADEMAMIX s1_vals = state1[0] s3_vals = state1[1] s2_vals = state2 m1 = s1_vals * beta1 + (1.0 - beta1) * g_vals m2 = s3_vals * beta3 + (1.0 - beta3) * g_vals nu = s2_vals * beta2 + (1.0 - beta2) * g_vals * g_vals correction1 = 1.0 - beta1**step correction2 = sqrt(1.0 - beta2**step) if weight_decay > 0.0: p_vals = p_vals * (1.0 - lr * weight_decay) mixed_momentum = (m1 / correction1) + (alpha * m2) adaptive_term = (torch.sqrt(nu) / correction2) + eps p_vals = p_vals - lr * (mixed_momentum / adaptive_term) state1[0].copy_(m1) state1[1].copy_(m2) state2.copy_(nu) elif optimizer_id == 0: # MOMENTUM if step == 1: s1_vals = g_vals else: s1_vals = state1 * beta1 + g_vals update_val = update_scale * (-lr * s1_vals) p_vals = p_vals + update_val state1.copy_(s1_vals) elif optimizer_id == 4: # LION momentum_update = state1 * beta1 + (1.0 - beta1) * g_vals update_val = update_scale * lr * torch.sign(momentum_update) p_vals = p_vals - update_val s1_vals = state1 * beta2 + (1.0 - beta2) * g_vals state1.copy_(s1_vals) elif optimizer_id == 1: # RMSPROP s1_vals = state1 * beta1 + (1.0 - beta1) * g_vals * g_vals update_val = update_scale * lr * g_vals / (torch.sqrt(s1_vals) + eps) p_vals = p_vals - update_val state1.copy_(s1_vals) elif optimizer_id == 2: # ADAGRAD s1_vals = state1 + g_vals * g_vals update_val = lr * g_vals / (torch.sqrt(s1_vals) + eps) p_vals = p_vals - update_val state1.copy_(s1_vals) p.copy_(p_vals) @register_kernel("bitsandbytes::optimizer_update_32bit", "default") def _( optimizer_name: str, g: torch.Tensor, p: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], unorm_vec: Optional[torch.Tensor], max_unorm: float, param_norm: float, beta1: float, beta2: float, beta3: float, alpha: float, eps: float, weight_decay: float, step: int, lr: float, gnorm_scale: float = 1.0, skip_zeros=False, ) -> None: """ 32-bit optimizer implemented by PyTorch with @torch.compile """ if skip_zeros: raise NotImplementedError("skip_zeros is not supported yet") optimizer_id = name2optimizer_id[optimizer_name] if optimizer_name == "lion": _optimizer_update_32bit( g, p, state1, state2, unorm_vec, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, lr, gnorm_scale, optimizer_id, ) if max_unorm > 0.0: unorm_vec.zero_() _optimizer_precondition_32bit( g, p, state1, state2, unorm_vec, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, optimizer_id ) else: if max_unorm > 0.0: unorm_vec.zero_() _optimizer_precondition_32bit( g, p, state1, state2, unorm_vec, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, optimizer_id ) _optimizer_update_32bit( g, p, state1, state2, unorm_vec, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, lr, gnorm_scale, optimizer_id, ) bitsandbytes/bitsandbytes/backends/hpu/000077500000000000000000000000001506752665600206345ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/hpu/__init__.py000066400000000000000000000000001506752665600227330ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/hpu/ops.py000066400000000000000000000026751506752665600220210ustar00rootroot00000000000000from collections.abc import Sequence import math import torch from ..._ops import register_kernel from ..utils import GAUDI_SW_VER # convert btw standard 4-bit compression format and ipex compression format # needed for backward compatibility with older versions of gaudi sw def _reverse_4bit_compress_format(weight: torch.Tensor): out_1 = (weight & 0xF0) >> 4 out_2 = (weight & 0xF) << 4 out = out_1 | out_2 return out @register_kernel("bitsandbytes::dequantize_4bit", "hpu") def _( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, shape: Sequence[int], dtype: torch.dtype, ) -> torch.Tensor: torch._check_is_size(blocksize) torch._check(quant_type == "nf4", lambda: f"quant_type must be nf4, got {quant_type}") torch._check( A.dtype in [torch.bfloat16, torch.uint8], lambda: f"quant_storage supports uint8 or bfloat16, but got {A.dtype}", ) # Enable non uint8 dtype if A.dtype != torch.uint8: A = A.view(torch.uint8) A = A.reshape(-1) if GAUDI_SW_VER and (GAUDI_SW_VER.major < 1 or GAUDI_SW_VER.minor < 22): A = _reverse_4bit_compress_format(A) # HPU dequantization function for NF4 quantized tensors. out_dq = torch.ops.hpu.dequantize_nf4( A, absmax.to(dtype), blocksize, out_shape=(math.prod(shape),), out_dtype=dtype, ) output = out_dq.reshape(shape) return output bitsandbytes/bitsandbytes/backends/triton/000077500000000000000000000000001506752665600213575ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/triton/__init__.py000066400000000000000000000000001506752665600234560ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/triton/kernels_4bit.py000066400000000000000000000522631506752665600243260ustar00rootroot00000000000000import torch import triton import triton.language as tl # Triton implementation of similar CUDA kernel to avoid loading code from csrc/kernels.cu::dQuantizeFP4 # @triton.autotune( # configs=[ # triton.Config({"SPLIT_NUM_BLOCKS": 1, "grf_mode": "auto"}, num_stages=4, num_warps=32), # triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=4, num_warps=32), # triton.Config({"SPLIT_NUM_BLOCKS": 1}), # triton.Config({"SPLIT_NUM_BLOCKS": 2}), # triton.Config({"SPLIT_NUM_BLOCKS": 4}), # triton.Config({"SPLIT_NUM_BLOCKS": 8}), # ], # key=["n_elements"], # ) @triton.jit def quantize_fp4_blockwise_kernel( A_ptr, absmax_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr, SPLIT_NUM_BLOCKS: tl.constexpr, ): PAIRED_SPLIT_NUM_BLOCKS: tl.constexpr = SPLIT_NUM_BLOCKS * 2 block_start_idx = tl.program_id(0) * PAIRED_SPLIT_NUM_BLOCKS thread_idx = tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS * BLOCK_SIZE) offsets = block_start_idx * BLOCK_SIZE + thread_idx mask = offsets < n_elements A = tl.load(A_ptr + offsets, mask=mask, other=0.0) # To be able process several blocks -> (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE) A_reshaped = tl.reshape(A, (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE)) # Calculating absamax for each block absmax = tl.max(tl.abs(A_reshaped), axis=1) tl.store(absmax_ptr + block_start_idx + tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS), absmax) A_normalized = A_reshaped / absmax[:, None] A_normalized = tl.clamp(A_normalized, -1.0, 1.0) sign = tl.where(A_normalized < 0, 0b1000, 0b0000) A_absf = tl.abs(A_normalized) result = tl.where( A_absf > 0.29166667, tl.where( A_absf > 0.583333, tl.where(A_absf > 0.8333333, 0b011, 0b010), tl.where(A_absf > 0.4166667, 0b101, 0b100) ), tl.where( A_absf > 0.0859375, tl.where(A_absf > 0.20833333, 0b0111, 0b0110), tl.where(A_absf > 0.00260417, 0b0001, 0b0000), ), ) quantized = (result ^ sign).to(tl.uint8) quantized = quantized.reshape((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE // 2, 2)) left, right = quantized.split() packed = left << 4 | (right & 0xF) packed_flat = tl.reshape(packed, (BLOCK_SIZE * SPLIT_NUM_BLOCKS,)) out_offsets = block_start_idx * BLOCK_SIZE // 2 + tl.arange(0, SPLIT_NUM_BLOCKS * BLOCK_SIZE) out_mask = out_offsets < n_elements // 2 tl.store(out_ptr + out_offsets, packed_flat, mask=out_mask) # Triton implementation of similar CUDA kernel to avoid loading code from csrc/kernels.cu::dQuantizeNF4 # @triton.autotune( # configs=[ # triton.Config({"SPLIT_NUM_BLOCKS": 1, "grf_mode": "auto"}, num_stages=4, num_warps=32), # triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=4, num_warps=32), # triton.Config({"SPLIT_NUM_BLOCKS": 1}), # triton.Config({"SPLIT_NUM_BLOCKS": 2}), # triton.Config({"SPLIT_NUM_BLOCKS": 4}), # triton.Config({"SPLIT_NUM_BLOCKS": 8}), # ], # key=["n_elements"], # ) @triton.jit def quantize_nf4_blockwise_kernel( A_ptr, absmax_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr, SPLIT_NUM_BLOCKS: tl.constexpr, ): PAIRED_SPLIT_NUM_BLOCKS: tl.constexpr = SPLIT_NUM_BLOCKS * 2 block_start_idx = tl.program_id(0) * PAIRED_SPLIT_NUM_BLOCKS thread_idx = tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS * BLOCK_SIZE) offsets = block_start_idx * BLOCK_SIZE + thread_idx mask = offsets < n_elements A = tl.load(A_ptr + offsets, mask=mask, other=0.0) # To be able process several blocks -> (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE) A_reshaped = tl.reshape(A, (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE)) # Calculating absamax for each block absmax = tl.max(tl.abs(A_reshaped), axis=1) tl.store(absmax_ptr + block_start_idx + tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS), absmax) A_normalized = A_reshaped / absmax[:, None] A_normalized = tl.clamp(A_normalized, -1.0, 1.0) result = tl.where( A_normalized > 0.03979014977812767, tl.where( A_normalized > 0.3893125355243683, tl.where( A_normalized > 0.6427869200706482, tl.where(A_normalized > 0.8614784181118011, 0b1111, 0b1110), tl.where(A_normalized > 0.5016634166240692, 0b1101, 0b1100), ), tl.where( A_normalized > 0.2035212516784668, tl.where(A_normalized > 0.2920137718319893, 0b1011, 0b1010), tl.where(A_normalized > 0.1202552504837513, 0b1001, 0b1000), ), ), tl.where( A_normalized > -0.33967943489551544, tl.where( A_normalized > -0.13791173323988914, tl.where(A_normalized > -0.045525018125772476, 0b0111, 0b0110), tl.where(A_normalized > -0.23460740596055984, 0b0101, 0b0100), ), tl.where( A_normalized > -0.6106329262256622, tl.where(A_normalized > -0.4599952697753906, 0b0011, 0b0010), tl.where(A_normalized > -0.8480964004993439, 0b0001, 0b0000), ), ), ) quantized = result.to(tl.uint8) quantized = quantized.reshape((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE // 2, 2)) left, right = quantized.split() packed = left << 4 | (right & 0xF) packed_flat = tl.reshape(packed, (BLOCK_SIZE * SPLIT_NUM_BLOCKS,)) out_offsets = block_start_idx * BLOCK_SIZE // 2 + tl.arange(0, SPLIT_NUM_BLOCKS * BLOCK_SIZE) out_mask = out_offsets < n_elements // 2 tl.store(out_ptr + out_offsets, packed_flat, mask=out_mask) def quantize_4bit_blockwise_triton(A, blocksize, quant_type, blocks, absmax, num_elements, quantized_out): # grid = lambda META: (triton.cdiv(blocks, META["SPLIT_NUM_BLOCKS"]),) split_num_blocks = 4 grid = (triton.cdiv(blocks, split_num_blocks),) if quant_type == "fp4": quantize_fp4_blockwise_kernel[grid]( A_ptr=A, absmax_ptr=absmax, out_ptr=quantized_out, n_elements=num_elements, BLOCK_SIZE=blocksize, SPLIT_NUM_BLOCKS=split_num_blocks, ) else: quantize_nf4_blockwise_kernel[grid]( A_ptr=A, absmax_ptr=absmax, out_ptr=quantized_out, n_elements=num_elements, BLOCK_SIZE=blocksize, SPLIT_NUM_BLOCKS=split_num_blocks, ) return quantized_out, absmax @triton.jit def dequant_4bit_body_util(a, offsets, quant_ptr, absmax_ptr, n_elems, QUANT_BLOCK: tl.constexpr): PAIRED_QUANT_BLOCK: tl.constexpr = QUANT_BLOCK // 2 mask = offsets < n_elems higher = a & 0xF # lower 4bits lower = a >> 4 abs_offsets = offsets // PAIRED_QUANT_BLOCK absmax = tl.load(absmax_ptr + abs_offsets, mask=mask, other=1.0, eviction_policy="evict_last") # apply conversion lower_4 = tl.load(quant_ptr + lower, eviction_policy="evict_last") higher_4 = tl.load(quant_ptr + higher, eviction_policy="evict_last") mul_high = higher_4 * absmax mul_low = lower_4 * absmax out_dq = tl.interleave(mul_low, mul_high) return out_dq # Triton implementation of similar CUDA kernel to avoid loading code from csrc/kernels.cu::dDequantizeFP4Tree @triton.jit def dequantize_fp4_tree(val, absmax): # val: tl.tensor (uint8) # absmax: tl.tensor (float32/float16) # 00001100 00001011 00001001 00001111 sign = tl.where((val & 0b1000) == 0b1000, -1.0, 1.0) # -1 third_bit = (val & 0b0100) == 0b0100 # True second_bit = (val & 0b0010) == 0b0010 # False first_bit = (val & 0b0001) == 0b0001 # False branch1 = tl.where( second_bit, tl.where(first_bit, 0.25, 0.16666667), # 1111, 1110 tl.where(first_bit, 0.5, 0.33333333), # 1101, 1100 ) branch2 = tl.where( second_bit, tl.where(first_bit, 1.0, 0.66666667), # 1011, 1010 tl.where(first_bit, 0.00520833, 0.0), # 1001, 1000 ) out = tl.where(third_bit, branch1, branch2) return out * sign * absmax @triton.jit def dequant_fp4_body_util(a, offsets, absmax_ptr, n_elems, QUANT_BLOCK: tl.constexpr): PAIRED_QUANT_BLOCK: tl.constexpr = QUANT_BLOCK // 2 mask = offsets < n_elems higher = a & 0xF lower = a >> 4 abs_offsets = offsets // PAIRED_QUANT_BLOCK absmax = tl.load(absmax_ptr + abs_offsets, mask=mask, other=1.0, eviction_policy="evict_last") mul_high = dequantize_fp4_tree(higher, absmax) mul_low = dequantize_fp4_tree(lower, absmax) out_dq = tl.interleave(mul_low, mul_high) return out_dq # Triton implementation of similar CUDA kernel to avoid loading code from csrc/kernels.cu::dDequantizeNF4 @triton.jit def dequantize_nf4_tree(val): # val: tl.tensor (uint8) cond0 = (val & 0b1000) == 0b1000 cond1 = (val & 0b0100) == 0b0100 cond2 = (val & 0b0010) == 0b0010 cond3 = (val & 0b0001) == 0b0001 # Positive branch (val & 0b1000) == 8 branch_pos = tl.where( cond1, tl.where( cond2, tl.where(cond3, 1.0, 0.7229568362236023), # 1111, 1110 tl.where(cond3, 0.5626170039176941, 0.44070982933044434), # 1101, 1100 ), tl.where( cond2, tl.where(cond3, 0.33791524171829224, 0.24611230194568634), # 1011, 1010 tl.where(cond3, 0.16093020141124725, 0.07958029955625534), # 1001, 1000 ), ) # Negative branch (val & 0b1000) == 0 branch_neg = tl.where( cond1, tl.where( cond2, tl.where(cond3, 0.0, -0.09105003625154495), # 0111, 0110 tl.where(cond3, -0.18477343022823334, -0.28444138169288635), # 0101, 0100 ), tl.where( cond2, tl.where(cond3, -0.39491748809814453, -0.5250730514526367), # 0011, 0010 tl.where(cond3, -0.6961928009986877, -1.0), # 0001, 0000 ), ) return tl.where(cond0, branch_pos, branch_neg) @triton.jit def dequant_nf4_body_util(a, offsets, absmax_ptr, n_elems, QUANT_BLOCK: tl.constexpr): PAIRED_QUANT_BLOCK: tl.constexpr = QUANT_BLOCK // 2 mask = offsets < n_elems higher = a & 0xF # lower 4bits lower = a >> 4 abs_offsets = offsets // PAIRED_QUANT_BLOCK absmax = tl.load(absmax_ptr + abs_offsets, mask=mask, other=1.0, eviction_policy="evict_last") mul_high = dequantize_nf4_tree(higher) * absmax mul_low = dequantize_nf4_tree(lower) * absmax out_dq = tl.interleave(mul_low, mul_high) return out_dq # All such kernels are similar, so maybe code can be generalised. # @triton.autotune( # configs=[ # # # triton.Config({'SPLIT_SIZE': 64}), # # # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'large'}, num_stages=2, num_warps=32), # # # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'auto'}, num_stages=2, num_warps=32), # # # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'large'}, num_stages=4, num_warps=32), # # # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'auto'}, num_stages=4, num_warps=32), # triton.Config({'SPLIT_SIZE': 128}), # triton.Config({'SPLIT_SIZE': 128}, num_warps = 32, num_stages = 2), # # # triton.Config({'SPLIT_SIZE': 128}, num_warps = 4, num_stages = 4), # # # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'large'}, num_stages=2, num_warps=32), # # # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'auto'}, num_stages=2, num_warps=32), # # # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'large'}, num_stages=4, num_warps=32), # # # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'auto'}, num_stages=4, num_warps=32), # triton.Config({'SPLIT_SIZE': 256}), # triton.Config({'SPLIT_SIZE': 256}, num_warps = 32, num_stages = 2), # # triton.Config({'SPLIT_SIZE': 256}, num_warps = 4, num_stages = 4), # triton.Config({'SPLIT_SIZE': 512}), # triton.Config({'SPLIT_SIZE': 512}, num_warps = 32, num_stages = 2), # # triton.Config({'SPLIT_SIZE': 512}, num_warps = 4, num_stages = 4), # # # # triton.Config({'SPLIT_SIZE': 512, 'grf_mode': 'large'}, num_stages=2, num_warps=32), # # # # triton.Config({'SPLIT_SIZE': 512, 'grf_mode': 'auto'}, num_stages=2, num_warps=32), # # # # triton.Config({'SPLIT_SIZE': 512, 'grf_mode': 'large'}, num_stages=4, num_warps=32), # # # # triton.Config({'SPLIT_SIZE': 512, 'grf_mode': 'auto'}, num_stages=4, num_warps=32), # # # triton.Config({'SPLIT_SIZE': 1024}), # # # # triton.Config({'SPLIT_SIZE': 2048}), # # # # triton.Config({'SPLIT_SIZE': 4096}), # # # # triton.Config({'SPLIT_SIZE': 8192}), # # # # triton.Config({'SPLIT_SIZE': 16384}), # ], # key=['num_paired_elements'], # ) @triton.jit def dequant_4bit_kernel( a_ptr, c_ptr, quant_ptr, absmax_ptr, num_paired_elements, QUANT_BLOCK: tl.constexpr, SPLIT_SIZE: tl.constexpr ): pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. block_start = pid * SPLIT_SIZE offsets = block_start + tl.arange(0, SPLIT_SIZE) mask = offsets < num_paired_elements a = tl.load(a_ptr + offsets, mask, eviction_policy="evict_first") out_dq = dequant_4bit_body_util( a=a, offsets=offsets, quant_ptr=quant_ptr, absmax_ptr=absmax_ptr, n_elems=num_paired_elements, QUANT_BLOCK=QUANT_BLOCK, ) out_block_start = pid * SPLIT_SIZE * 2 offs = out_block_start + tl.arange(0, SPLIT_SIZE * 2) mask = offs < num_paired_elements * 2 tl.store(c_ptr + offs, out_dq, mask) # @triton.autotune( # configs=[ # triton.Config({'SPLIT_SIZE': 128}, num_warps = 32, num_stages = 2), # triton.Config({'SPLIT_SIZE': 256}), # triton.Config({'SPLIT_SIZE': 256}, num_warps = 32, num_stages = 2), # triton.Config({'SPLIT_SIZE': 512}), # triton.Config({'SPLIT_SIZE': 512}, num_warps = 32, num_stages = 2), # triton.Config({'SPLIT_SIZE': 1024}, num_warps = 32, num_stages = 2), # ], # key=['num_paired_elements'], # ) @triton.jit def dequant_fp4_kernel( a_ptr, c_ptr, absmax_ptr, num_paired_elements, QUANT_BLOCK: tl.constexpr, SPLIT_SIZE: tl.constexpr ): pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. block_start = pid * SPLIT_SIZE offsets = block_start + tl.arange(0, SPLIT_SIZE) mask = offsets < num_paired_elements a = tl.load(a_ptr + offsets, mask, eviction_policy="evict_first") out_dq = dequant_fp4_body_util( a=a, offsets=offsets, absmax_ptr=absmax_ptr, n_elems=num_paired_elements, QUANT_BLOCK=QUANT_BLOCK, ) out_block_start = pid * SPLIT_SIZE * 2 offs = out_block_start + tl.arange(0, SPLIT_SIZE * 2) mask = offs < num_paired_elements * 2 tl.store(c_ptr + offs, out_dq, mask) # @triton.autotune( # configs=[ # triton.Config({'SPLIT_SIZE': 128}, num_warps = 32, num_stages = 2), # triton.Config({'SPLIT_SIZE': 256}), # triton.Config({'SPLIT_SIZE': 256}, num_warps = 32, num_stages = 2), # triton.Config({'SPLIT_SIZE': 512}), # triton.Config({'SPLIT_SIZE': 512}, num_warps = 32, num_stages = 2), # triton.Config({'SPLIT_SIZE': 1024}, num_warps = 32, num_stages = 2), # ], # key=['num_paired_elements'], # ) @triton.jit def dequant_nf4_kernel( a_ptr, c_ptr, absmax_ptr, num_paired_elements, QUANT_BLOCK: tl.constexpr, SPLIT_SIZE: tl.constexpr ): pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. block_start = pid * SPLIT_SIZE offsets = block_start + tl.arange(0, SPLIT_SIZE) mask = offsets < num_paired_elements a = tl.load(a_ptr + offsets, mask, eviction_policy="evict_first") out_dq = dequant_nf4_body_util( a=a, offsets=offsets, absmax_ptr=absmax_ptr, n_elems=num_paired_elements, QUANT_BLOCK=QUANT_BLOCK, ) out_block_start = pid * SPLIT_SIZE * 2 offs = out_block_start + tl.arange(0, SPLIT_SIZE * 2) mask = offs < num_paired_elements * 2 tl.store(c_ptr + offs, out_dq, mask) def dequantize_4bit_impl( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, dtype: torch.dtype, out: torch.Tensor, ) -> None: # It's will be processed as an array, so # actual length is row * col # Elements are in uint8 format, so interleaved # so total amount of data is 2 * elem_count number_of_paired_elements = A.numel() # we assume that split_size > quant_blocksize SPLIT_SIZE = 256 # grid = lambda META: (triton.cdiv(number_of_paired_elements, META['SPLIT_SIZE']), ) grid = (triton.cdiv(number_of_paired_elements, SPLIT_SIZE),) if quant_type == "fp4": dequant_fp4_kernel[grid](A, out, absmax, number_of_paired_elements, blocksize, SPLIT_SIZE) else: dequant_nf4_kernel[grid](A, out, absmax, number_of_paired_elements, blocksize, SPLIT_SIZE) def dequantize_4bit_impl_passing_code( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, code: torch.Tensor, dtype: torch.dtype, out: torch.Tensor, ) -> None: number_of_paired_elements = A.numel() # we assume that split_size > quant_blocksize SPLIT_SIZE = 256 # grid = lambda META: (triton.cdiv(number_of_paired_elements, META['SPLIT_SIZE']), ) grid = (triton.cdiv(number_of_paired_elements, SPLIT_SIZE),) dequant_4bit_kernel[grid](A, out, code, absmax, number_of_paired_elements, blocksize, SPLIT_SIZE) ######################### Fallback dequantization functions ######################### ## for debug ## # @triton.autotune( # configs=[ # # triton.Config({'SPLIT_NUM_BLOCKS': 1, 'grf_mode': 'large'}, num_stages=2, num_warps=32), # # triton.Config({'SPLIT_NUM_BLOCKS': 1, 'grf_mode': 'auto'}, num_stages=2, num_warps=32), # # triton.Config({'SPLIT_NUM_BLOCKS': 1, 'grf_mode': 'large'}, num_stages=4, num_warps=32), # # # # # triton.Config({"SPLIT_NUM_BLOCKS": 1, "grf_mode": "auto"}, num_stages=4, num_warps=32), # # # triton.Config({"SPLIT_NUM_BLOCKS": 2}), # # triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "large"}, num_stages=2, num_warps=32), # # # triton.Config({'SPLIT_NUM_BLOCKS': 2, 'grf_mode': 'large'}, num_stages=4, num_warps=32), # # triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=2, num_warps=32), # # triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=4, num_warps=32), # # triton.Config({"SPLIT_NUM_BLOCKS": 4, "grf_mode": "large"}, num_stages=2, num_warps=32), # # triton.Config({"SPLIT_NUM_BLOCKS": 4, "grf_mode": "large"}, num_stages=4, num_warps=32), # # triton.Config({'SPLIT_NUM_BLOCKS': 8, 'grf_mode': 'large'}, num_stages=2, num_warps=32), # ], # key=["n_elements", "BLOCK_SIZE"], # ) @triton.jit def quantize_4bit_blockwise_kernel( A_ptr, code_ptr, absmax_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr, CODE_SIZE: tl.constexpr, SPLIT_NUM_BLOCKS: tl.constexpr, ): PAIRED_SPLIT_NUM_BLOCKS: tl.constexpr = SPLIT_NUM_BLOCKS * 2 block_start_idx = tl.program_id(0) * PAIRED_SPLIT_NUM_BLOCKS thread_idx = tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS * BLOCK_SIZE) offsets = block_start_idx * BLOCK_SIZE + thread_idx mask = offsets < n_elements A = tl.load(A_ptr + offsets, mask=mask, other=0.0) # To be able process several blocks -> (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE) A_reshaped = tl.reshape(A, (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE)) # Calculating absamax for each block absmax = tl.max(tl.abs(A_reshaped), axis=1) tl.store(absmax_ptr + block_start_idx + tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS), absmax) A_normalized = A_reshaped / absmax[:, None] A_normalized = tl.clamp(A_normalized, -1.0, 1.0) lower_pivot = tl.zeros((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE), dtype=tl.int32) upper_pivot = tl.full((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE), CODE_SIZE - 1, dtype=tl.int32) for _ in range(4): # ceil(log2(code_size)) = 4, actually, in general case should be input parameter pivot = (lower_pivot + upper_pivot) // 2 val = tl.load(code_ptr + pivot) is_higher = A_normalized > val # code[pivot] lower_pivot = tl.where(is_higher, pivot, lower_pivot) upper_pivot = tl.where(is_higher, upper_pivot, pivot) # Choose closest level lower_val = tl.load(code_ptr + lower_pivot) upper_val = tl.load(code_ptr + upper_pivot) lower_dist = tl.abs(A_normalized - lower_val) upper_dist = tl.abs(A_normalized - upper_val) quantized = tl.where(lower_dist <= upper_dist, lower_pivot, upper_pivot).to(tl.uint8) quantized = quantized.reshape((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE // 2, 2)) quantized = quantized.to(tl.uint8, bitcast=True) left, right = quantized.split() packed = left << 4 | (right & 0xF) # Reduce don't guarantee the order of the elements passed to unite_2_int4 # packed = tl.reduce(quantized, axis=2, combine_fn=unite_2_int4) # packed = packed.to(tl.uint8, bitcast=True) packed_flat = tl.reshape(packed, (BLOCK_SIZE * SPLIT_NUM_BLOCKS,)) out_offsets = block_start_idx * BLOCK_SIZE // 2 + tl.arange(0, SPLIT_NUM_BLOCKS * BLOCK_SIZE) out_mask = out_offsets < n_elements // 2 tl.store(out_ptr + out_offsets, packed_flat, mask=out_mask) bitsandbytes/bitsandbytes/backends/triton/kernels_8bit_quant.py000066400000000000000000000146161506752665600255420ustar00rootroot00000000000000import torch import triton import triton.language as tl # @triton.autotune( # configs=[ # # triton.Config({'SPLIT_SIZE': 64}), # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'large'}, num_stages=2, num_warps=32), # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'auto'}, num_stages=2, num_warps=32), # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'large'}, num_stages=4, num_warps=32), # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'auto'}, num_stages=4, num_warps=32), # # triton.Config({'SPLIT_SIZE': 128}), # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'large'}, num_stages=2, num_warps=32), # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'auto'}, num_stages=2, num_warps=32), # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'large'}, num_stages=4, num_warps=32), # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'auto'}, num_stages=4, num_warps=32), # triton.Config({"SPLIT_SIZE": 256}), # # triton.Config({'SPLIT_SIZE': 256, 'grf_mode': 'large'}, num_stages=2, num_warps=32), # # triton.Config({'SPLIT_SIZE': 256, 'grf_mode': 'auto'}, num_stages=2, num_warps=32), # triton.Config({"SPLIT_SIZE": 512}), # # triton.Config({'SPLIT_SIZE': 1024}), # ], # key=["num_paired_elements", "QUANT_BLOCK"], # ) @triton.jit def dequant_8bit_kernel( a_ptr, out_ptr, code_ptr, absmax_ptr, n, QUANT_BLOCK: tl.constexpr, SPLIT_SIZE: tl.constexpr, ): pid = tl.program_id(axis=0) block_start = pid * SPLIT_SIZE offsets = block_start + tl.arange(0, SPLIT_SIZE) mask = offsets < n out_dq = dequant_8bit_blockwise_kernel_util(a_ptr, offsets, code_ptr, absmax_ptr, mask, QUANT_BLOCK) tl.store(out_ptr + offsets, out_dq, mask) def dequant_8bit_blockwise( a: torch.Tensor, absmax: torch.Tensor, quant_state_code: torch.Tensor, quant_blocksize: int = 64, dtype: torch.dtype = None, out: torch.Tensor = None, ): n = a.numel() if out is None: if dtype is None: raise ValueError("If out is None, dtype must be specified") out = torch.empty_like(a, dtype=dtype, device=a.device) SPLIT_SIZE = 256 # grid = lambda META: (triton.cdiv(number_of_paired_elements, META["SPLIT_SIZE"]),) grid = (triton.cdiv(n, SPLIT_SIZE),) dequant_8bit_kernel[grid]( a, out, quant_state_code, absmax, n, quant_blocksize, SPLIT_SIZE, ) return out # @triton.autotune( # configs=[ # triton.Config({"SPLIT_NUM_BLOCKS": 1, "grf_mode": "auto"}, num_stages=4, num_warps=32), # triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=4, num_warps=32), # triton.Config({"SPLIT_NUM_BLOCKS": 1}), # triton.Config({"SPLIT_NUM_BLOCKS": 2}), # ], # key=["n_elements"], # ) @triton.jit def quantize_8bit_blockwise_kernel( A_ptr, code_ptr, absmax_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr, CODE_SIZE: tl.constexpr, SPLIT_NUM_BLOCKS: tl.constexpr, ): block_start_idx = tl.program_id(0) * SPLIT_NUM_BLOCKS thread_idx = tl.arange(0, SPLIT_NUM_BLOCKS * BLOCK_SIZE) offsets = block_start_idx * BLOCK_SIZE + thread_idx mask = offsets < n_elements A = tl.load(A_ptr + offsets, mask=mask, other=0.0) quantized, absmax = quantize_8bit_blockwise_kernel_util(A, code_ptr, CODE_SIZE, BLOCK_SIZE, SPLIT_NUM_BLOCKS) tl.store(absmax_ptr + block_start_idx + tl.arange(0, SPLIT_NUM_BLOCKS), absmax) tl.store(out_ptr + offsets, quantized, mask=mask) def quantize_blockwise_triton(A, code, blocksize, absmax=None, out=None): n = A.numel() blocks = -(n // -blocksize) if absmax is None: absmax = torch.empty((blocks,), device=A.device, dtype=A.dtype) if out is None: out = torch.empty_like(A.flatten(), dtype=torch.uint8) split_num_blocks = 1 grid = (triton.cdiv(blocks, split_num_blocks),) # grid = lambda META: (triton.cdiv(blocks, META["SPLIT_NUM_BLOCKS"]),) quantize_8bit_blockwise_kernel[grid]( A_ptr=A, code_ptr=code, absmax_ptr=absmax, out_ptr=out, n_elements=n, BLOCK_SIZE=blocksize, CODE_SIZE=code.numel(), SPLIT_NUM_BLOCKS=split_num_blocks, # num_warps=1, # num_stages=2, ) out = out.reshape(A.shape) return out, absmax @triton.jit def quantize_8bit_blockwise_kernel_util( a, code_ptr, CODE_SIZE: tl.constexpr, BLOCK_SIZE: tl.constexpr, N_PER_TH: tl.constexpr, ): # To be able process several blocks -> (BLOCK_SIZE, SPLIT_NUM_BLOCKS) a_reshaped = tl.reshape(a, (N_PER_TH, BLOCK_SIZE)) # Calculating absmax for each block absmax = tl.max(tl.abs(a_reshaped), axis=1) a_normalized = a_reshaped / absmax[:, None] a_normalized = tl.clamp(a_normalized, -1.0, 1.0) lower_pivot = tl.zeros((N_PER_TH, BLOCK_SIZE), dtype=tl.int32) upper_pivot = tl.full((N_PER_TH, BLOCK_SIZE), CODE_SIZE - 1, dtype=tl.int32) # ceil(log2(code_size)) = 8, actually, in general case should be input parameter for _ in range(8): pivot = (lower_pivot + upper_pivot) // 2 val = tl.load(code_ptr + pivot) is_higher = a_normalized > val # code[pivot] lower_pivot = tl.where(is_higher, pivot, lower_pivot) upper_pivot = tl.where(is_higher, upper_pivot, pivot) # Choose closest level lower_val = tl.load(code_ptr + lower_pivot) upper_val = tl.load(code_ptr + upper_pivot) lower_dist = tl.abs(a_normalized - lower_val) upper_dist = tl.abs(a_normalized - upper_val) quantized = tl.where(lower_dist <= upper_dist, lower_pivot, upper_pivot).to(tl.uint8) # too slow approach # diff = tl.abs(A_normalized[:, :, None] - code[None, None, :]) # quantized = tl.argmin(diff, axis=2).to(tl.uint8) quantized_flat = tl.reshape(quantized, (BLOCK_SIZE * N_PER_TH,)) return quantized_flat, absmax @triton.jit def dequant_8bit_blockwise_kernel_util( a_ptr, offsets, code_ptr, absmax_ptr, mask, BLOCK_SIZE: tl.constexpr, ): a = tl.load(a_ptr + offsets, mask, other=0).to(tl.uint8) scaled_int8 = tl.load(code_ptr + a, mask) # Load scales absmax_offsets = offsets // BLOCK_SIZE absmax = tl.load(absmax_ptr + absmax_offsets, mask=mask, other=0.0, eviction_policy="evict_last") # Apply scales out_dq = scaled_int8 * absmax return out_dq bitsandbytes/bitsandbytes/backends/triton/kernels_optim.py000066400000000000000000001110771506752665600246130ustar00rootroot00000000000000import math from typing import Optional import torch import triton import triton.language as tl # from triton.language.extra import libdevice from .kernels_8bit_quant import ( dequant_8bit_blockwise, dequant_8bit_blockwise_kernel_util, quantize_8bit_blockwise_kernel_util, quantize_blockwise_triton, ) MOMENTUM = 0 RMSPROP = 1 ADAGRAD = 2 ADAM = 3 # LION should be larger than MOMENTUM, RMSPROP, ADAGRAD due to comparison in kernels LION = 4 ADEMAMIX = 5 name2optimizer_id = { "momentum": MOMENTUM, "rmsprop": RMSPROP, "adagrad": ADAGRAD, "adam": ADAM, "lion": LION, "ademamix": ADEMAMIX, } @triton.jit def _optimizer_precondition_2state_32bit( g_ptr, p_ptr, state1_ptr, state2_ptr, unorm_ptr, beta1: tl.constexpr, beta2: tl.constexpr, eps: tl.constexpr, weight_decay: tl.constexpr, step, beta1_step, beta2_step, lr, gnorm_scale: tl.constexpr, n_elements, OPTIMIZER_ID: tl.constexpr, BLOCK_SIZE: tl.constexpr, N_PER_TH: tl.constexpr, ): """Preprocessing optimizer, computing update norm (2-state optimizer)""" pid = tl.program_id(axis=0) block_start_idx = pid * N_PER_TH offsets = block_start_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE * N_PER_TH) mask = offsets < n_elements g_vals = tl.load(g_ptr + offsets, mask=mask, other=0.0) s1_vals = tl.load(state1_ptr + offsets, mask=mask, other=0.0) s2_vals = tl.load(state2_ptr + offsets, mask=mask, other=0.0) g_vals = gnorm_scale * g_vals correction1 = 1.0 / (1.0 - beta1_step) correction2 = 1.0 / (1.0 - beta2_step) if OPTIMIZER_ID == 3: # ADAM s1_vals = s1_vals * beta1 + (1.0 - beta1) * g_vals s2_vals = s2_vals * beta2 + (1.0 - beta2) * g_vals * g_vals s1_vals = s1_vals * correction1 s2_vals = s2_vals * correction2 update_vals = s1_vals / (tl.sqrt(s2_vals) + eps) update_norm = update_vals * update_vals elif OPTIMIZER_ID == 5: # ADEMAMIX update_norm = s1_vals total_norm = tl.sum(tl.where(mask, update_norm, 0.0)) tl.atomic_add(unorm_ptr, total_norm) @triton.jit def _optimizer_precondition_1state_32bit( g_ptr, p_ptr, state1_ptr, state2_ptr, unorm_ptr, beta1: tl.constexpr, beta2: tl.constexpr, eps: tl.constexpr, weight_decay, step, beta1_step, beta2_step, lr, gnorm_scale: tl.constexpr, n_elements, OPTIMIZER_ID: tl.constexpr, BLOCK_SIZE: tl.constexpr, N_PER_TH: tl.constexpr, ): """Preprocessing optimizer, computing update norm (1-state optimizer)""" pid = tl.program_id(axis=0) block_start_idx = pid * N_PER_TH offsets = block_start_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE * N_PER_TH) mask = offsets < n_elements g_vals = tl.load(g_ptr + offsets, mask=mask, other=0.0) s1_vals = tl.load(state1_ptr + offsets, mask=mask, other=0.0) g_vals = gnorm_scale * g_vals if OPTIMIZER_ID == 0: # MOMENTUM if step == 1: s1_vals = g_vals else: s1_vals = s1_vals * beta1 + g_vals update_norm = s1_vals * s1_vals elif OPTIMIZER_ID == 4: # LION s1_vals = s1_vals * beta2 + (1.0 - beta2) * g_vals update_norm = s1_vals elif OPTIMIZER_ID == 1: # RMSPROP s1_vals = s1_vals * beta1 + (1.0 - beta1) * g_vals * g_vals update_vals = g_vals / (tl.sqrt(s1_vals) + eps) update_norm = update_vals * update_vals elif OPTIMIZER_ID == 2: # ADAGRAD s1_vals = s1_vals + g_vals * g_vals update_vals = g_vals / (tl.sqrt(s1_vals) + eps) update_norm = update_vals * update_vals total_norm = tl.sum(tl.where(mask, update_norm, 0.0)) tl.atomic_add(unorm_ptr, total_norm) @triton.jit def _optimizer_update_2state_32bit_triton_kernel( g_ptr, p_ptr, state1_ptr, state2_ptr, unorm_ptr, max_unorm: tl.constexpr, param_norm, beta1: tl.constexpr, beta2: tl.constexpr, beta3, alpha, eps: tl.constexpr, weight_decay: tl.constexpr, step, beta1_step, beta2_step, lr, gnorm_scale: tl.constexpr, skip_zeros, n_elements, OPTIMIZER_ID: tl.constexpr, BLOCK_SIZE: tl.constexpr, N_PER_TH: tl.constexpr, ): """2-state optimizer kernel""" pid = tl.program_id(axis=0) block_start_idx = pid * N_PER_TH offsets = block_start_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE * N_PER_TH) mask = offsets < n_elements g_vals = tl.load(g_ptr + offsets, mask=mask, other=0.0).to(tl.float32) p_vals = tl.load(p_ptr + offsets, mask=mask, other=0.0).to(tl.float32) s1_vals = tl.load(state1_ptr + offsets, mask=mask, other=0.0) s2_vals = tl.load(state2_ptr + offsets, mask=mask, other=0.0) if OPTIMIZER_ID == 5: # ADEMAMIX s3_vals = tl.load(state1_ptr + n_elements + offsets, mask=mask, other=0.0) g_vals = gnorm_scale * g_vals update_scale = 1.0 if max_unorm > 0.0: current_unorm = tl.sqrt(tl.load(unorm_ptr)) if current_unorm > max_unorm * param_norm: update_scale = (max_unorm * param_norm) / current_unorm if OPTIMIZER_ID == 3: # ADAM s1_vals = s1_vals * beta1 + (1.0 - beta1) * g_vals s2_vals = s2_vals * beta2 + (1.0 - beta2) * g_vals * g_vals correction1 = 1.0 - beta1_step correction2 = tl.sqrt(1.0 - beta2_step) step_size = -lr * correction2 / correction1 if weight_decay > 0.0: p_vals = p_vals * (1.0 - lr * weight_decay) update_val = update_scale * step_size * (s1_vals / (tl.sqrt(s2_vals) + eps * correction2)) p_vals = p_vals + update_val elif OPTIMIZER_ID == 5: # ADEMAMIX s1_vals = s1_vals * beta1 + (1.0 - beta1) * g_vals # m1 s3_vals = s3_vals * beta3 + (1.0 - beta3) * g_vals # m2 s2_vals = s2_vals * beta2 + (1.0 - beta2) * g_vals * g_vals # nu correction1 = 1.0 - beta1_step correction2 = tl.sqrt(1.0 - beta2_step) if weight_decay > 0.0: p_vals = p_vals * (1.0 - lr * weight_decay) mixed_momentum = (s1_vals / correction1) + (alpha * s3_vals) adaptive_term = (tl.sqrt(s2_vals) / correction2) + eps p_vals = p_vals - lr * (mixed_momentum / adaptive_term) tl.store(p_ptr + offsets, p_vals, mask=mask) tl.store(state1_ptr + offsets, s1_vals, mask=mask) tl.store(state2_ptr + offsets, s2_vals, mask=mask) if OPTIMIZER_ID == 5: # ADEMAMIX tl.store(state1_ptr + n_elements + offsets, s3_vals, mask=mask) @triton.jit def _optimizer_update_1state_32bit_triton_kernel( g_ptr, p_ptr, state1_ptr, state2_ptr, unorm_ptr, max_unorm: tl.constexpr, param_norm, beta1: tl.constexpr, beta2: tl.constexpr, beta3, alpha, eps: tl.constexpr, weight_decay: tl.constexpr, step, beta1_step, beta2_step, lr, gnorm_scale: tl.constexpr, skip_zeros, n_elements, OPTIMIZER_ID: tl.constexpr, BLOCK_SIZE: tl.constexpr, N_PER_TH: tl.constexpr, ): """1-state optimizer kernel""" pid = tl.program_id(axis=0) block_start_idx = pid * N_PER_TH offsets = block_start_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE * N_PER_TH) mask = offsets < n_elements g_vals = tl.load(g_ptr + offsets, mask=mask, other=0.0).to(tl.float32) p_vals = tl.load(p_ptr + offsets, mask=mask, other=0.0).to(tl.float32) s1_vals = tl.load(state1_ptr + offsets, mask=mask, other=0.0) g_vals = gnorm_scale * g_vals if weight_decay > 0.0: g_vals = g_vals + p_vals * weight_decay update_scale = 1.0 if max_unorm > 0.0: current_unorm = tl.sqrt(tl.load(unorm_ptr)) if current_unorm > max_unorm * param_norm + eps: update_scale = (max_unorm * param_norm + eps) / current_unorm if OPTIMIZER_ID == 0: # MOMENTUM if step == 1: s1_vals = g_vals else: s1_vals = s1_vals * beta1 + g_vals update_val = update_scale * (-lr * s1_vals) p_vals = p_vals + update_val elif OPTIMIZER_ID == 4: # LION momentum_update = s1_vals * beta1 + (1.0 - beta1) * g_vals update_val = update_scale * lr * tl.where(momentum_update > 0, 1.0, tl.where(momentum_update < 0, -1.0, 0.0)) p_vals = p_vals - update_val s1_vals = s1_vals * beta2 + (1.0 - beta2) * g_vals elif OPTIMIZER_ID == 1: # RMSPROP s1_vals = s1_vals * beta1 + (1.0 - beta1) * g_vals * g_vals update_val = update_scale * lr * g_vals / (tl.sqrt(s1_vals) + eps) p_vals = p_vals - update_val elif OPTIMIZER_ID == 2: # ADAGRAD s1_vals = s1_vals + g_vals * g_vals update_val = lr * g_vals / (tl.sqrt(s1_vals) + eps) p_vals = p_vals - update_val tl.store(p_ptr + offsets, p_vals, mask=mask) tl.store(state1_ptr + offsets, s1_vals, mask=mask) name2optimizer_32bit_fn = { "adam": { "preprocess": _optimizer_precondition_2state_32bit, "update": _optimizer_update_2state_32bit_triton_kernel, }, "ademamix": { "preprocess": _optimizer_precondition_2state_32bit, "update": _optimizer_update_2state_32bit_triton_kernel, }, "momentum": { "preprocess": _optimizer_precondition_1state_32bit, "update": _optimizer_update_1state_32bit_triton_kernel, }, "rmsprop": { "preprocess": _optimizer_precondition_1state_32bit, "update": _optimizer_update_1state_32bit_triton_kernel, }, "adagrad": { "preprocess": _optimizer_precondition_1state_32bit, "update": _optimizer_update_1state_32bit_triton_kernel, }, "lion": { "preprocess": _optimizer_precondition_1state_32bit, "update": _optimizer_update_1state_32bit_triton_kernel, }, } def optimizer_update_32bit_impl( optimizer_name: str, g: torch.Tensor, p: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], unorm_vec: Optional[torch.Tensor], max_unorm: float, param_norm: float, beta1: float, beta2: float, beta3: float, alpha: float, eps: float, weight_decay: float, step: int, lr: float, gnorm_scale: float = 1.0, skip_zeros=False, ) -> None: """ 32-bit optimizer implemented by Triton """ if skip_zeros: raise NotImplementedError("skip_zeros is not supported on XPU yet") BLOCK_SIZE = 256 N_PER_TH = 1 # Number of blocks processed per thread. grid = (triton.cdiv(p.numel(), BLOCK_SIZE * N_PER_TH),) optimizer_id = name2optimizer_id[optimizer_name] fn_preprocess = name2optimizer_32bit_fn[optimizer_name]["preprocess"] fn_update = name2optimizer_32bit_fn[optimizer_name]["update"] # In torch=2.7 on XPU there is an issue with libdevice.pow, leading to an error. # For backwards compatibility we precompute the bias correction factors. beta1_step = beta1**step beta2_step = beta2**step if optimizer_name == "lion": fn_update[grid]( g, p, state1, state2, unorm_vec, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, beta1_step, beta2_step, lr, gnorm_scale, skip_zeros, p.numel(), optimizer_id, BLOCK_SIZE, N_PER_TH, num_warps=2, ) if max_unorm > 0.0: unorm_vec.zero_() fn_preprocess[grid]( g, p, state1, state2, unorm_vec, beta1, beta2, eps, weight_decay, step, beta1_step, beta2_step, lr, gnorm_scale, p.numel(), optimizer_id, BLOCK_SIZE, N_PER_TH, num_warps=2, ) else: if max_unorm > 0.0: unorm_vec.zero_() fn_preprocess[grid]( g, p, state1, state2, unorm_vec, beta1, beta2, eps, weight_decay, step, beta1_step, beta2_step, lr, gnorm_scale, p.numel(), optimizer_id, BLOCK_SIZE, N_PER_TH, num_warps=2, ) fn_update[grid]( g, p, state1, state2, unorm_vec, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, beta1_step, beta2_step, lr, gnorm_scale, skip_zeros, p.numel(), optimizer_id, BLOCK_SIZE, N_PER_TH, num_warps=2, ) ########################################### # Pure torch implementation for reference # ########################################### @torch.compile def _dequantize_blockwise_pytorch( A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype, ) -> torch.Tensor: """ Pure PyTorch reference implementation for block-wise dequantization. """ if A.numel() == 0: return torch.empty_like(A, dtype=dtype) A_flat = A.flatten() num_elements = A_flat.numel() dequantized_flat = code.to(A.device)[A_flat.long()].to(dtype) num_blocks = math.ceil(num_elements / blocksize) pad_len = num_blocks * blocksize - num_elements if pad_len > 0: dequantized_flat = torch.nn.functional.pad(dequantized_flat, (0, pad_len)) dequantized_blocks = dequantized_flat.reshape(num_blocks, blocksize) rescaled_blocks = dequantized_blocks * absmax.unsqueeze(1).to(dtype) rescaled_flat = rescaled_blocks.flatten() if pad_len > 0: rescaled_flat = rescaled_flat[:-pad_len] return rescaled_flat.reshape(A.shape) @torch.compile def _quantize_blockwise_pytorch( A: torch.Tensor, code: torch.Tensor, blocksize: int, ) -> tuple[torch.Tensor, torch.Tensor]: """ Pure PyTorch reference implementation for block-wise quantization. """ if A.numel() == 0: return torch.empty_like(A, dtype=torch.uint8), torch.empty(0, dtype=torch.float32, device=A.device) A_flat = A.flatten() num_elements = A_flat.numel() num_blocks = math.ceil(num_elements / blocksize) pad_len = num_blocks * blocksize - num_elements if pad_len > 0: A_flat = torch.nn.functional.pad(A_flat, (0, pad_len)) A_blocks = A_flat.reshape(num_blocks, blocksize) absmax = torch.max(torch.abs(A_blocks), dim=1, keepdim=True)[0] absmax[absmax == 0] = 1.0 scaled_blocks = A_blocks / absmax # Inefficient but straightforward quantization, takes a lot of memory diff = torch.abs(scaled_blocks.unsqueeze(2) - code.to(A.device)) quantized_indices = torch.argmin(diff, dim=2).to(torch.uint8) quantized_flat = quantized_indices.flatten() if pad_len > 0: quantized_flat = quantized_flat[:-pad_len] return quantized_flat.reshape(A.shape), absmax.flatten() # Main updated function def optimizer_update_8bit_blockwise_pytorch( p: torch.Tensor, g: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], beta1: float, beta2: float, beta3: float, # ADEMIX alpha: float, # ADEMIX eps: float, step: int, lr: float, qmap1: torch.Tensor, qmap2: Optional[torch.Tensor], absmax1: torch.Tensor, absmax2: Optional[torch.Tensor], weight_decay: float, gnorm_scale: float, skip_zeros: bool, # ADEMIX *, optimizer_name: str, ) -> None: """ Pure PyTorch implementation of the 8-bit block-wise optimizer update step. This version ensures high-precision updates for float16 parameters. """ if skip_zeros: raise ValueError("skip_zeros is not supported on XPU yet.") blocksize = 256 with torch.no_grad(): # Dequantize states to perform updates in 32-bit precision if optimizer_name == "ademamix" and absmax1.ndim == 2: # For AdEMAMix, state1 holds two EMAs, so absmax1 is stacked. s1_1_fp32 = _dequantize_blockwise_pytorch(state1[0], absmax1[0], qmap1, blocksize, torch.float32) s1_2_fp32 = _dequantize_blockwise_pytorch(state1[1], absmax1[1], qmap1, blocksize, torch.float32) state1_fp32 = torch.stack([s1_1_fp32, s1_2_fp32]) else: state1_fp32 = _dequantize_blockwise_pytorch(state1, absmax1, qmap1, blocksize, torch.float32) state2_fp32 = None if state2 is not None: state2_fp32 = _dequantize_blockwise_pytorch(state2, absmax2, qmap2, blocksize, torch.float32) grad = g.float() * gnorm_scale # Create a 32-bit copy of the parameter for high-precision updates p_fp32 = p.data.float() if optimizer_name == "adam": state1_fp32.mul_(beta1).add_(grad, alpha=1.0 - beta1) state2_fp32.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) bias_correction1 = 1.0 - beta1**step bias_correction2 = 1.0 - beta2**step denom = (state2_fp32.sqrt() / math.sqrt(bias_correction2)).add_(eps) if weight_decay > 0.0: p_fp32.mul_(1.0 - lr * weight_decay) p_fp32.addcdiv_(state1_fp32, denom, value=-lr / bias_correction1) elif optimizer_name == "ademamix": m1_fp32, m2_fp32 = state1_fp32[0], state1_fp32[1] nu_fp32 = state2_fp32 m1_fp32.mul_(beta1).add_(grad, alpha=1.0 - beta1) m2_fp32.mul_(beta3).add_(grad, alpha=1.0 - beta3) nu_fp32.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) bias_correction1 = 1.0 - beta1**step bias_correction2 = math.sqrt(1.0 - beta2**step) update = (m1_fp32 / bias_correction1 + alpha * m2_fp32) / (nu_fp32.sqrt() / bias_correction2 + eps) if weight_decay > 0.0: p_fp32.mul_(1.0 - lr * weight_decay) p_fp32.add_(update, alpha=-lr) state1_fp32 = torch.stack([m1_fp32, m2_fp32]) elif optimizer_name == "momentum": grad.add_(p_fp32, alpha=weight_decay) if step == 1: state1_fp32.copy_(grad) else: state1_fp32.mul_(beta1).add_(grad) p_fp32.add_(state1_fp32, alpha=-lr) elif optimizer_name == "rmsprop": grad.add_(p_fp32, alpha=weight_decay) state1_fp32.mul_(beta1).addcmul_(grad, grad, value=1.0 - beta1) p_fp32.addcdiv_(grad, state1_fp32.sqrt().add_(eps), value=-lr) elif optimizer_name == "lion": if weight_decay > 0.0: p_fp32.mul_(1.0 - lr * weight_decay) update_dir = torch.sign(state1_fp32.mul(beta1) + grad.mul(1.0 - beta1)) p_fp32.add_(update_dir, alpha=-lr) state1_fp32.mul_(beta2).add_(grad, alpha=1.0 - beta2) elif optimizer_name == "adagrad": grad.add_(p_fp32, alpha=weight_decay) state1_fp32.addcmul_(grad, grad, value=1.0) p_fp32.addcdiv_(grad, state1_fp32.sqrt().add_(eps), value=-lr) else: raise NotImplementedError( f"Pure PyTorch implementation for optimizer '{optimizer_name}' is not available." ) # Copy the updated 32-bit parameter back to the original tensor p.data.copy_(p_fp32) # Re-quantize states and update state tensors in-place if optimizer_name == "ademamix": new_m1_8bit, new_absmax_m1 = _quantize_blockwise_pytorch(state1_fp32[0], qmap1, blocksize) new_m2_8bit, new_absmax_m2 = _quantize_blockwise_pytorch(state1_fp32[1], qmap1, blocksize) state1[0].copy_(new_m1_8bit) state1[1].copy_(new_m2_8bit) absmax1[0].copy_(new_absmax_m1) absmax1[1].copy_(new_absmax_m2) new_state2_8bit, new_absmax2 = _quantize_blockwise_pytorch(state2_fp32, qmap2, blocksize) state2.copy_(new_state2_8bit) absmax2.copy_(new_absmax2) else: new_state1_8bit, new_absmax1 = _quantize_blockwise_pytorch(state1_fp32, qmap1, blocksize) state1.copy_(new_state1_8bit) absmax1.copy_(new_absmax1) if state2_fp32 is not None: new_state2_8bit, new_absmax2 = _quantize_blockwise_pytorch(state2_fp32, qmap2, blocksize) state2.copy_(new_state2_8bit) absmax2.copy_(new_absmax2) ####################################### # Mixed torch + triton implementation # ####################################### # Much more memory efficient due to using triton for quantization/dequantization def optimizer_update_8bit_blockwise_triton_quant( p: torch.Tensor, g: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], beta1: float, beta2: float, beta3: float, # ADEMIX alpha: float, # ADEMIX eps: float, step: int, lr: float, qmap1: torch.Tensor, qmap2: Optional[torch.Tensor], absmax1: torch.Tensor, absmax2: Optional[torch.Tensor], weight_decay: float, gnorm_scale: float, skip_zeros: bool, # ADEMIX *, optimizer_name: str, ) -> None: """ Pure PyTorch implementation of the 8-bit block-wise optimizer update step. This version ensures high-precision updates for float16 parameters. """ if skip_zeros and not torch.any(g): return blocksize = 256 grad = g.float() * gnorm_scale with torch.no_grad(): # Create a 32-bit copy of the parameter for high-precision updates p_fp32 = p.data.float() # Dequantize states to perform updates in 32-bit precision if optimizer_name == "ademamix" and absmax1.ndim == 2: # For AdEMAMix, state1 holds two EMAs, so absmax1 is stacked. s1_1_fp32 = dequant_8bit_blockwise(state1[0], absmax1[0], qmap1, blocksize, dtype=torch.float32) s1_2_fp32 = dequant_8bit_blockwise(state1[1], absmax1[1], qmap1, blocksize, dtype=torch.float32) state1_fp32 = torch.stack([s1_1_fp32, s1_2_fp32]) else: state1_fp32 = dequant_8bit_blockwise(state1, absmax1, qmap1, blocksize, dtype=torch.float32) state2_fp32 = None if state2 is not None: state2_fp32 = dequant_8bit_blockwise(state2, absmax2, qmap2, blocksize, dtype=torch.float32) # Apply optimizer-specific update logic if optimizer_name == "adam": if weight_decay > 0.0: p_fp32.mul_(1.0 - lr * weight_decay) state1_fp32.mul_(beta1).add_(grad, alpha=1.0 - beta1) state2_fp32.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) bias_correction1 = 1.0 - beta1**step bias_correction2 = 1.0 - beta2**step denom = (state2_fp32.sqrt() / math.sqrt(bias_correction2)).add_(eps) p_fp32.addcdiv_(state1_fp32, denom, value=-lr / bias_correction1) elif optimizer_name == "ademamix": m1_fp32, m2_fp32 = state1_fp32[0], state1_fp32[1] nu_fp32 = state2_fp32 m1_fp32.mul_(beta1).add_(grad, alpha=1.0 - beta1) m2_fp32.mul_(beta3).add_(grad, alpha=1.0 - beta3) nu_fp32.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) bias_correction1 = 1.0 - beta1**step bias_correction2 = math.sqrt(1.0 - beta2**step) update = (m1_fp32 / bias_correction1 + alpha * m2_fp32) / (nu_fp32.sqrt() / bias_correction2 + eps) if weight_decay > 0.0: p_fp32.mul_(1.0 - lr * weight_decay) p_fp32.add_(update, alpha=-lr) state1_fp32 = torch.stack([m1_fp32, m2_fp32]) elif optimizer_name == "momentum": grad.add_(p_fp32, alpha=weight_decay) if step == 1: state1_fp32.copy_(grad) else: state1_fp32.mul_(beta1).add_(grad) p_fp32.add_(state1_fp32, alpha=-lr) elif optimizer_name == "rmsprop": grad.add_(p_fp32, alpha=weight_decay) state1_fp32.mul_(beta1).addcmul_(grad, grad, value=1.0 - beta1) p_fp32.addcdiv_(grad, state1_fp32.sqrt().add_(eps), value=-lr) elif optimizer_name == "lion": if weight_decay > 0.0: p_fp32.mul_(1.0 - lr * weight_decay) update_dir = torch.sign(state1_fp32.mul(beta1) + grad.mul(1.0 - beta1)) p_fp32.add_(update_dir, alpha=-lr) state1_fp32.mul_(beta2).add_(grad, alpha=1.0 - beta2) elif optimizer_name == "adagrad": grad.add_(p_fp32, alpha=weight_decay) state1_fp32.addcmul_(grad, grad, value=1.0) p_fp32.addcdiv_(grad, state1_fp32.sqrt().add_(eps), value=-lr) else: raise NotImplementedError( f"Pure PyTorch implementation for optimizer '{optimizer_name}' is not available." ) # Copy the updated 32-bit parameter back to the original tensor p.data.copy_(p_fp32) # Re-quantize states and update state tensors in-place if optimizer_name == "ademamix": new_m1_8bit, new_absmax_m1 = quantize_blockwise_triton(state1_fp32[0], qmap1, blocksize) new_m2_8bit, new_absmax_m2 = quantize_blockwise_triton(state1_fp32[1], qmap1, blocksize) state1[0].copy_(new_m1_8bit) state1[1].copy_(new_m2_8bit) absmax1[0].copy_(new_absmax_m1) absmax1[1].copy_(new_absmax_m2) new_state2_8bit, new_absmax2 = quantize_blockwise_triton(state2_fp32, qmap2, blocksize) state2.copy_(new_state2_8bit) absmax2.copy_(new_absmax2) else: new_state1_8bit, new_absmax1 = quantize_blockwise_triton(state1_fp32, qmap1, blocksize) state1.copy_(new_state1_8bit) absmax1.copy_(new_absmax1) if state2_fp32 is not None: new_state2_8bit, new_absmax2 = quantize_blockwise_triton(state2_fp32, qmap2, blocksize) state2.copy_(new_state2_8bit) absmax2.copy_(new_absmax2) ######################### # Triton implementation # ######################### @triton.jit def _optimizer_update_1state_8bit_blockwise_triton_kernel( # Tensors p_ptr, g_ptr, state1_ptr, state2_ptr, beta1: tl.constexpr, beta2: tl.constexpr, beta3, alpha, eps: tl.constexpr, step, beta1_step, beta2_step, lr, qmap1_ptr, qmap2_ptr, absmax1_ptr, absmax2_ptr, weight_decay, gnorm_scale, # Meta-parameters n_elements, BLOCK_SIZE_N: tl.constexpr, N_PER_TH: tl.constexpr, OPTIMIZER_ID: tl.constexpr, ): """ Triton kernel for 8-bit optimizers that use one momentum state. Supports: Momentum, RMSprop, Adagrad, Lion. """ # 1. Boilerplate: pid, offsets, mask pid = tl.program_id(axis=0) block_start_idx = pid * N_PER_TH offsets = block_start_idx * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N * N_PER_TH) mask = offsets < n_elements # 2. Load and dequantize tensors g = tl.load(g_ptr + offsets, mask=mask, other=0.0).to(tl.float32) * gnorm_scale p = tl.load(p_ptr + offsets, mask=mask, other=0.0).to(tl.float32) s1 = dequant_8bit_blockwise_kernel_util(state1_ptr, offsets, qmap1_ptr, absmax1_ptr, mask, BLOCK_SIZE_N) # 3. Optimizer-specific updates # LION if weight_decay > 0.0 and OPTIMIZER_ID == 2: p *= 1.0 - lr * weight_decay # Apply weight decay for momentum, rmsprop, adagrad elif weight_decay > 0.0: g += p * weight_decay # Momentum update if OPTIMIZER_ID == 0: # MOMENTUM if step == 1: s1 = g else: s1 = s1 * beta1 + g p -= lr * s1 # RMSprop update elif OPTIMIZER_ID == 1: # RMSPROP s1 = s1 * beta1 + (1.0 - beta1) * g * g p -= lr * (g / (tl.sqrt(s1) + eps)) # Adagrad update elif OPTIMIZER_ID == 2: # ADAGRAD s1 += g * g p -= lr * (g / (tl.sqrt(s1) + eps)) # Lion update elif OPTIMIZER_ID == 4: # LION val = s1 * beta1 + (1.0 - beta1) * g update = tl.where(val > 0.0, 1.0, tl.where(val < 0.0, -1.0, 0.0)) p -= lr * update s1 = s1 * beta2 + (1.0 - beta2) * g # 4. Store updated parameter and requantized state tl.store(p_ptr + offsets, p.to(p_ptr.dtype.element_ty), mask=mask) s1_codes, new_absmax1 = quantize_8bit_blockwise_kernel_util(s1, qmap1_ptr, 256, BLOCK_SIZE_N, N_PER_TH) tl.store(state1_ptr + offsets, s1_codes, mask=mask) tl.store(absmax1_ptr + block_start_idx + tl.arange(0, N_PER_TH), new_absmax1) @triton.jit def _optimizer_update_2state_8bit_blockwise_triton_kernel( # Tensors p_ptr, g_ptr, state1_ptr, state2_ptr, beta1: tl.constexpr, beta2: tl.constexpr, # ademamix changes alpha and beta3 beta3, # ademamix changes alpha and beta3 alpha, eps: tl.constexpr, step, beta1_step, beta2_step, lr, qmap1_ptr, qmap2_ptr, absmax1_ptr, absmax2_ptr, weight_decay: tl.constexpr, gnorm_scale: tl.constexpr, # Meta-parameters n_elements, BLOCK_SIZE_N: tl.constexpr, N_PER_TH: tl.constexpr, OPTIMIZER_ID: tl.constexpr, ): """ Triton kernel for 8-bit optimizers that use two momentum states. Supports: Adam, AdEMAMix. """ # 1. Boilerplate: pid, offsets, mask pid = tl.program_id(axis=0) block_start_idx = pid * N_PER_TH offsets = block_start_idx * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N * N_PER_TH) mask = offsets < n_elements # 2. Load and dequantize tensors g = tl.load(g_ptr + offsets, mask=mask, other=0.0).to(tl.float32) * gnorm_scale p = tl.load(p_ptr + offsets, mask=mask, other=0.0).to(tl.float32) # 3. Optimizer-specific updates if OPTIMIZER_ID == 3: # ADAM s1 = dequant_8bit_blockwise_kernel_util(state1_ptr, offsets, qmap1_ptr, absmax1_ptr, mask, BLOCK_SIZE_N) s2 = dequant_8bit_blockwise_kernel_util(state2_ptr, offsets, qmap2_ptr, absmax2_ptr, mask, BLOCK_SIZE_N) s1 = s1 * beta1 + (1.0 - beta1) * g s2 = s2 * beta2 + (1.0 - beta2) * g * g # In torch=2.7 on XPU there is an issue with libdevice.pow, leading to an error. # For backwards compatibility we precompute the bias correction factors. # bias_correction1 = 1.0 - libdevice.pow(beta1, step) # bias_correction2 = 1.0 - libdevice.pow(beta2, step) bias_correction1 = 1.0 - beta1_step bias_correction2 = 1.0 - beta2_step if weight_decay > 0.0: p *= 1.0 - lr * weight_decay denom = tl.sqrt(s2) / tl.sqrt(bias_correction2) + eps p -= (lr / bias_correction1) * (s1 / denom) # Store updated parameter tl.store(p_ptr + offsets, p.to(p_ptr.dtype.element_ty), mask=mask) # Requantize and store states s1_codes, new_absmax1 = quantize_8bit_blockwise_kernel_util(s1, qmap1_ptr, 256, BLOCK_SIZE_N, N_PER_TH) tl.store(state1_ptr + offsets, s1_codes, mask=mask) tl.store(absmax1_ptr + block_start_idx + tl.arange(0, N_PER_TH), new_absmax1) s2_codes, new_absmax2 = quantize_8bit_blockwise_kernel_util(s2, qmap2_ptr, 256, BLOCK_SIZE_N, N_PER_TH) tl.store(state2_ptr + offsets, s2_codes, mask=mask) tl.store(absmax2_ptr + block_start_idx + tl.arange(0, N_PER_TH), new_absmax2) elif OPTIMIZER_ID == 5: # ADEMAMIX # AdEMAMix has a stacked state1 (m1, m2) and state2 (nu) m1 = dequant_8bit_blockwise_kernel_util(state1_ptr, offsets, qmap1_ptr, absmax1_ptr, mask, BLOCK_SIZE_N) m2 = dequant_8bit_blockwise_kernel_util( state1_ptr + n_elements, offsets, qmap1_ptr, absmax1_ptr + n_elements // BLOCK_SIZE_N, mask, BLOCK_SIZE_N, ) nu = dequant_8bit_blockwise_kernel_util(state2_ptr, offsets, qmap2_ptr, absmax2_ptr, mask, BLOCK_SIZE_N) m1 = m1 * beta1 + (1.0 - beta1) * g m2 = m2 * beta3 + (1.0 - beta3) * g nu = nu * beta2 + (1.0 - beta2) * g * g # In torch=2.7 on XPU there is an issue with libdevice.pow, leading to an error. # For backwards compatibility we precompute the bias correction factors. # bias_correction1 = 1.0 - libdevice.pow(beta1, step) # bias_correction2 = tl.sqrt(1.0 - libdevice.pow(beta2, step)) bias_correction1 = 1.0 - beta1_step bias_correction2 = tl.sqrt(1.0 - beta2_step) update = (m1 / bias_correction1 + alpha * m2) / (tl.sqrt(nu) / bias_correction2 + eps) if weight_decay > 0.0: p *= 1.0 - lr * weight_decay p -= lr * update # Store updated parameter tl.store(p_ptr + offsets, p.to(p_ptr.dtype.element_ty), mask=mask) # Requantize and store all three states m1_codes, new_absmax_m1 = quantize_8bit_blockwise_kernel_util(m1, qmap1_ptr, 256, BLOCK_SIZE_N, N_PER_TH) tl.store(state1_ptr + offsets, m1_codes, mask=mask) tl.store(absmax1_ptr + block_start_idx + tl.arange(0, N_PER_TH), new_absmax_m1) m2_codes, new_absmax_m2 = quantize_8bit_blockwise_kernel_util(m2, qmap1_ptr, 256, BLOCK_SIZE_N, N_PER_TH) tl.store(state1_ptr + n_elements + offsets, m2_codes, mask=mask) tl.store( absmax1_ptr + block_start_idx + tl.arange(0, N_PER_TH) + n_elements // BLOCK_SIZE_N, new_absmax_m2, ) nu_codes, new_absmax_nu = quantize_8bit_blockwise_kernel_util(nu, qmap2_ptr, 256, BLOCK_SIZE_N, N_PER_TH) tl.store(state2_ptr + offsets, nu_codes, mask=mask) tl.store(absmax2_ptr + block_start_idx + tl.arange(0, N_PER_TH), new_absmax_nu) name2optimizer_fn = { "momentum": _optimizer_update_1state_8bit_blockwise_triton_kernel, "rmsprop": _optimizer_update_1state_8bit_blockwise_triton_kernel, "adagrad": _optimizer_update_1state_8bit_blockwise_triton_kernel, "adam": _optimizer_update_2state_8bit_blockwise_triton_kernel, "lion": _optimizer_update_1state_8bit_blockwise_triton_kernel, "ademamix": _optimizer_update_2state_8bit_blockwise_triton_kernel, } def optimizer_update_8bit_blockwise_impl( optimizer_name: str, g: torch.Tensor, p: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], beta1: float, beta2: float, beta3: float, alpha: float, eps: float, step: int, lr: float, qmap1: torch.Tensor, qmap2: Optional[torch.Tensor], absmax1: torch.Tensor, absmax2: Optional[torch.Tensor], weight_decay: float = 0.0, gnorm_scale: float = 1.0, skip_zeros=False, ) -> None: if skip_zeros: raise NotImplementedError("skip_zeros is not supported on XPU yet") if optimizer_name == "ademamix": # Handle AdEMAMIX's stacked state tensors if state1.dim() < 2 or state1.shape[0] != 2: raise ValueError( f"For ademamix, state1 must be a stacked tensor of shape (2, ...), but got {state1.shape}" ) if absmax1.dim() < 2 or absmax1.shape[0] != 2: raise ValueError( f"For ademamix, absmax1 must be a stacked tensor of shape (2, ...), but got {absmax1.shape}" ) BLOCK_SIZE = 256 N_PER_TH = 1 # Number of blocks processed per thread. grid = (triton.cdiv(p.numel(), BLOCK_SIZE * N_PER_TH),) fn = name2optimizer_fn[optimizer_name] optimizer_id = name2optimizer_id[optimizer_name] # In torch=2.7 on XPU there is an issue with libdevice.pow, leading to an error. # For backwards compatibility we precompute the bias correction factors. beta1_step = beta1**step beta2_step = beta2**step fn[grid]( p, g, state1, state2, beta1, beta2, beta3, alpha, eps, step, beta1_step, beta2_step, lr, qmap1, qmap2, absmax1, absmax2, weight_decay, gnorm_scale, p.numel(), BLOCK_SIZE_N=BLOCK_SIZE, N_PER_TH=N_PER_TH, OPTIMIZER_ID=optimizer_id, num_warps=2, ) # optimizer_update_8bit_blockwise_impl = optimizer_update_8bit_blockwise_pytorch # optimizer_update_8bit_blockwise_impl = torch.compile(optimizer_update_8bit_blockwise_pytorch_impl) # optimizer_update_8bit_blockwise_impl = optimizer_update_8bit_blockwise_triton_quant # optimizer_update_8bit_blockwise_impl = torch.compile(optimizer_update_8bit_blockwise_triton_quant) optimizer_update_8bit_blockwise_impl = optimizer_update_8bit_blockwise_impl bitsandbytes/bitsandbytes/backends/triton/ops.py000066400000000000000000000232721506752665600225400ustar00rootroot00000000000000from collections.abc import Sequence from typing import Optional import torch from . import kernels_4bit, kernels_8bit_quant, kernels_optim # currently codes unused, kept for reference # Should be the same for quant/dequant # from bitsandbytes.functional import get_4bit_type # _FP4_QUANT_TABLE = get_4bit_type("fp4", device="xpu") # _NF4_QUANT_TABLE = get_4bit_type("nf4", device="xpu") device_type = torch.accelerator.current_accelerator().type if hasattr(torch, "accelerator") else "cuda" torch_accelerator_module = getattr(torch, device_type, torch.cuda) def quantize_blockwise(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]: torch._check_is_size(blocksize) # torch._check(A.dtype == torch.float32, lambda: f"A must be float32 on xpu, got {A.dtype}") with torch_accelerator_module.device(A.device): out, absmax = kernels_8bit_quant.quantize_blockwise_triton(A, code, blocksize) return out, absmax.float() def dequantize_blockwise( A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype ) -> torch.Tensor: torch._check_is_size(blocksize) torch._check(A.dtype == torch.uint8, lambda: f"A must be uint8, got {A.dtype}") # torch._check(dtype == torch.float32, lambda: f"dtype must be float32 on xpu, got {dtype}") with torch_accelerator_module.device(A.device): out = kernels_8bit_quant.dequant_8bit_blockwise( A, absmax, code, blocksize, dtype=dtype, ) return out def dequantize_blockwise_inplace( A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype, out: torch.Tensor, ) -> None: torch._check_is_size(blocksize) torch._check(A.dtype == torch.uint8, lambda: f"A must be uint8, got {A.dtype}") torch._check(out.shape == A.shape, lambda: f"Expected out.shape == {A.shape}, got {out.shape}") torch._check(out.device == A.device, lambda: f"Expected out.device == {A.device}, got {out.device}") torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}") with torch_accelerator_module.device(A.device): kernels_8bit_quant.dequant_8bit_blockwise( A, absmax, code, blocksize, dtype=dtype, out=out, ) def quantize_4bit( A: torch.Tensor, blocksize: int, quant_type: str, quant_storage: torch.dtype ) -> tuple[torch.Tensor, torch.Tensor]: torch._check_is_size(blocksize) # torch._check(quant_type == "nf4", lambda: f"quant_type must be nf4 on CPU, got {quant_type}") torch._check( A.dtype in [torch.bfloat16, torch.float16, torch.float32], lambda: f"Blockwise 4bit quantization only supports 16/32-bit floats, but got {A.dtype}", ) n = A.numel() # TODO: Support when weight matrix is not divisible by blocksize # torch._check(n % blocksize == 0, lambda: f"n must be divisible by blocksize, got {n} and {blocksize}") blocks = -(n // -(blocksize * 2)) absmax = torch.empty((blocks * 2,), device=A.device, dtype=A.dtype) out = torch.empty((n // 2, 1), device=A.device, dtype=torch.uint8) with torch_accelerator_module.device(A.device): kernels_4bit.quantize_4bit_blockwise_triton( A, blocksize, quant_type, blocks, absmax, num_elements=n, quantized_out=out ) packed = out if quant_storage != torch.uint8: packed = out.squeeze().view(quant_storage).unsqueeze(1) return packed, absmax.float() def dequantize_4bit( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, shape: Sequence[int], dtype: torch.dtype, ) -> torch.Tensor: torch._check_is_size(blocksize) # torch._check(quant_type == "nf4", lambda: f"quant_type must be nf4 on XPU, got {quant_type}") torch._check( dtype in [torch.bfloat16, torch.float16, torch.float32], lambda: f"Blockwise 4bit dequantization only supports 16/32-bit floats, but got {dtype}", ) # torch._check( # A.dtype == torch.uint8, # lambda: f"Blockwise 4bit dequantization on XPU only supports uint8 storage, got {A.dtype}", # ) # Check if this is fine and fast if A.dtype != torch.uint8: A = A.squeeze().view(torch.uint8).unsqueeze(1) out = torch.empty(shape, dtype=dtype, device=A.device) with torch_accelerator_module.device(A.device): kernels_4bit.dequantize_4bit_impl(A, absmax, blocksize, quant_type, dtype, out=out) return out def dequantize_4bit_inplace( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, shape: Sequence[int], dtype: torch.dtype, out: torch.Tensor, ) -> None: torch._check(out.shape == shape, lambda: f"Expected out.shape == {shape}, got {out.shape}") torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}") with torch_accelerator_module.device(A.device): kernels_4bit.dequantize_4bit_impl(A, absmax, blocksize, quant_type, dtype, out=out) def gemv_4bit( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int, ) -> torch.Tensor: if B.dtype != torch.uint8: B = B.squeeze().view(torch.uint8).unsqueeze(1) B_dq_triton = torch.empty(shapeB, dtype=A.dtype, device=A.device) with torch_accelerator_module.device(A.device): kernels_4bit.dequantize_4bit_impl_passing_code( B, absmax, blocksize, code, dtype=A.dtype, out=B_dq_triton, ) return torch.nn.functional.linear( A, B_dq_triton, bias=None, ) # optimizer_update_8bit_blockwise_impl = kernels_optim.optimizer_update_8bit_blockwise_pytorch # optimizer_update_8bit_blockwise_impl = torch.compile(kernels_optim.optimizer_update_8bit_blockwise_pytorch) # 60ms # optimizer_update_8bit_blockwise_impl = kernels_optim.optimizer_update_8bit_blockwise_triton_quant #2.8ms # optimizer_update_8bit_blockwise_impl = torch.compile(kernels_optim.optimizer_update_8bit_blockwise_triton_quant) # 2.3ms optimizer_update_8bit_blockwise_impl = kernels_optim.optimizer_update_8bit_blockwise_impl # ~0.95ms for adam def optimizer_update_8bit_blockwise( optimizer_name: str, g: torch.Tensor, p: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], beta1: float, beta2: float, beta3: float, alpha: float, eps: float, step: int, lr: float, qmap1: torch.Tensor, qmap2: Optional[torch.Tensor], absmax1: torch.Tensor, absmax2: Optional[torch.Tensor], weight_decay: float = 0.0, gnorm_scale: float = 1.0, skip_zeros=False, ) -> None: # torch._check( # g.numel() == p.numel(), # lambda: f"g and p must have the same number of elements, got {g.numel()} and {p.numel()}", # ) # compute_dtypes = [torch.float16, torch.bfloat16, torch.float32] # torch._check( # g.dtype in compute_dtypes, # lambda: f"g must be bfloat16, float16, or float32, got {g.dtype}", # ) # torch._check( # g.dtype == p.dtype, # lambda: f"Expected all tensors to have the same dtype, got g.dtype={g.dtype}, p.dtype={p.dtype}", # ) # torch._check( # state1.dtype == torch.uint8, # lambda: f"state1 must be uint8, got {state1.dtype}", # ) # torch._check( # qmap1.dtype == absmax1.dtype == torch.float32, # lambda: f"Expected qmap1 and absmax1 to be float32, got qmap1.dtype={qmap1.dtype}, absmax1.dtype={absmax1.dtype}", # ) # if state2 is not None: # torch._check( # state2.dtype == torch.uint8, # lambda: f"state2 must be uint8, got {state2.dtype}", # ) # torch._check( # qmap2.dtype == absmax2.dtype == torch.float32, # lambda: f"Expected qmap2 and absmax2 to be float32, got qmap2.dtype={qmap2.dtype}, absmax2.dtype={absmax2.dtype}", # ) with torch_accelerator_module.device(state1.device): optimizer_update_8bit_blockwise_impl( optimizer_name=optimizer_name, g=g, p=p, state1=state1, state2=state2, beta1=beta1, beta2=beta2, beta3=beta3, alpha=alpha, eps=eps, step=step, lr=lr, qmap1=qmap1, qmap2=qmap2, absmax1=absmax1, absmax2=absmax2, weight_decay=weight_decay, gnorm_scale=gnorm_scale, skip_zeros=skip_zeros, ) def optimizer_update_32bit( optimizer_name: str, g: torch.Tensor, p: torch.Tensor, state1: torch.Tensor, state2: Optional[torch.Tensor], unorm_vec: Optional[torch.Tensor], max_unorm: float, param_norm: float, beta1: float, beta2: float, beta3: float, alpha: float, eps: float, weight_decay: float, step: int, lr: float, gnorm_scale: float, skip_zeros=False, ) -> None: with torch_accelerator_module.device(state1.device): kernels_optim.optimizer_update_32bit_impl( optimizer_name=optimizer_name, g=g, p=p, state1=state1, state2=state2, unorm_vec=unorm_vec, max_unorm=max_unorm, param_norm=param_norm, beta1=beta1, beta2=beta2, beta3=beta3, alpha=alpha, eps=eps, weight_decay=weight_decay, step=step, lr=lr, gnorm_scale=gnorm_scale, skip_zeros=skip_zeros, ) bitsandbytes/bitsandbytes/backends/utils.py000066400000000000000000000034211506752665600215520ustar00rootroot00000000000000import subprocess from packaging import version import torch try: import triton # noqa: F401 import triton.language as tl # noqa: F401 triton_available = True except ImportError: triton_available = False _NF4_QUANT_TABLE = torch.tensor( [ -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453, -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0, 0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224, 0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0, ], dtype=torch.float32, device="xpu" if hasattr(torch, "xpu") and torch.xpu.is_available() else "cpu", # Only cpu/xpu use this table for now. ) _FP4_QUANT_TABLE = torch.tensor( [ 0.0000, 0.0052, 0.6667, 1.0000, 0.3333, 0.5000, 0.1667, 0.2500, 0.0000, -0.0052, -0.6667, -1.0000, -0.3333, -0.5000, -0.1667, -0.2500, ], dtype=torch.float32, device="xpu" if hasattr(torch, "xpu") and torch.xpu.is_available() else "cpu", # Only cpu/xpu use this table for now. ) CODE = {"nf4": _NF4_QUANT_TABLE, "fp4": _FP4_QUANT_TABLE} def get_gaudi_sw_version(): """ Returns the installed version of Gaudi SW. """ output = subprocess.run( "pip list | grep habana-torch-plugin", shell=True, text=True, capture_output=True, ) # If grep return nothing if not output.stdout.strip(): return None return version.parse(output.stdout.split("\n")[0].split()[-1]) GAUDI_SW_VER = get_gaudi_sw_version() bitsandbytes/bitsandbytes/backends/xpu/000077500000000000000000000000001506752665600206545ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/xpu/__init__.py000066400000000000000000000000001506752665600227530ustar00rootroot00000000000000bitsandbytes/bitsandbytes/backends/xpu/ops.py000066400000000000000000000175541506752665600220430ustar00rootroot00000000000000from collections.abc import Sequence import ctypes as ct import logging from packaging import version import torch from bitsandbytes.functional import _get_tensor_stream, get_ptr from ..._ops import register_kernel from ...cextension import ErrorHandlerMockBNBNativeLibrary, lib from ..utils import triton_available logger = logging.getLogger(__name__) # _int_mm is available in torch starting from 2.9 version if version.parse(torch.__version__).release >= version.parse("2.9").release: @register_kernel("bitsandbytes::int8_linear_matmul", "xpu") def _(A: torch.Tensor, B: torch.Tensor): return torch._int_mm( A.reshape(-1, A.shape[-1]), B.t(), ).reshape(*A.shape[:-1], B.shape[0]) def _dequantize_4bit_impl( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, dtype: torch.dtype, out: torch.Tensor, ) -> None: args = ( None, get_ptr(A), get_ptr(absmax), get_ptr(out), ct.c_int(blocksize), ct.c_int(out.numel()), _get_tensor_stream(A), ) if dtype == torch.bfloat16: if quant_type == "fp4": lib.cdequantize_blockwise_bf16_fp4(*args) else: lib.cdequantize_blockwise_bf16_nf4(*args) elif dtype == torch.float16: if quant_type == "fp4": lib.cdequantize_blockwise_fp16_fp4(*args) else: lib.cdequantize_blockwise_fp16_nf4(*args) elif dtype == torch.float32: if quant_type == "fp4": lib.cdequantize_blockwise_fp32_fp4(*args) else: lib.cdequantize_blockwise_fp32_nf4(*args) def _dequantize_blockwise_impl( A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype, out: torch.Tensor ) -> None: args = ( get_ptr(code), get_ptr(A), get_ptr(absmax), get_ptr(out), ct.c_int(blocksize), ct.c_int(A.numel()), _get_tensor_stream(A), ) if dtype == torch.float16: lib.cdequantize_blockwise_fp16(*args) elif dtype == torch.bfloat16: lib.cdequantize_blockwise_bf16(*args) elif dtype == torch.float32: lib.cdequantize_blockwise_fp32(*args) def _gemv_4bit_impl( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int, out: torch.Tensor, ) -> None: m = ct.c_int32(1) n = ct.c_int32(shapeB[0]) k = ct.c_int32(shapeB[1]) lda = m ldb = ct.c_int32((A.shape[-1] + 1) // 2) ldc = m stream = _get_tensor_stream(A) if A.dtype == torch.float16: lib.cgemv_4bit_inference_fp16( m, n, k, get_ptr(A), get_ptr(B), get_ptr(absmax), get_ptr(code), get_ptr(out), lda, ldb, ldc, ct.c_int32(blocksize), stream, ) elif A.dtype == torch.bfloat16: lib.cgemv_4bit_inference_bf16( m, n, k, get_ptr(A), get_ptr(B), get_ptr(absmax), get_ptr(code), get_ptr(out), lda, ldb, ldc, ct.c_int32(blocksize), stream, ) elif A.dtype == torch.float32: lib.cgemv_4bit_inference_fp32( m, n, k, get_ptr(A), get_ptr(B), get_ptr(absmax), get_ptr(code), get_ptr(out), lda, ldb, ldc, ct.c_int32(blocksize), stream, ) # SYCL should be faster for xpu, so at first checking if it is available. if not isinstance(lib, ErrorHandlerMockBNBNativeLibrary): logger.info("Register sycl bitsandbytes kernels for XPU") # TODO: Remove the triton register when quantization sycl kernel is ready. if triton_available: from ..triton import ops as triton_ops register_kernel("bitsandbytes::quantize_blockwise", "xpu")(triton_ops.quantize_blockwise) register_kernel("bitsandbytes::quantize_4bit", "xpu")(triton_ops.quantize_4bit) register_kernel("bitsandbytes::optimizer_update_8bit_blockwise", "xpu")( triton_ops.optimizer_update_8bit_blockwise ) register_kernel("bitsandbytes::optimizer_update_32bit", "xpu")(triton_ops.optimizer_update_32bit) @register_kernel("bitsandbytes::dequantize_4bit", "xpu") def _( A: torch.Tensor, absmax: torch.Tensor, blocksize: int, quant_type: str, shape: Sequence[int], dtype: torch.dtype, ) -> torch.Tensor: out = torch.empty(shape, dtype=dtype, device=A.device) _dequantize_4bit_impl(A, absmax, blocksize, quant_type, dtype, out=out) return out @register_kernel("bitsandbytes::dequantize_blockwise", "xpu") def _( A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype ) -> torch.Tensor: out = torch.empty_like(A, dtype=dtype) _dequantize_blockwise_impl(A, absmax, code, blocksize, dtype, out=out) return out @register_kernel("bitsandbytes::dequantize_blockwise.out", "xpu") def _( A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype, out: torch.Tensor, ) -> None: torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}") torch._check(out.shape == A.shape, lambda: f"Expected out.shape == {A.shape}, got {out.shape}") _dequantize_blockwise_impl(A, absmax, code, blocksize, dtype, out=out) @register_kernel("bitsandbytes::gemv_4bit", "xpu") def _( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int, ) -> torch.Tensor: shape = (*A.shape[:-1], shapeB[0]) out = torch.empty(shape, device=A.device, dtype=A.dtype) _gemv_4bit_impl(A, B, shapeB, absmax, code, blocksize, out=out) return out @register_kernel("bitsandbytes::gemv_4bit.out", "xpu") def _( A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int, out: torch.Tensor, ) -> None: torch._check( out.shape == (*A.shape[:-1], shapeB[0]), lambda: f"Expected out.shape == {(*A.shape[:-1], shapeB[0])}, got {out.shape}", ) torch._check(out.dtype == A.dtype, lambda: f"Expected out.dtype == {A.dtype}, got {out.dtype}") _gemv_4bit_impl(A, B, shapeB, absmax, code, blocksize, out=out) elif triton_available: logger.info("Register triton bitsandbytes kernels for XPU") from ..triton import ops as triton_ops register_kernel("bitsandbytes::quantize_blockwise", "xpu")(triton_ops.quantize_blockwise) register_kernel("bitsandbytes::dequantize_blockwise.out", "xpu")(triton_ops.dequantize_blockwise_inplace) register_kernel("bitsandbytes::dequantize_blockwise", "xpu")(triton_ops.dequantize_blockwise) register_kernel("bitsandbytes::quantize_4bit", "xpu")(triton_ops.quantize_4bit) register_kernel("bitsandbytes::dequantize_4bit.out", "xpu")(triton_ops.dequantize_4bit_inplace) register_kernel("bitsandbytes::dequantize_4bit", "xpu")(triton_ops.dequantize_4bit) register_kernel("bitsandbytes::gemv_4bit", "xpu")(triton_ops.gemv_4bit) register_kernel("bitsandbytes::optimizer_update_8bit_blockwise", "xpu")(triton_ops.optimizer_update_8bit_blockwise) register_kernel("bitsandbytes::optimizer_update_32bit", "xpu")(triton_ops.optimizer_update_32bit) else: logger.warning("Register pytorch bitsandbytes kernels for XPU because no native library or triton packages found.") bitsandbytes/bitsandbytes/cextension.py000066400000000000000000000325571506752665600210330ustar00rootroot00000000000000import ctypes as ct import functools import logging import os from pathlib import Path import re from typing import Optional import torch from bitsandbytes.consts import DYNAMIC_LIBRARY_SUFFIX, PACKAGE_DIR from bitsandbytes.cuda_specs import CUDASpecs, get_cuda_specs, get_cuda_version_tuple, get_rocm_gpu_arch logger = logging.getLogger(__name__) def get_cuda_bnb_library_path(cuda_specs: CUDASpecs) -> Path: """ Get the disk path to the CUDA BNB native library specified by the given CUDA specs, taking into account the `BNB_CUDA_VERSION` override environment variable. The library is not guaranteed to exist at the returned path. """ prefix = "rocm" if torch.version.hip else "cuda" library_name = f"libbitsandbytes_{prefix}{cuda_specs.cuda_version_string}{DYNAMIC_LIBRARY_SUFFIX}" override_value = os.environ.get("BNB_CUDA_VERSION") if override_value: library_name = re.sub(r"cuda\d+", f"cuda{override_value}", library_name, count=1) if torch.version.hip: raise RuntimeError( f"BNB_CUDA_VERSION={override_value} detected for ROCm!! \n" f"Clear the variable and retry: export BNB_CUDA_VERSION=\n" ) logger.warning( f"WARNING: BNB_CUDA_VERSION={override_value} environment variable detected; loading {library_name}.\n" "This can be used to load a bitsandbytes version built with a CUDA version that is different from the PyTorch CUDA version.\n" "If this was unintended set the BNB_CUDA_VERSION variable to an empty string: export BNB_CUDA_VERSION=\n" ) return PACKAGE_DIR / library_name class BNBNativeLibrary: _lib: ct.CDLL compiled_with_cuda = False def __init__(self, lib: ct.CDLL): self._lib = lib @functools.cache # noqa: B019 def __getattr__(self, name): fn = getattr(self._lib, name, None) if fn is not None: return fn def throw_on_call(*args, **kwargs): raise RuntimeError( f"Method '{name}' not available in CPU-only version of bitsandbytes.\n" "Reinstall with GPU support or use CUDA-enabled hardware." ) return throw_on_call def __getitem__(self, item): return self.__getattr__(item) class CudaBNBNativeLibrary(BNBNativeLibrary): compiled_with_cuda = True def __init__(self, lib: ct.CDLL): super().__init__(lib) lib.get_context.restype = ct.c_void_p lib.get_cusparse.restype = ct.c_void_p lib.cget_managed_ptr.restype = ct.c_void_p def get_available_cuda_binary_versions() -> list[str]: """Get formatted CUDA versions from existing library files using cuda_specs logic""" lib_pattern = f"libbitsandbytes_{BNB_BACKEND.lower()}*{DYNAMIC_LIBRARY_SUFFIX}" versions = [] for lib in Path(__file__).parent.glob(lib_pattern): pattern = rf"{BNB_BACKEND.lower()}(\d+)" match = re.search(pattern, lib.name) if match: ver_code = int(match.group(1)) major = ver_code // 10 minor = ver_code % 10 versions.append(f"{major}.{minor}") return sorted(versions) def parse_cuda_version(version_str: str) -> str: """Convert raw version string (e.g. '118' from env var) to formatted version (e.g. '11.8')""" if version_str.isdigit(): return f"{version_str[:-1]}.{version_str[-1]}" return version_str # fallback as safety net class ErrorHandlerMockBNBNativeLibrary(BNBNativeLibrary): """ Mock library handler that defers errors until native methods are called. This class serves as a fallback when the native bitsandbytes library fails to load. It captures the original error and generates detailed troubleshooting guidance. Key behaviors: - Allows attribute access and method assignment without immediate errors - Throws a RuntimeError with diagnostic information only when a native method is called, as otherwise it would error out on import, breaking backward compatibility - Handles both missing CUDA dependencies and version mismatch scenarios Error scenarios covered: 1. Missing shared library dependencies (e.g., libcudart.so not in LD_LIBRARY_PATH or through PyTorch CUDA installation) 2. CUDA version mismatch between PyTorch and available pre-compiled binaries 3. Completely missing pre-compiled binaries when CUDA is detected 4. Custom BNB_CUDA_VERSION override but mismatch 5. CPU-only installation attempts when GPU functionality is requested """ def __init__(self, error_msg: str): self.error_msg = error_msg self.user_cuda_version = get_cuda_version_tuple() self.available_versions = get_available_cuda_binary_versions() self.override_value = os.environ.get("BNB_CUDA_VERSION") self.requested_version = ( parse_cuda_version(self.override_value) if self.override_value else f"{self.user_cuda_version[0]}.{self.user_cuda_version[1]}" if self.user_cuda_version else "unknown" ) # Pre-generate the error message based on error type if "cannot open shared object file" in error_msg: self.formatted_error = self._format_dependency_error() else: # lib loading errors self.formatted_error = self._format_lib_error_message( available_versions=self.available_versions, user_cuda_version=f"{self.user_cuda_version[0]}.{self.user_cuda_version[1]}" if self.user_cuda_version else "unknown", original_error=f"Original error: {self.error_msg}\n" if self.error_msg else "", requested_version=self.requested_version, ) def _format_lib_error_message( self, available_versions: list[str], user_cuda_version: str, original_error: str = "", requested_version: Optional[str] = None, ) -> str: """Format detailed error message for library loading failures""" analysis = "" no_cpu_lib_found = "libbitsandbytes_cpu.so: cannot open" in original_error no_cuda_lib_found = f"{BNB_BACKEND} binary not found" in original_error if no_cpu_lib_found: analysis = "\n🚨 Failed to load CPU-only bitsandbytes library 🚨\n\n" elif no_cuda_lib_found: version_list_str = "\n - " + "\n - ".join(available_versions) if available_versions else "NONE" analysis = ( ( f"\n🚨 {BNB_BACKEND} VERSION MISMATCH 🚨\n" f"Requested {BNB_BACKEND} version: {requested_version}\n" f"Detected PyTorch {BNB_BACKEND} version: {user_cuda_version}\n" f"Available pre-compiled versions: {version_list_str}\n\n" "This means:\n" "The version you're trying to use is NOT distributed with this package\n\n" ) if available_versions else "\n🚨 Forgot to compile the bitsandbytes library? 🚨\n" "1. You're not using the package but checked-out the source code\n" "2. You MUST compile from source\n\n" ) base_msg = "Attempted to use bitsandbytes native library functionality but it's not available.\n\n" troubleshooting = ( ( f"This typically happens when:\n" f"1. bitsandbytes doesn't ship with a pre-compiled binary for your {BNB_BACKEND} version\n" f"2. The library wasn't compiled properly during installation from source\n\n" ) if no_cuda_lib_found else f"This typically happens when you checked the code out from source and your torch installation doesn't detect {BNB_BACKEND} on your machine.\n\n" ) note = ( ( f"To make bitsandbytes work, the compiled library version MUST exactly match the linked {BNB_BACKEND} version.\n" f"If your {BNB_BACKEND} version doesn't have a pre-compiled binary, you MUST compile from source.\n\n" ) if no_cuda_lib_found else "" ) compile_instructions = ( ("COMPILE FROM SOURCE for CPU-only:\n `cmake -DCOMPUTE_BACKEND=cpu -S . && make`\n\n") if not no_cuda_lib_found else ( "You have two options:\n" "1. COMPILE FROM SOURCE (required if no binary exists):\n" " https://huggingface.co/docs/bitsandbytes/main/en/installation#cuda-compile\n" "2. Use BNB_CUDA_VERSION to specify a DIFFERENT CUDA version from the detected one, which is installed on your machine and matching an available pre-compiled version listed above\n\n" ) if not HIP_ENVIRONMENT else ( "You can COMPILE FROM SOURCE as mentioned here:\n" " https://huggingface.co/docs/bitsandbytes/main/en/installation?backend=AMD+ROCm#amd-gpu\n" ) ) diagnostics = ( f"🔍 Run this command for detailed diagnostics:\n" f"python -m bitsandbytes\n\n" f"If you've tried everything and still have issues:\n" f"1. Include ALL version info (operating system, bitsandbytes, pytorch, {BNB_BACKEND.lower()}, python)\n" f"2. Describe what you've tried in detail\n" f"3. Open an issue with this information:\n" f" https://github.com/bitsandbytes-foundation/bitsandbytes/issues\n\n" ) return f"{analysis}{base_msg}{troubleshooting}{note}{compile_instructions}{original_error}\n{diagnostics}" def _format_dependency_error(self) -> str: """Format error message for missing shared libraries""" # Extract missing library name from error error_parts = self.error_msg.split(":") missing_lib = error_parts[0].strip() if len(error_parts) > 0 else "unknown library" cuda_major_version = ( self.requested_version.split(".")[0] if "." in self.requested_version else self.requested_version ) return ( f"\n🚨 {BNB_BACKEND} SETUP ERROR: Missing dependency: {missing_lib} 🚨\n\n" f"{BNB_BACKEND} {cuda_major_version}.x runtime libraries were not found in the LD_LIBRARY_PATH.\n\n" f"To fix this, make sure that:\n" f"1. You have installed {BNB_BACKEND} {cuda_major_version}.x toolkit on your system\n" f"2. The {BNB_BACKEND} runtime libraries are in your LD_LIBRARY_PATH\n\n" f"You can add them with (and persist the change by adding the line to your .bashrc):\n" f" export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/{BNB_BACKEND.lower()}-{cuda_major_version}.x/\ {'lib64' if not HIP_ENVIRONMENT else 'lib'}\n\n" f"Original error: {self.error_msg}\n\n" f"🔍 Run this command for detailed diagnostics:\n" f"python -m bitsandbytes\n\n" f"If you've tried everything and still have issues:\n" f"1. Include ALL version info (operating system, bitsandbytes, pytorch, {BNB_BACKEND.lower()}, python)\n" f"2. Describe what you've tried in detail\n" f"3. Open an issue with this information:\n" f" https://github.com/bitsandbytes-foundation/bitsandbytes/issues\n\n" ) def __getattr__(self, name): """Return a dummy function that throws when called, rather than on attribute access""" def throw_on_call(*args, **kwargs): raise RuntimeError(f"{self.formatted_error}Native code method attempted to call: lib.{name}()") return throw_on_call def __getitem__(self, name): return self.__getattr__(name) def get_native_library() -> BNBNativeLibrary: """ Load CUDA library XOR CPU, as the latter contains a subset of symbols of the former. """ cuda_specs = get_cuda_specs() binary_path = PACKAGE_DIR / f"libbitsandbytes_cpu{DYNAMIC_LIBRARY_SUFFIX}" if cuda_specs: cuda_binary_path = get_cuda_bnb_library_path(cuda_specs) if not cuda_binary_path.exists(): raise RuntimeError(f"Configured {BNB_BACKEND} binary not found at {cuda_binary_path}") binary_path = cuda_binary_path if torch._C._has_xpu: binary_path = PACKAGE_DIR / f"libbitsandbytes_xpu{DYNAMIC_LIBRARY_SUFFIX}" logger.debug(f"Loading bitsandbytes native library from: {binary_path}") # Try to load the library - any errors will propagate up dll = ct.cdll.LoadLibrary(str(binary_path)) if hasattr(dll, "get_context"): # only a CUDA-built library exposes this return CudaBNBNativeLibrary(dll) return BNBNativeLibrary(dll) ROCM_GPU_ARCH = get_rocm_gpu_arch() HIP_ENVIRONMENT = False BNB_BACKEND = "CPU" if torch.version.hip: HIP_ENVIRONMENT = True BNB_BACKEND = "ROCm" elif torch.cuda.is_available(): BNB_BACKEND = "CUDA" elif torch._C._has_xpu: BNB_BACKEND = "XPU" try: lib = get_native_library() except Exception as e: if BNB_BACKEND in ("CPU", "XPU"): lib = ErrorHandlerMockBNBNativeLibrary("XPU/CPU can run without native library.") else: error_msg = str(e) logger.error( f"bitsandbytes library load error: {error_msg}", exc_info=True, ) # create a mock with error messaging as fallback lib = ErrorHandlerMockBNBNativeLibrary(error_msg) bitsandbytes/bitsandbytes/consts.py000066400000000000000000000005741506752665600201570ustar00rootroot00000000000000from pathlib import Path import platform DYNAMIC_LIBRARY_SUFFIX = { "Darwin": ".dylib", "Linux": ".so", "Windows": ".dll", }.get(platform.system(), ".so") PACKAGE_DIR = Path(__file__).parent PACKAGE_GITHUB_URL = "https://github.com/TimDettmers/bitsandbytes" NONPYTORCH_DOC_URL = "https://github.com/TimDettmers/bitsandbytes/blob/main/docs/source/nonpytorchcuda.mdx" bitsandbytes/bitsandbytes/cuda_specs.py000066400000000000000000000055501506752665600207560ustar00rootroot00000000000000import dataclasses from functools import lru_cache import logging import re import subprocess from typing import Optional import torch @dataclasses.dataclass(frozen=True) class CUDASpecs: highest_compute_capability: tuple[int, int] cuda_version_string: str cuda_version_tuple: tuple[int, int] @property def has_imma(self) -> bool: return torch.version.hip or self.highest_compute_capability >= (7, 5) def get_compute_capabilities() -> list[tuple[int, int]]: return sorted(torch.cuda.get_device_capability(torch.cuda.device(i)) for i in range(torch.cuda.device_count())) @lru_cache(None) def get_cuda_version_tuple() -> Optional[tuple[int, int]]: """Get CUDA/HIP version as a tuple of (major, minor).""" try: if torch.version.cuda: version_str = torch.version.cuda elif torch.version.hip: version_str = torch.version.hip else: return None parts = version_str.split(".") if len(parts) >= 2: return tuple(map(int, parts[:2])) return None except (AttributeError, ValueError, IndexError): return None def get_cuda_version_string() -> Optional[str]: """Get CUDA/HIP version as a string.""" version_tuple = get_cuda_version_tuple() if version_tuple is None: return None major, minor = version_tuple return f"{major * 10 + minor}" def get_cuda_specs() -> Optional[CUDASpecs]: """Get CUDA/HIP specifications.""" if not torch.cuda.is_available(): return None try: compute_capabilities = get_compute_capabilities() if not compute_capabilities: return None version_tuple = get_cuda_version_tuple() if version_tuple is None: return None version_string = get_cuda_version_string() if version_string is None: return None return CUDASpecs( highest_compute_capability=compute_capabilities[-1], cuda_version_string=version_string, cuda_version_tuple=version_tuple, ) except Exception: return None def get_rocm_gpu_arch() -> str: """Get ROCm GPU architecture.""" logger = logging.getLogger(__name__) try: if torch.version.hip: result = subprocess.run(["rocminfo"], capture_output=True, text=True) match = re.search(r"Name:\s+gfx([a-zA-Z\d]+)", result.stdout) if match: return "gfx" + match.group(1) else: return "unknown" else: return "unknown" except Exception as e: logger.error(f"Could not detect ROCm GPU architecture: {e}") if torch.cuda.is_available(): logger.warning( """ ROCm GPU architecture detection failed despite ROCm being available. """, ) return "unknown" bitsandbytes/bitsandbytes/diagnostics/000077500000000000000000000000001506752665600205755ustar00rootroot00000000000000bitsandbytes/bitsandbytes/diagnostics/__init__.py000066400000000000000000000000001506752665600226740ustar00rootroot00000000000000bitsandbytes/bitsandbytes/diagnostics/cuda.py000066400000000000000000000176151506752665600220750ustar00rootroot00000000000000from collections.abc import Iterable, Iterator import logging import os from pathlib import Path import torch from bitsandbytes.cextension import HIP_ENVIRONMENT, get_cuda_bnb_library_path from bitsandbytes.cuda_specs import CUDASpecs from bitsandbytes.diagnostics.utils import print_dedented CUDART_PATH_PREFERRED_ENVVARS = ("CONDA_PREFIX", "LD_LIBRARY_PATH") CUDART_PATH_IGNORED_ENVVARS = { "DBUS_SESSION_BUS_ADDRESS", # hardware related "GOOGLE_VM_CONFIG_LOCK_FILE", # GCP: requires elevated permissions, causing problems in VMs and Jupyter notebooks "HOME", # Linux shell default "LESSCLOSE", "LESSOPEN", # related to the `less` command "MAIL", # something related to emails "OLDPWD", "PATH", # this is for finding binaries, not libraries "PWD", # PWD: this is how the shell keeps track of the current working dir "SHELL", # binary for currently invoked shell "SSH_AUTH_SOCK", # SSH stuff, therefore unrelated "SSH_TTY", "TMUX", # Terminal Multiplexer "XDG_DATA_DIRS", # XDG: Desktop environment stuff "XDG_GREETER_DATA_DIR", # XDG: Desktop environment stuff "XDG_RUNTIME_DIR", "_", # current Python interpreter } CUDA_RUNTIME_LIB_PATTERNS = ( ("libamdhip64.so*",) if HIP_ENVIRONMENT else ( "cudart64*.dll", # Windows "libcudart*.so*", # libcudart.so, libcudart.so.11.0, libcudart.so.12.0, libcudart.so.12.1, libcudart.so.12.2 etc. "nvcuda*.dll", # Windows ) ) logger = logging.getLogger(__name__) def find_cuda_libraries_in_path_list(paths_list_candidate: str) -> Iterable[Path]: for dir_string in paths_list_candidate.split(os.pathsep): if not dir_string: continue if os.sep not in dir_string: continue try: dir = Path(dir_string) try: if not dir.exists(): logger.warning(f"The directory listed in your path is found to be non-existent: {dir}") continue except OSError: # Assume an esoteric error trying to poke at the directory pass for lib_pattern in CUDA_RUNTIME_LIB_PATTERNS: for pth in dir.glob(lib_pattern): if pth.is_file() and not pth.is_symlink(): yield pth except (OSError, PermissionError): pass def is_relevant_candidate_env_var(env_var: str, value: str) -> bool: return ( env_var in CUDART_PATH_PREFERRED_ENVVARS # is a preferred location or ( os.sep in value # might contain a path and env_var not in CUDART_PATH_IGNORED_ENVVARS # not ignored and "CONDA" not in env_var # not another conda envvar and "BASH_FUNC" not in env_var # not a bash function defined via envvar and "\n" not in value # likely e.g. a script or something? ) ) def get_potentially_lib_path_containing_env_vars() -> dict[str, str]: return {env_var: value for env_var, value in os.environ.items() if is_relevant_candidate_env_var(env_var, value)} def find_cudart_libraries() -> Iterator[Path]: """ Searches for a cuda installations, in the following order of priority: 1. active conda env 2. LD_LIBRARY_PATH 3. any other env vars, while ignoring those that - are known to be unrelated - don't contain the path separator `/` If multiple libraries are found in part 3, we optimistically try one, while giving a warning message. """ candidate_env_vars = get_potentially_lib_path_containing_env_vars() for envvar in CUDART_PATH_PREFERRED_ENVVARS: if envvar in candidate_env_vars: directory = candidate_env_vars[envvar] yield from find_cuda_libraries_in_path_list(directory) candidate_env_vars.pop(envvar) for env_var, value in candidate_env_vars.items(): yield from find_cuda_libraries_in_path_list(value) def _print_cuda_diagnostics(cuda_specs: CUDASpecs) -> None: print( f"PyTorch settings found: CUDA_VERSION={cuda_specs.cuda_version_string}, " f"Highest Compute Capability: {cuda_specs.highest_compute_capability}.", ) binary_path = get_cuda_bnb_library_path(cuda_specs) if not binary_path.exists(): print_dedented( f""" Library not found: {binary_path}. Maybe you need to compile it from source? """, ) # 7.5 is the minimum CC for int8 tensor cores if not cuda_specs.has_imma: print_dedented( """ WARNING: Compute capability < 7.5 detected! Only slow 8-bit matmul is supported for your GPU! If you run into issues with 8-bit matmul, you can try 4-bit quantization: https://huggingface.co/blog/4bit-transformers-bitsandbytes """, ) def _print_hip_diagnostics(cuda_specs: CUDASpecs) -> None: print(f"PyTorch settings found: ROCM_VERSION={cuda_specs.cuda_version_string}") binary_path = get_cuda_bnb_library_path(cuda_specs) if not binary_path.exists(): print_dedented( f""" Library not found: {binary_path}. Maybe you need to compile it from source? If you compiled from source, check that ROCm version in PyTorch Settings matches your ROCm install. If not, reinstall PyTorch for your ROCm version and rebuild bitsandbytes. """, ) hip_major, hip_minor = cuda_specs.cuda_version_tuple if (hip_major, hip_minor) < (6, 1): print_dedented( """ WARNING: bitsandbytes is fully supported only from ROCm 6.1. """, ) def print_diagnostics(cuda_specs: CUDASpecs) -> None: if HIP_ENVIRONMENT: _print_hip_diagnostics(cuda_specs) else: _print_cuda_diagnostics(cuda_specs) def _print_cuda_runtime_diagnostics() -> None: cudart_paths = list(find_cudart_libraries()) if not cudart_paths: print("CUDA SETUP: WARNING! CUDA runtime files not found in any environmental path.") elif len(cudart_paths) > 1: print_dedented( f""" Found duplicate CUDA runtime files (see below). We select the PyTorch default CUDA runtime, which is {torch.version.cuda}, but this might mismatch with the CUDA version that is needed for bitsandbytes. To override this behavior set the `BNB_CUDA_VERSION=` environmental variable. For example, if you want to use the CUDA version 122, BNB_CUDA_VERSION=122 python ... OR set the environmental variable in your .bashrc: export BNB_CUDA_VERSION=122 In the case of a manual override, make sure you set LD_LIBRARY_PATH, e.g. export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-11.2, """, ) for pth in cudart_paths: print(f"* Found CUDA runtime at: {pth}") def _print_hip_runtime_diagnostics() -> None: cudart_paths = list(find_cudart_libraries()) if not cudart_paths: print("WARNING! ROCm runtime files not found in any environmental path.") elif len(cudart_paths) > 1: print_dedented( f""" Found duplicate ROCm runtime files (see below). We select the PyTorch default ROCm runtime, which is {torch.version.hip}, but this might mismatch with the ROCm version that is needed for bitsandbytes. To resolve it, install PyTorch built for the ROCm version you want to use and set LD_LIBRARY_PATH to your ROCm install path, e.g. export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/rocm-6.1.2/lib, """, ) for pth in cudart_paths: print(f"* Found ROCm runtime at: {pth}") def print_runtime_diagnostics() -> None: if HIP_ENVIRONMENT: _print_hip_runtime_diagnostics() else: _print_cuda_runtime_diagnostics() bitsandbytes/bitsandbytes/diagnostics/main.py000066400000000000000000000064231506752665600221000ustar00rootroot00000000000000import importlib import platform import sys import traceback import torch from bitsandbytes import __version__ as bnb_version from bitsandbytes.cextension import BNB_BACKEND from bitsandbytes.consts import PACKAGE_GITHUB_URL from bitsandbytes.cuda_specs import get_cuda_specs from bitsandbytes.diagnostics.cuda import ( print_diagnostics, ) from bitsandbytes.diagnostics.utils import print_dedented, print_header _RELATED_PACKAGES = [ "accelerate", "diffusers", "numpy", "pip", "peft", "safetensors", "transformers", "triton", "trl", ] def sanity_check(): from bitsandbytes.optim import Adam p = torch.nn.Parameter(torch.rand(10, 10).cuda()) a = torch.rand(10, 10).cuda() p1 = p.data.sum().item() adam = Adam([p]) out = a * p loss = out.sum() loss.backward() adam.step() p2 = p.data.sum().item() assert p1 != p2 def get_package_version(name: str) -> str: try: version = importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: version = "not found" return version def show_environment(): """Simple utility to print out environment information.""" print(f"Platform: {platform.platform()}") if platform.system() == "Linux": print(f" libc: {'-'.join(platform.libc_ver())}") print(f"Python: {platform.python_version()}") print(f"PyTorch: {torch.__version__}") print(f" CUDA: {torch.version.cuda or 'N/A'}") print(f" HIP: {torch.version.hip or 'N/A'}") print(f" XPU: {getattr(torch.version, 'xpu', 'N/A') or 'N/A'}") print("Related packages:") for pkg in _RELATED_PACKAGES: version = get_package_version(pkg) print(f" {pkg}: {version}") def main(): print_header(f"bitsandbytes v{bnb_version}") show_environment() print_header("") cuda_specs = get_cuda_specs() if cuda_specs: print_diagnostics(cuda_specs) # TODO: There's a lot of noise in this; needs improvement. # print_cuda_runtime_diagnostics() if not torch.cuda.is_available(): print(f"PyTorch says {BNB_BACKEND} is not available. Possible reasons:") print(f"1. {BNB_BACKEND} driver not installed") print("2. Using a CPU-only PyTorch build") print("3. No GPU detected") else: print(f"Checking that the library is importable and {BNB_BACKEND} is callable...") try: sanity_check() print("SUCCESS!") return except RuntimeError as e: if "not available in CPU-only" in str(e): print( f"WARNING: {__package__} is currently running as CPU-only!\n" "Therefore, 8-bit optimizers and GPU quantization are unavailable.\n\n" f"If you think that this is so erroneously,\nplease report an issue!", ) else: raise e except Exception: traceback.print_exc() print_dedented( f""" Above we output some debug information. Please provide this info when creating an issue via {PACKAGE_GITHUB_URL}/issues/new/choose WARNING: Please be sure to sanitize sensitive info from the output before posting it. """, ) sys.exit(1) bitsandbytes/bitsandbytes/diagnostics/utils.py000066400000000000000000000004341506752665600223100ustar00rootroot00000000000000import textwrap HEADER_WIDTH = 60 def print_header(txt: str, width: int = HEADER_WIDTH, filler: str = "=") -> None: txt = f" {txt} " if txt else "" print(txt.center(width, filler)) def print_dedented(text): print("\n".join(textwrap.dedent(text).strip().split("\n"))) bitsandbytes/bitsandbytes/functional.py000066400000000000000000002174761506752665600210230ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from collections.abc import Iterable import ctypes as ct import itertools from math import prod from typing import Any, Optional, Union import numpy as np import torch from torch import Tensor from typing_extensions import deprecated from bitsandbytes.utils import pack_dict_to_tensor, unpack_tensor_to_dict from .cextension import HIP_ENVIRONMENT, lib name2qmap = {} """C FUNCTIONS FOR OPTIMIZERS""" str2optimizer8bit = { "adam": ( lib.cadam_static_8bit_grad_32, lib.cadam_static_8bit_grad_16, ), "momentum": ( lib.cmomentum_static_8bit_grad_32, lib.cmomentum_static_8bit_grad_16, ), "rmsprop": ( lib.crmsprop_static_8bit_grad_32, lib.crmsprop_static_8bit_grad_16, ), "lion": ( lib.clion_static_8bit_grad_32, lib.clion_static_8bit_grad_16, ), "lamb": ( lib.cadam_static_8bit_grad_32, lib.cadam_static_8bit_grad_16, ), "lars": ( lib.cmomentum_static_8bit_grad_32, lib.cmomentum_static_8bit_grad_16, ), } class GlobalPageManager: _instance = None def __init__(self): raise RuntimeError("Call get_instance() instead") def initialize(self): self.paged_tensors = [] @classmethod def get_instance(cls): if cls._instance is None: cls._instance = cls.__new__(cls) cls._instance.initialize() return cls._instance def prefetch_all(self, to_cpu=False): # assume the first added, will be the # ones that are used first, so swap them in last # in the case they are evicted again for t in self.paged_tensors[::-1]: prefetch_tensor(t, to_cpu) class CUBLAS_Context: _instance = None def __init__(self): raise RuntimeError("Call get_instance() instead") def initialize(self): self.context = {} @classmethod def get_instance(cls): if cls._instance is None: cls._instance = cls.__new__(cls) cls._instance.initialize() return cls._instance def get_context(self, device): if device.index not in self.context: prev_device = torch.cuda.current_device() torch.cuda.set_device(device) self.context[device.index] = ct.c_void_p(lib.get_context()) torch.cuda.set_device(prev_device) return self.context[device.index] class Cusparse_Context: _instance = None def __init__(self): raise RuntimeError("Call get_instance() instead") def initialize(self): self.context = ct.c_void_p(lib.get_cusparse()) @classmethod def get_instance(cls): if cls._instance is None: cls._instance = cls.__new__(cls) cls._instance.initialize() return cls._instance FIRST_CUDA_DEVICE = torch.device("cuda", index=0) # When multiple GPUs are present, we use a context manager to # switch to the correct device of a tensor before invoking our CUDA # kernels in the C++ library. However, when there's only one device # there is no need to incur the overhead of cudaGetDevice/cudaSetDevice. if torch.cuda.device_count() > 1: def _cuda_device_of(a: torch.Tensor): return torch.cuda.device_of(a) else: import contextlib def _cuda_device_of(a: torch.Tensor): return contextlib.nullcontext() def get_paged(*shape, dtype=torch.float32, device=FIRST_CUDA_DEVICE): num_bytes = dtype.itemsize * prod(shape) cuda_ptr = lib.cget_managed_ptr(ct.c_size_t(num_bytes)) c_ptr = ct.cast(cuda_ptr, ct.POINTER(ct.c_int)) new_array = np.ctypeslib.as_array(c_ptr, shape=shape) out = torch.frombuffer(new_array, dtype=dtype, count=prod(shape)).view(shape) out.is_paged = True out.page_deviceid = device.index return out def prefetch_tensor(A: torch.Tensor, to_cpu=False): assert A.is_paged, "Only paged tensors can be prefetched!" if to_cpu: deviceid = -1 else: deviceid = A.page_deviceid lib.cprefetch(get_ptr(A), ct.c_size_t(A.nbytes), ct.c_int32(deviceid)) def elementwise_func(func_name, A, B, value, prefetch=True): func = None if A.dtype == torch.float32: func = getattr(lib, f"c{func_name}_fp32", None) cvalue = ct.c_float(value) elif A.dtype == torch.uint8: func = getattr(lib, f"c{func_name}_uint8", None) cvalue = ct.c_uint8(value) if func is None: raise NotImplementedError(f"Function not implemented: {func_name}") is_managed = getattr(A, "is_managed", False) if is_managed and prefetch: prefetch_tensor(A) if B is not None: prefetch_tensor(B) func(get_ptr(A), get_ptr(B), cvalue, ct.c_int64(A.numel())) if A.is_paged or B.is_paged: # paged function are fully asynchronous # if we return from this function, we want to the tensor # to be in the correct state, that is the final state after the # operation occurred. So we synchronize. torch.cuda.synchronize() def fill(A, value, device=None, prefetch=True): elementwise_func("fill", A, None, value) def _mul(A, B, device=None): elementwise_func("_mul", A, B, 0) def create_linear_map(signed=True, total_bits=8, add_zero=True): sign = -1.0 if signed else 0.0 total_values = 2**total_bits if add_zero or total_bits < 8: # add a zero # since we simulate less bits by having zeros in the data type, we # we need to center the quantization around zero and as such lose # a single value total_values = 2**total_bits if not signed else 2**total_bits - 1 values = torch.linspace(sign, 1.0, total_values) gap = 256 - values.numel() if gap == 0: return values else: l = values.numel() // 2 # noqa: E741 return torch.Tensor(values[:l].tolist() + [0] * gap + values[l:].tolist()) def create_normal_map(offset=0.9677083, use_extra_value=True): try: from scipy.stats import norm except ImportError as ie: raise ImportError( "Scipy is required for `create_normal_map`. Install `bitsandbytes` with the `[test]` extra.", ) from ie if use_extra_value: # one more positive value, this is an asymmetric type v1 = norm.ppf(torch.linspace(offset, 0.5, 9)[:-1]).tolist() v2 = [0] * (256 - 15) ## we have 15 non-zero values in this data type v3 = (-norm.ppf(torch.linspace(offset, 0.5, 8)[:-1])).tolist() else: v1 = norm.ppf(torch.linspace(offset, 0.5, 8)[:-1]).tolist() v2 = [0] * (256 - 14) ## we have 14 non-zero values in this data type v3 = (-norm.ppf(torch.linspace(offset, 0.5, 8)[:-1])).tolist() v = v1 + v2 + v3 values = torch.Tensor(v) values = values.sort().values values /= values.max() assert values.numel() == 256 return values def create_fp8_map(signed=True, exponent_bits=5, precision_bits=2, total_bits=8): e = exponent_bits p = precision_bits has_sign = 1 if signed else 0 assert e + p == total_bits - has_sign # the exponent is biased to 2^(e-1) -1 == 0 evalues = [] for i, val in enumerate(range(-(2 ** (exponent_bits - has_sign)), 2 ** (exponent_bits - has_sign), 1)): evalues.append(2**val) values = [] lst = list(itertools.product([0, 1], repeat=precision_bits)) # for ev in evalues: bias = 2 ** (exponent_bits - 1) for evalue in range(2 ** (exponent_bits)): for bit_pattern in lst: value = 1 if evalue != 0 else 0 for i, pval in enumerate(list(bit_pattern)): value += pval * (2 ** -(i + 1)) if evalue == 0: # subnormals value = value * 2**-(bias) else: # normals value = value * 2 ** -(evalue - bias - 1) values.append(value) if signed: values.append(-value) assert len(values) == 2**total_bits values.sort() if total_bits < 8: gap = 256 - len(values) for i in range(gap): values.append(0) values.sort() code = torch.tensor(values) code /= code.max() return code def create_dynamic_map(signed=True, max_exponent_bits=7, total_bits=8): """ Creates the dynamic quantiztion map. The dynamic data type is made up of a dynamic exponent and fraction. As the exponent increase from 0 to -7 the number of bits available for the fraction shrinks. This is a generalization of the dynamic type where a certain number of the bits and be reserved for the linear quantization region (the fraction). n determines the maximum number of exponent bits. For more details see (8-Bit Approximations for Parallelism in Deep Learning)[https://arxiv.org/abs/1511.04561] """ data = [] # these are additional items that come from the case # where all the exponent bits are zero and no # indicator bit is present non_sign_bits = total_bits - 1 additional_items = 2 ** (non_sign_bits - max_exponent_bits) - 1 for i in range(max_exponent_bits): fraction_items = int( 2 ** (i + non_sign_bits - max_exponent_bits) + 1 if signed else 2 ** (i + non_sign_bits - max_exponent_bits + 1) + 1, ) boundaries = torch.linspace(0.1, 1, fraction_items, dtype=torch.float32) means = (boundaries[:-1] + boundaries[1:]) / 2.0 data += ((10 ** (-(max_exponent_bits - 1) + i)) * means).tolist() if signed: data += (-(10 ** (-(max_exponent_bits - 1) + i)) * means).tolist() if additional_items > 0: boundaries = torch.linspace(0.1, 1, additional_items + 1, dtype=torch.float32) means = (boundaries[:-1] + boundaries[1:]) / 2.0 data += ((10 ** (-(max_exponent_bits - 1) + i)) * means).tolist() if signed: data += (-(10 ** (-(max_exponent_bits - 1) + i)) * means).tolist() data.append(0) data.append(1.0) assert len(data) == 2**total_bits gap = 256 - len(data) for i in range(gap): data.append(0) data.sort() return torch.tensor(data, dtype=torch.float32) def is_on_gpu(tensors: Iterable[Optional[torch.Tensor]]): """Verifies that the input tensors are all on the same device. An input tensor may also be marked as `paged`, in which case the device placement is ignored. Args: tensors (`Iterable[Optional[torch.Tensor]]`): A list of tensors to verify. Raises: `RuntimeError`: Raised when the verification fails. Returns: `Literal[True]` """ on_gpu = True gpu_ids = set() for t in tensors: # NULL pointers and paged tensors are OK. if t is not None and not getattr(t, "is_paged", False): on_gpu &= t.device.type != "cpu" gpu_ids.add((t.device.type, t.device.index)) if not on_gpu: raise RuntimeError( f"All input tensors need to be on the same GPU, but found some tensors to not be on a GPU:\n {[(t.shape, t.device) for t in tensors]}", ) if len(gpu_ids) > 1: raise RuntimeError( f"Input tensors need to be on the same GPU, but found the following tensor and device combinations:\n {[(t.shape, t.device) for t in tensors]}", ) return on_gpu def _get_tensor_stream(tensor: Tensor) -> ct.c_void_p: # We use the raw stream for performance reasons. if tensor.device.type == "xpu": return ct.c_void_p(torch._C._xpu_getCurrentRawStream(tensor.device.index)) return ct.c_void_p(torch._C._cuda_getCurrentRawStream(tensor.device.index)) def get_ptr(A: Optional[Tensor]) -> Optional[ct.c_void_p]: """Gets the memory address of the first element of a tenso Args: A (`Optional[Tensor]`): A PyTorch tensor. Returns: `Optional[ct.c_void_p]`: A pointer to the underlying tensor data. """ if A is None: return None return ct.c_void_p(A.data_ptr()) class QuantState: """container for quantization state components to work with Params4bit and similar classes""" valid_quant_types = ("fp4", "nf4") valid_qs_type_keys = [f"bitsandbytes__{x}" for x in valid_quant_types] valid_qs_keys = [ "absmax", "quant_map", "nested_absmax", "nested_quant_map", "quant_state", "quant_type", "blocksize", "dtype", "shape", "nested_blocksize", "nested_dtype", "nested_offset", ] def __init__( self, absmax, shape=None, code=None, blocksize=None, quant_type=None, dtype=None, offset=None, state2=None, ): self.absmax = absmax self.shape = shape self.code = code self.dtype = dtype self.blocksize = blocksize self.quant_type = quant_type self.offset = offset self.state2 = state2 self.nested = state2 is not None def __getitem__(self, idx): """ ensures compatibility with older quant state scheme with nested lists. assumes the following layout: state = [qabsmax, input_shape, A.dtype, blocksize, [offset, state2], quant_type] state2 = [absmax, input_shape, A.dtype, blocksize, None, quant_type] """ if self.nested: list_repr = [ self.absmax, self.shape, self.dtype, self.blocksize, [self.offset, self.state2], self.quant_type, ] else: list_repr = [self.absmax, self.shape, self.dtype, self.blocksize, None, self.quant_type] return list_repr[idx] @classmethod def from_dict(cls, qs_dict: dict[str, Any], device: torch.device) -> "QuantState": """ unpacks components of state_dict into QuantState where necessary, convert into strings, torch.dtype, ints, etc. qs_dict: based on state_dict, with only relevant keys, striped of prefixes. item with key `quant_state.bitsandbytes__[nf4/fp4]` may contain minor and non-tensor quant state items. """ # unpacking tensor with non-tensor components qs_key = [k for k, v in qs_dict.items() if "quant_state" in k and isinstance(v, torch.Tensor)] if not len(qs_key) and "quant_type" not in qs_dict: raise ValueError("Expected packed or unpacked quant_state items, found neither") elif len(qs_key) != 1 or qs_key[0].split(".")[-1] not in cls.valid_qs_type_keys: raise ValueError( f"There should be exactly one `quant_state` item with ending from {cls.valid_qs_type_keys}.\nDetected {qs_key}.", ) # unpacking minor and non-tensor quant state items if necessary if len(qs_key) == 1: first_qs_key = qs_key[0] qs_dict.update(unpack_tensor_to_dict(qs_dict.pop(first_qs_key))) qs_dict = {k.split(".")[-1]: v for k, v in qs_dict.items()} # strip prefixes assert set(qs_dict.keys()).issubset(cls.valid_qs_keys) if "nested_absmax" in qs_dict: offset = torch.tensor(float(qs_dict["nested_offset"])).to(device) state2 = cls( absmax=qs_dict["nested_absmax"].to(device), blocksize=qs_dict["nested_blocksize"], code=qs_dict["nested_quant_map"].to(device), dtype=getattr(torch, qs_dict["nested_dtype"]), ) else: offset, state2 = None, None quant_state = cls( quant_type=qs_dict["quant_type"], absmax=qs_dict["absmax"].to(device), blocksize=qs_dict["blocksize"], code=qs_dict["quant_map"].to(device), dtype=getattr(torch, qs_dict["dtype"]), shape=torch.Size(qs_dict["shape"]) if qs_dict["shape"] is not None else None, offset=offset, state2=state2, ) return quant_state def as_dict(self, packed=False): """ returns dict of tensors and strings to use in serialization via _save_to_state_dict() param: packed -- returns dict[str, torch.Tensor] for state_dict fit for safetensors saving """ qs_dict = { "quant_type": self.quant_type, "absmax": self.absmax, "blocksize": self.blocksize, "quant_map": self.code, "dtype": str(self.dtype).strip("torch."), "shape": tuple(self.shape), } if self.nested: qs_dict.update( { "nested_absmax": self.state2.absmax, "nested_blocksize": self.state2.blocksize, "nested_quant_map": self.state2.code.clone(), # un-shared to avoid restoring it after shared tensors are removed by safetensors "nested_dtype": str(self.state2.dtype).strip("torch."), "nested_offset": self.offset.item(), }, ) if not packed: return qs_dict # packed format allows serialization of non-tensor components, critical for saving in safetensors format qs_packed_dict = {k: v for k, v in qs_dict.items() if isinstance(v, torch.Tensor)} non_tensor_dict = {k: v for k, v in qs_dict.items() if not isinstance(v, torch.Tensor)} qs_packed_dict["quant_state." + "bitsandbytes__" + self.quant_type] = pack_dict_to_tensor(non_tensor_dict) return qs_packed_dict def to(self, device): # make sure the quantization state is on the right device self.code = self.code.to(device) self.absmax = self.absmax.to(device) if self.nested: self.offset = self.offset.to(device) self.state2.absmax = self.state2.absmax.to(device) self.state2.code = self.state2.code.to(device) def __eq__(self, other): if not isinstance(other, QuantState): return False return ( torch.allclose(self.absmax, other.absmax, atol=1e-6) and self.shape == other.shape and torch.allclose(self.code, other.code, atol=1e-6) and self.dtype == other.dtype and self.blocksize == other.blocksize and self.quant_type == other.quant_type and ( self.offset == other.offset if self.offset is not None and other.offset is not None else self.offset is other.offset ) and ( self.state2 == other.state2 if self.state2 is not None and other.state2 is not None else self.state2 is other.state2 ) ) def quantize_blockwise( A: torch.Tensor, code: Optional[torch.Tensor] = None, absmax: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, blocksize=4096, nested=False, ) -> tuple[torch.Tensor, QuantState]: """Quantize a tensor in blocks of values. The input tensor is quantized by dividing it into blocks of `blocksize` values. The the absolute maximum value within these blocks is calculated for scaling the non-linear quantization. Args: A (`torch.Tensor`): The input tensor. Supports `float16`, `bfloat16`, or `float32` datatypes. code (`torch.Tensor`, *optional*): A mapping describing the low-bit data type. Defaults to a signed 8-bit dynamic type. For more details, see (8-Bit Approximations for Parallelism in Deep Learning)[https://arxiv.org/abs/1511.04561]. absmax (`torch.Tensor`, *optional*): A tensor to use to store the absmax values. out (`torch.Tensor`, *optional*): A tensor to use to store the result. blocksize (`int`, *optional*): The size of the blocks. Defaults to 4096. Valid values are 64, 128, 256, 512, 1024, 2048, and 4096. nested (`bool`, *optional*): Whether to additionally quantize the absmax values. Defaults to False. Raises: ValueError: Raised when the input data type is not supported. Returns: `Tuple[torch.Tensor, QuantState]`: A tuple containing the quantization results. - `torch.Tensor`: The quantized tensor. - [`QuantState`]: The state object used to undo the quantization. """ if code is None: if "dynamic" not in name2qmap: name2qmap["dynamic"] = create_dynamic_map().to(A.device) code = name2qmap["dynamic"] _out, _absmax = torch.ops.bitsandbytes.quantize_blockwise.default( A, code.to(A.device), blocksize, ) if nested: offset = _absmax.mean() _absmax -= offset qabsmax, state2 = quantize_blockwise(_absmax, blocksize=blocksize, nested=False) quant_state = QuantState( absmax=qabsmax, code=code.to(A.device, copy=True), blocksize=blocksize, dtype=A.dtype, offset=offset, state2=state2, ) else: quant_state = QuantState(absmax=_absmax, code=code.to(A.device, copy=True), blocksize=blocksize, dtype=A.dtype) # TODO(matthewdouglas): Deprecate out kwarg out = out.copy_(_out) if out is not None else _out # TODO(matthewdouglas): Deprecate absmax kwarg if absmax is not None: quant_state.absmax = absmax.copy_(quant_state.absmax) return out, quant_state def dequantize_blockwise( A: torch.Tensor, quant_state: Optional[QuantState] = None, absmax: Optional[torch.Tensor] = None, code: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, blocksize: int = 4096, nested=False, ) -> torch.Tensor: """Dequantize a tensor in blocks of values. The input tensor is dequantized by dividing it into blocks of `blocksize` values. The the absolute maximum value within these blocks is used for scaling the non-linear dequantization. Args: A (`torch.Tensor`): The quantized input tensor. quant_state ([`QuantState`], *optional*): The quantization state as returned by [`quantize_blockwise`]. Required if `absmax` is not provided. absmax (`torch.Tensor`, *optional*): A tensor containing the scaling values. Required if `quant_state` is not provided and ignored otherwise. code (`torch.Tensor`, *optional*): A mapping describing the low-bit data type. Defaults to a signed 8-bit dynamic type. For more details, see (8-Bit Approximations for Parallelism in Deep Learning)[https://arxiv.org/abs/1511.04561]. Ignored when `quant_state` is provided. out (`torch.Tensor`, *optional*): A tensor to use to store the result. blocksize (`int`, *optional*): The size of the blocks. Defaults to 4096. Valid values are 64, 128, 256, 512, 1024, 2048, and 4096. Ignored when `quant_state` is provided. Raises: ValueError: Raised when the input data type is not supported. Returns: `torch.Tensor`: The dequantized tensor. The datatype is indicated by `quant_state.dtype` and defaults to `torch.float32`. """ assert quant_state is not None or absmax is not None if code is None and quant_state is None: if "dynamic" not in name2qmap: name2qmap["dynamic"] = create_dynamic_map().to(A.device) code = name2qmap["dynamic"] if quant_state is None: quant_state = QuantState(absmax=absmax, code=code, blocksize=blocksize, dtype=torch.float32) absmax = quant_state.absmax if quant_state.nested: absmax = dequantize_blockwise(quant_state.absmax, quant_state.state2) absmax += quant_state.offset if absmax.dtype != torch.float32: absmax = absmax.float() if out is not None: torch.ops.bitsandbytes.dequantize_blockwise.out( A, absmax, quant_state.code.to(A.device), quant_state.blocksize, quant_state.dtype, out=out, ) return out return torch.ops.bitsandbytes.dequantize_blockwise.default( A, absmax, quant_state.code.to(A.device), quant_state.blocksize, quant_state.dtype, ) def get_4bit_type(typename, device=None, blocksize=64): if device is None: device = "cuda" data = None if typename == "nf4": """ Implements the NF4 data type. Constructs a quantization data type where each bin has equal area under a standard normal distribution N(0, 1) that is normalized into the range [-1, 1]. For more information read the paper: QLoRA: Efficient Finetuning of Quantized LLMs (https://arxiv.org/abs/2305.14314) Implementation of the NF4 data type in bitsandbytes can be found in the `create_normal_map` function in the `functional.py` file: https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L236. """ data = [ -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453, -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0, 0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224, 0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0, ] elif typename == "fp4": # 0b000 = 0 # 0b001 = 0.0625 # 0b010 = 8 # 0b011 = 12 # 0b100 = 4 # 0b101 = 6 # 0b110 = 2 # 0b111 = 3 # can also be created with bnb.functional.create_fp8_map(signed=True, exponent_bits=2, precision_bits=1, total_bits=4) data = [0, 0.0625, 8.0, 12.0, 4.0, 6.0, 2.0, 3.0, -0, -0.0625, -8.0, -12.0, -4.0, -6.0, -2.0, -3.0] elif typename == "int4": data = [7, 6, 5, 4, 3, 2, 1, 0, -0, -1, -2, -3, -4, -5, -6, -7] elif typename == "af4": # Taken from: NF4 Isn't Information Theoretically Optimal (and that's Good) # https://arxiv.org/abs/2306.06965 if blocksize == 64: data = [ -1.0, -0.69441008, -0.51243739, -0.3736951, -0.25607552, -0.14982478, -0.04934812, 0.0, 0.04273164, 0.12934483, 0.21961274, 0.31675666, 0.42563882, 0.55496234, 0.72424863, 1.0, ][::-1] else: raise NotImplementedError("4-bit AbnormalFloats currently only support blocksize 64.") if data is None: raise NotImplementedError(f"Typename {typename} not supported") data = torch.tensor(data, device=device) data.div_(data.abs().max()) assert data.numel() == 16 return data def quantize_fp4( A: torch.Tensor, absmax: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, blocksize=None, compress_statistics=False, quant_storage=torch.uint8, ): if blocksize is None: blocksize = 64 if not HIP_ENVIRONMENT else 128 return quantize_4bit(A, absmax, out, blocksize, compress_statistics, "fp4", quant_storage) def quantize_nf4( A: torch.Tensor, absmax: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, blocksize=None, compress_statistics=False, quant_storage=torch.uint8, ): if blocksize is None: blocksize = 64 if not HIP_ENVIRONMENT else 128 return quantize_4bit(A, absmax, out, blocksize, compress_statistics, "nf4", quant_storage) def quantize_4bit( A: torch.Tensor, absmax: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, blocksize=None, compress_statistics=False, quant_type="fp4", quant_storage=torch.uint8, ) -> tuple[torch.Tensor, QuantState]: """Quantize tensor A in blocks of 4-bit values. Quantizes tensor A by dividing it into blocks which are independently quantized. Args: A (`torch.Tensor`): The input tensor. Supports `float16`, `bfloat16`, or `float32` datatypes. absmax (`torch.Tensor`, *optional*): A tensor to use to store the absmax values. out (`torch.Tensor`, *optional*): A tensor to use to store the result. blocksize (`int`, *optional*): The size of the blocks. Defaults to 128 on ROCm and 64 otherwise. Valid values are 64, 128, 256, 512, 1024, 2048, and 4096. compress_statistics (`bool`, *optional*): Whether to additionally quantize the absmax values. Defaults to False. quant_type (`str`, *optional*): The data type to use: `nf4` or `fp4`. Defaults to `fp4`. quant_storage (`torch.dtype`, *optional*): The dtype of the tensor used to store the result. Defaults to `torch.uint8`. Raises: ValueError: Raised when the input data type is not supported. Returns: Tuple[`torch.Tensor`, `QuantState`]: A tuple containing the quantization results. - `torch.Tensor`: The quantized tensor with packed 4-bit values. - [`QuantState`]: The state object used to undo the quantization. """ if blocksize is None: blocksize = 64 if not HIP_ENVIRONMENT else 128 input_shape = A.shape _out, _absmax = torch.ops.bitsandbytes.quantize_4bit.default( A, blocksize, quant_type, quant_storage, ) code = get_4bit_type(quant_type, device=A.device) if compress_statistics: offset = _absmax.mean() qabsmax, state2 = quantize_blockwise(_absmax - offset, blocksize=256) del _absmax state = QuantState( absmax=qabsmax, shape=input_shape, dtype=A.dtype, blocksize=blocksize, code=code, quant_type=quant_type, offset=offset, state2=state2, ) else: state = QuantState( absmax=_absmax, shape=input_shape, dtype=A.dtype, blocksize=blocksize, code=code, quant_type=quant_type, ) # TODO(matthewdouglas): Deprecate out kwarg out = out.copy_(_out) if out is not None else _out # TODO(matthewdouglas): Deprecate absmax kwarg if absmax is not None: state.absmax = absmax.copy_(state.absmax) return out, state def dequantize_fp4( A: torch.Tensor, quant_state: Optional[QuantState] = None, absmax: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, blocksize: Optional[int] = None, ) -> torch.Tensor: if blocksize is None: blocksize = 64 if not HIP_ENVIRONMENT else 128 return dequantize_4bit(A, quant_state, absmax, out, blocksize, "fp4") def dequantize_nf4( A: torch.Tensor, quant_state: Optional[QuantState] = None, absmax: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, blocksize: Optional[int] = None, ) -> torch.Tensor: if blocksize is None: blocksize = 64 if not HIP_ENVIRONMENT else 128 return dequantize_4bit(A, quant_state, absmax, out, blocksize, "nf4") def dequantize_4bit( A: torch.Tensor, quant_state: Optional[QuantState] = None, absmax: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, blocksize: Optional[int] = None, quant_type="fp4", ) -> torch.Tensor: """Dequantizes a packed 4-bit quantized tensor. The input tensor is dequantized by dividing it into blocks of `blocksize` values. The the absolute maximum value within these blocks is used for scaling the non-linear dequantization. Args: A (`torch.Tensor`): The quantized input tensor. quant_state ([`QuantState`], *optional*): The quantization state as returned by [`quantize_4bit`]. Required if `absmax` is not provided. absmax (`torch.Tensor`, *optional*): A tensor containing the scaling values. Required if `quant_state` is not provided and ignored otherwise. out (`torch.Tensor`, *optional*): A tensor to use to store the result. blocksize (`int`, *optional*): The size of the blocks. Defaults to 128 on ROCm and 64 otherwise. Valid values are 64, 128, 256, 512, 1024, 2048, and 4096. quant_type (`str`, *optional*): The data type to use: `nf4` or `fp4`. Defaults to `fp4`. Raises: ValueError: Raised when the input data type or blocksize is not supported. Returns: `torch.Tensor`: The dequantized tensor. """ if blocksize is None: blocksize = 64 if not HIP_ENVIRONMENT else 128 if quant_state is None: assert absmax is not None and out is not None quant_state = QuantState( absmax=absmax, shape=out.shape, dtype=out.dtype, blocksize=blocksize, quant_type=quant_type, ) else: absmax = quant_state.absmax if quant_state.nested: absmax = dequantize_blockwise(quant_state.absmax, quant_state.state2) absmax += quant_state.offset if absmax.dtype != torch.float32: absmax = absmax.float() if out is not None: torch.ops.bitsandbytes.dequantize_4bit.out( A, absmax, quant_state.blocksize, quant_state.quant_type, quant_state.shape, quant_state.dtype, out=out ) else: out = torch.ops.bitsandbytes.dequantize_4bit.default( A, absmax, quant_state.blocksize, quant_state.quant_type, quant_state.shape, quant_state.dtype, ) if A.shape[0] == 1: # is transposed, transpose back return out.t() return out @deprecated("This function is deprecated and will be removed in a future release.", category=FutureWarning) def quantize( A: Tensor, code: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, ) -> tuple[Tensor, tuple[Tensor, Tensor]]: if code is None: if "dynamic" not in name2qmap: name2qmap["dynamic"] = create_dynamic_map().to(A.device) code = name2qmap["dynamic"] code = code.to(A.device) absmax = torch.abs(A).max() if absmax.dtype != torch.float32: absmax = absmax.float() inp = A / absmax out = quantize_no_absmax(inp, code, out) return out, (absmax, code) @deprecated("This function is deprecated and will be removed in a future release.", category=FutureWarning) def dequantize( A: Tensor, state: Optional[tuple[Tensor, Tensor]] = None, absmax: Optional[torch.Tensor] = None, code: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, ) -> Tensor: assert state is not None or absmax is not None if code is None and state is None: if "dynamic" not in name2qmap: name2qmap["dynamic"] = create_dynamic_map().to(A.device) code = name2qmap["dynamic"] code = code.to(A.device) if state is None: state = (absmax, code) out = dequantize_no_absmax(A, state[1], out) return out * state[0] @deprecated("This function is deprecated and will be removed in a future release.", category=FutureWarning) def quantize_no_absmax(A: Tensor, code: Tensor, out: Optional[torch.Tensor] = None) -> Tensor: """ Quantizes input tensor to 8-bit. Quantizes the 32-bit input tensor `A` to the 8-bit output tensor `out` using the quantization map `code`. Parameters ---------- A : torch.Tensor The input tensor. code : torch.Tensor The quantization map. out : torch.Tensor, optional The output tensor. Needs to be of type byte. Returns ------- torch.Tensor: Quantized 8-bit tensor. """ with _cuda_device_of(A): if out is None: out = torch.zeros_like(A, dtype=torch.uint8) is_on_gpu([A, out]) lib.cquantize(get_ptr(code), get_ptr(A), get_ptr(out), ct.c_int(A.numel())) return out @deprecated("This function is deprecated and will be removed in a future release.", category=FutureWarning) def dequantize_no_absmax(A: Tensor, code: Tensor, out: Optional[torch.Tensor] = None) -> Tensor: """ Dequantizes the 8-bit tensor to 32-bit. Dequantizes the 8-bit tensor `A` to the 32-bit tensor `out` via the quantization map `code`. Parameters ---------- A : torch.Tensor The 8-bit input tensor. code : torch.Tensor The quantization map. out : torch.Tensor The 32-bit output tensor. Returns ------- torch.Tensor: 32-bit output tensor. """ with _cuda_device_of(A): if out is None: out = torch.zeros_like(A, dtype=torch.float32) is_on_gpu([code, A, out]) stream = _get_tensor_stream(A) lib.cdequantize(get_ptr(code), get_ptr(A), get_ptr(out), ct.c_int(A.numel()), stream) return out def optimizer_update_32bit( optimizer_name: str, g: Tensor, p: Tensor, state1: Tensor, beta1: float, eps: float, step: int, lr: float, state2: Optional[torch.Tensor] = None, beta2: float = 0.0, beta3: float = 0.0, alpha: float = 0.0, weight_decay: float = 0.0, gnorm_scale: float = 1.0, unorm_vec: Optional[torch.Tensor] = None, max_unorm: float = 0.0, skip_zeros=False, ) -> None: """ Performs an inplace optimizer update with one or two optimizer states. Universal optimizer update for 32-bit state and 32/16-bit gradients/weights. Parameters ---------- optimizer_name : str The name of the optimizer: {adam}. g : torch.Tensor Gradient tensor. p : torch.Tensor Parameter tensor. state1 : torch.Tensor Optimizer state 1. beta1 : float Optimizer beta1. eps : float Optimizer epsilon. weight_decay : float Weight decay. step : int Current optimizer step. lr : float The learning rate. state2 : torch.Tensor Optimizer state 2. beta2 : float Optimizer beta2. beta3 : float Optimizer beta3. alpha : float Optimizer alpha. gnorm_scale : float The factor to rescale the gradient to the max clip value. unorm_vec : torch.Tensor The tensor for the update norm. max_unorm : float The maximum update norm relative to the weight norm. skip_zeros : bool Whether to skip zero-valued gradients or not (default: False). """ param_norm = 0.0 if max_unorm > 0.0: param_norm = torch.norm(p.data.float()) is_on_gpu([g, p, state1, state2, unorm_vec]) torch.ops.bitsandbytes.optimizer_update_32bit( optimizer_name, g, p, state1, state2, unorm_vec, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, lr, gnorm_scale, skip_zeros, ) @deprecated( "This function is deprecated and will be removed in a future release. " "Please use optimizer_update_8bit_blockwise instead. ", category=FutureWarning, ) def optimizer_update_8bit( optimizer_name: str, g: Tensor, p: Tensor, state1: Tensor, state2: Optional[torch.Tensor], beta1: float, beta2: float, eps: float, step: int, lr: float, qmap1: Tensor, qmap2: Optional[torch.Tensor], max1: Tensor, max2: Optional[torch.Tensor], new_max1: Tensor, new_max2: Optional[torch.Tensor], weight_decay: float = 0.0, gnorm_scale: float = 1.0, unorm_vec: Optional[torch.Tensor] = None, max_unorm: float = 0.0, ) -> None: """ Performs an inplace Adam update. Universal Adam update for 32/8-bit state and 32/16-bit gradients/weights. Uses AdamW formulation if weight decay > 0.0. Parameters ---------- optimizer_name : str The name of the optimizer. Choices {adam, momentum} g : torch.Tensor Gradient tensor. p : torch.Tensor Parameter tensor. state1 : torch.Tensor Adam state 1. state2 : torch.Tensor Adam state 2. beta1 : float Adam beta1. beta2 : float Adam beta2. eps : float Adam epsilon. weight_decay : float Weight decay. step : int Current optimizer step. lr : float The learning rate. qmap1 : torch.Tensor Quantization map for first Adam state. qmap2 : torch.Tensor Quantization map for second Adam state. max1 : torch.Tensor Max value for first Adam state update. max2 : torch.Tensor Max value for second Adam state update. new_max1 : torch.Tensor Max value for the next Adam update of the first state. new_max2 : torch.Tensor Max value for the next Adam update of the second state. gnorm_scale : float The factor to rescale the gradient to the max clip value. unorm_vec : torch.Tensor The tensor for the update norm. max_unorm : float The maximum update norm relative to the weight norm. """ param_norm = 0.0 if max_unorm > 0.0: param_norm = torch.norm(p.data.float()) with _cuda_device_of(g): is_on_gpu([g, p, state1, state2, unorm_vec, qmap1, qmap2, max1, max2, new_max1, new_max2]) if g.dtype == torch.float32 and state1.dtype == torch.uint8: str2optimizer8bit[optimizer_name][0]( get_ptr(p), get_ptr(g), get_ptr(state1), get_ptr(state2), get_ptr(unorm_vec), ct.c_float(max_unorm), ct.c_float(param_norm), ct.c_float(beta1), ct.c_float(beta2), ct.c_float(eps), ct.c_int32(step), ct.c_float(lr), get_ptr(qmap1), get_ptr(qmap2), get_ptr(max1), get_ptr(max2), get_ptr(new_max1), get_ptr(new_max2), ct.c_float(weight_decay), ct.c_float(gnorm_scale), ct.c_int32(g.numel()), ) elif g.dtype == torch.float16 and state1.dtype == torch.uint8: str2optimizer8bit[optimizer_name][1]( get_ptr(p), get_ptr(g), get_ptr(state1), get_ptr(state2), get_ptr(unorm_vec), ct.c_float(max_unorm), ct.c_float(param_norm), ct.c_float(beta1), ct.c_float(beta2), ct.c_float(eps), ct.c_int32(step), ct.c_float(lr), get_ptr(qmap1), get_ptr(qmap2), get_ptr(max1), get_ptr(max2), get_ptr(new_max1), get_ptr(new_max2), ct.c_float(weight_decay), ct.c_float(gnorm_scale), ct.c_int32(g.numel()), ) else: raise ValueError( f"Gradient+optimizer bit data type combination not supported: grad {g.dtype}, optimizer {state1.dtype}", ) def optimizer_update_8bit_blockwise( optimizer_name: str, g: Tensor, p: Tensor, state1: Tensor, state2: Optional[torch.Tensor], beta1: float, beta2: float, beta3: float, alpha: float, eps: float, step: int, lr: float, qmap1: Tensor, qmap2: Optional[torch.Tensor], absmax1: Tensor, absmax2: Optional[torch.Tensor], weight_decay: float = 0.0, gnorm_scale: float = 1.0, skip_zeros=False, ) -> None: is_on_gpu([p, g, state1, state2, qmap1, qmap2, absmax1, absmax2]) torch.ops.bitsandbytes.optimizer_update_8bit_blockwise( optimizer_name, g, p, state1, state2, beta1, beta2, beta3, alpha, eps, step, lr, qmap1, qmap2, absmax1, absmax2, weight_decay, gnorm_scale, skip_zeros, ) @deprecated("This function is deprecated and will be removed in a future release.", category=FutureWarning) def percentile_clipping(grad: Tensor, gnorm_vec: Tensor, step: int, percentile: int = 5): """Applies percentile clipping grad: torch.Tensor The gradient tensor. gnorm_vec: torch.Tensor Vector of gradient norms. 100 elements expected. step: int The current optimization steps (number of past gradient norms). """ with _cuda_device_of(grad): is_on_gpu([grad, gnorm_vec]) if grad.dtype == torch.float32: lib.cpercentile_clipping_g32( get_ptr(grad), get_ptr(gnorm_vec), ct.c_int32(step), ct.c_int32(grad.numel()), ) elif grad.dtype == torch.float16: lib.cpercentile_clipping_g16( get_ptr(grad), get_ptr(gnorm_vec), ct.c_int32(step), ct.c_int32(grad.numel()), ) else: raise ValueError(f"Gradient type {grad.dtype} not supported!") current_gnorm = torch.sqrt(gnorm_vec[step % 100]) vals, idx = torch.sort(gnorm_vec) clip_value = torch.sqrt(vals[percentile]) gnorm_scale = 1.0 if current_gnorm > clip_value: gnorm_scale = clip_value / current_gnorm return current_gnorm, clip_value, gnorm_scale def check_matmul(A, B, out, transposed_A, transposed_B, expected_type=torch.int8): if not torch.cuda.is_initialized(): torch.cuda.init() if A.dtype != expected_type or B.dtype != expected_type: raise TypeError(f"Expected torch.int8 input tensors A and B, but got {A.dtype} and {B.dtype}") sA = A.shape sB = B.shape tA = transposed_A tB = transposed_B correct = True if len(sA) == 2 and len(sB) == 2: if not tA and not tB and A.shape[1] != B.shape[0]: correct = False elif tA and not tB and A.shape[0] != B.shape[0]: correct = False elif tA and tB and A.shape[0] != B.shape[1]: correct = False elif not tA and tB and A.shape[1] != B.shape[1]: correct = False elif len(sA) == 3 and len(sB) == 2: if not tA and not tB and A.shape[2] != B.shape[0]: correct = False elif tA and not tB and A.shape[1] != B.shape[0]: correct = False elif tA and tB and A.shape[1] != B.shape[1]: correct = False elif not tA and tB and A.shape[2] != B.shape[1]: correct = False elif len(sA) == 3 and len(sB) == 3: if not tA and not tB and A.shape[2] != B.shape[1]: correct = False elif tA and not tB and A.shape[1] != B.shape[1]: correct = False elif tA and tB and A.shape[1] != B.shape[2]: correct = False elif not tA and tB and A.shape[2] != B.shape[2]: correct = False if out is not None: sout = out.shape # special case common in backprop if not correct and len(sA) == 3 and len(sB) == 3: if sout[0] == sA[2] and sout[1] == sB[2] and sA[0] == sB[0] and sA[1] == sB[1]: correct = True else: if len(sA) == 2 and len(sB) == 2: if not tA and not tB: sout = (sA[0], sB[1]) elif tA and tB: sout = (sA[1], sB[0]) elif tA and not tB: sout = (sA[1], sB[1]) elif not tA and tB: sout = (sA[0], sB[0]) elif len(sA) == 3 and len(sB) == 2: if not tA and not tB: sout = (sA[0], sA[1], sB[1]) elif tA and tB: sout = (sA[0], sA[2], sB[0]) elif tA and not tB: sout = (sA[0], sA[2], sB[1]) elif not tA and tB: sout = (sA[0], sA[1], sB[0]) elif len(sA) == 3 and len(sB) == 3: if not tA and not tB: sout = (sA[0], sA[1], sB[2]) elif tA and tB: sout = (sA[0], sA[2], sB[1]) elif tA and not tB: sout = (sA[0], sA[2], sB[2]) elif not tA and tB: sout = (sA[0], sA[1], sB[1]) if not correct: raise ValueError( f"Tensor dimensions incorrect for matrix mulitiplication: A x B: {sA} x {sB} with transpose for A x B: {tA} x {tB}.", ) return sout def gemv_4bit( A: Tensor, B: Tensor, out: Optional[torch.Tensor] = None, transposed_A=False, transposed_B=False, state=None, ): if state is None: raise ValueError("state cannot be None. gemv_4bit() requires the state from quantize_4bit()") absmax = state.absmax if state.nested: absmax = dequantize_blockwise(absmax, state.state2) + state.offset if out is not None: torch.ops.bitsandbytes.gemv_4bit.out( A, B, state.shape, absmax, state.code, state.blocksize, out=out, ) return out return torch.ops.bitsandbytes.gemv_4bit.default( A, B, state.shape, absmax, state.code, state.blocksize, ) def igemm( A: Tensor, B: Tensor, out: Optional[torch.Tensor] = None, transposed_A=False, transposed_B=False, ): sout = check_matmul(A, B, out, transposed_A, transposed_B) if out is None: out = torch.zeros(size=sout, dtype=torch.int32, device=A.device) if len(A.shape) == 3 and len(B.shape) == 3: if A.shape[0] == B.shape[0] and A.shape[2] == B.shape[1]: return batched_igemm(A, B, out) sA = A.shape sB = B.shape if transposed_A and len(sA) == 2: sA = (sA[1], sA[0]) elif transposed_A and len(sA) == 3: sA = (sA[0], sA[2], sA[0]) if transposed_B and len(sB) == 2: sB = (sB[1], sB[0]) elif transposed_B and len(sB) == 3: sB = (sB[0], sB[2], sB[0]) # this is a mess: cuBLAS expect column major, but PyTorch is row major. # So to perform the matrix multiplication, we have to treat A, B, and C matrices # (transpose of row major is column major) # This means we compute B^T A^T = C^T and we explicitly switch the dimensions of each of these # matrices in the input arguments for cuBLAS # column major: A @ B = C: [m, k] @ [k, n] = [m, n] # row major: B^T @ A^T = C^T: [m, k] @ [k, n] = [m, n] # column major with row major layout: B^T @ A^T = C^T: [k, m] @ [n, k] = [n, m] if len(sB) == 2: if B.stride()[0] == B.shape[1]: transposed_B = False elif B.stride()[1] == B.shape[0]: transposed_B = True if len(A.shape) == 2: if A.stride()[0] == A.shape[1]: transposed_A = False elif A.stride()[1] == A.shape[0]: transposed_A = True else: if A.stride()[1] == A.shape[2]: transposed_A = False elif A.stride()[2] == A.shape[1]: transposed_A = True if len(sA) == 2: n = sA[0] ldb = A.stride()[1 if transposed_A else 0] elif len(sA) == 3 and len(sB) == 2: n = sA[0] * sA[1] ldb = sA[2] m = sB[1] k = sB[0] lda = B.stride()[(1 if transposed_B else 0)] ldc = sB[1] elif len(sB) == 3: # special case assert len(sA) == 3 if not (sA[0] == sB[0] and sA[1] == sB[1]): raise ValueError( f"Only bsi,bso->io supported for tensor contractions, but dims for A x B were: {sA} x {sB}", ) transposed_A = True transposed_B = False m = sB[2] n = sA[2] k = sB[0] * sB[1] lda = m ldb = sA[2] ldc = m ptr = CUBLAS_Context.get_instance().get_context(A.device) # B^T @ A^T = C^T # [km, nk -> mn] is_on_gpu([B, A, out]) lib.cigemm( ptr, ct.c_bool(transposed_B), ct.c_bool(transposed_A), ct.c_int32(m), ct.c_int32(n), ct.c_int32(k), get_ptr(B), get_ptr(A), get_ptr(out), ct.c_int32(lda), ct.c_int32(ldb), ct.c_int32(ldc), ) return out def batched_igemm( A: Tensor, B: Tensor, out: Optional[torch.Tensor] = None, transposed_A=False, transposed_B=False, ): if not len(A.shape) == 3 or not len(B.shape) == 3: raise ValueError(f"Expected 3-dimensional tensors for bmm, but got shapes A and B: {A.shape} and {B.shape}") sout = check_matmul(A, B, out, transposed_A, transposed_B) if out is None: out = torch.zeros(size=sout, dtype=torch.int32, device=A.device) if B.is_contiguous(): lda = B.stride()[1] transposed_A = False else: s = B.stride() if s[0] != B.shape[0]: B = B.contiguous() lda = B.stride()[1] elif s[2] == B.shape[1]: transposed_A = True lda = B.stride()[2] else: if s[2] == 1: B = B.contiguous() lda = B.stride()[1] elif s[1] == 1: B = B.contiguous() lda = B.stride()[1] else: B = B.contiguous() lda = B.stride()[1] if A.is_contiguous(): ldb = A.stride()[1] transposed_B = False else: s = A.stride() if s[0] != A.shape[0]: A = A.contiguous() ldb = A.stride()[1] transposed_B = False elif s[2] == A.shape[1]: ldb = A.stride()[2] transposed_B = True else: A = A.contiguous() ldb = A.stride()[1] transposed_B = False # this is a mess: cuBLAS expect column major, but PyTorch is row major. # So to perform the matrix multiplication, we have to treat A, B, and C matrices # (transpose of row major is column major) # This means we compute B^T A^T = C^T and we explicitly switch the dimensions of each of these # matrices in the input arguments for cuBLAS # column major: A @ B = C: [batch, m, k] @ [batch, k, n] = [batch, m, n] # row major: B^T @ A^T = C^T: [batch, m, k] @ [batch, k, n] = [batch, m, n] # column major with row major layout: B^T @ A^T = C^T: [batch, k, m] @ [batch, n, k] = [batch, n, m] num_batch = A.shape[0] n = A.shape[1] m = B.shape[2] k = B.shape[1] ldc = m strideA = B.shape[1] * B.shape[2] strideB = A.shape[1] * A.shape[2] strideC = A.shape[1] * B.shape[2] ptr = CUBLAS_Context.get_instance().get_context(A.device) is_on_gpu([B, A, out]) lib.cbatched_igemm( ptr, ct.c_bool(transposed_B), ct.c_bool(transposed_A), ct.c_int32(m), ct.c_int32(n), ct.c_int32(k), get_ptr(B), get_ptr(A), get_ptr(out), ct.c_int32(lda), ct.c_int32(ldb), ct.c_int32(ldc), ct.c_long(strideA), ct.c_long(strideB), ct.c_long(strideC), ct.c_uint32(num_batch), ) return out def int8_linear_matmul(A: torch.Tensor, B: torch.Tensor, out: Optional[torch.Tensor] = None, dtype=torch.int32): """Performs an 8-bit integer matrix multiplication. A linear transformation is applied such that `out = A @ B.T`. When possible, integer tensor core hardware is utilized to accelerate the operation. Args: A (`torch.Tensor`): The first matrix operand with the data type `torch.int8`. B (`torch.Tensor`): The second matrix operand with the data type `torch.int8`. out (`torch.Tensor`, *optional*): A pre-allocated tensor used to store the result. dtype (`torch.dtype`, *optional*): The expected data type of the output. Defaults to `torch.int32`. Raises: `NotImplementedError`: The operation is not supported in the current environment. `RuntimeError`: Raised when the cannot be completed for any other reason. Returns: `torch.Tensor`: The result of the operation. """ if out is not None: torch.ops.bitsandbytes.int8_linear_matmul.out(A, B, out) return out return torch.ops.bitsandbytes.int8_linear_matmul.default(A, B) def int8_mm_dequant( A: torch.Tensor, row_stats: torch.Tensor, col_stats: torch.Tensor, out: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None, ): """Performs dequantization on the result of a quantized int8 matrix multiplication. Args: A (`torch.Tensor` with dtype `torch.int32`): The result of a quantized int8 matrix multiplication. row_stats (`torch.Tensor`): The row-wise quantization statistics for the lhs operand of the matrix multiplication. col_stats (`torch.Tensor`): The column-wise quantization statistics for the rhs operand of the matrix multiplication. out (`torch.Tensor`, *optional*): A pre-allocated tensor to store the output of the operation. bias (`torch.Tensor`, *optional*): An optional bias vector to add to the result. Returns: `torch.Tensor`: The dequantized result with an optional bias, with dtype `torch.float16`. """ result = torch.ops.bitsandbytes.int8_mm_dequant.default(A, row_stats, col_stats, dtype=torch.float16, bias=bias) # TODO(matthewdouglas): Deprecate out kwarg if out is not None: return out.copy_(result) return result @deprecated("This function is deprecated and will be removed in a future release.", category=FutureWarning) def get_colrow_absmax( A: torch.Tensor, row_stats: Optional[torch.Tensor] = None, col_stats: Optional[torch.Tensor] = None, nnz_block_ptr: Optional[torch.Tensor] = None, threshold=0.0, ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: """ "Determine the quantization statistics for input matrix `A` in accordance to the `LLM.int8()` algorithm. The row-wise and column-wise absmax values are determined. For more information, see the [LLM.int8() paper](https://arxiv.org/abs/2208.07339). This function is useful for training, but for inference it is advised to use [`get_row_absmax`] instead. The column-wise quantization scales are not typically needed in inference scenarios. Args: A (`torch.Tensor` with dtype `torch.float16`): Input tensor. row_stats (`torch.Tensor`, *optional*): If provided, calculation of row statistics is skipped. col_stats (`torch.Tensor`, *optional*): If provided, calculation of column statistics is skipped. nnz_block_ptr (`torch.Tensor`, *optional*): Not used. threshold (`float`, *optional*): An optional threshold for sparse decomposition of outlier features. No outliers are held back when 0.0. Defaults to 0.0. Returns: `Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]`: A tuple containing quantization statistics. - `torch.Tensor` with dtype `torch.float32`: The row-wise quantization statistics. - `torch.Tensor` with dtype `torch.float32`: The column-wise quantization statistics. - `torch.Tensor` with dtype `torch.bool`, *optional*: A mask indicating the locations of outliers in the input tensor. """ assert A.is_floating_point() outlier_mask = None if row_stats is None or col_stats is None: absA = A.abs().view(-1, A.shape[-1]) if threshold > 0.0: # Filter outliers from stats when enabled outlier_mask = absA >= threshold absA.masked_fill_(outlier_mask, 0.0) if row_stats is None: # shape [rows]; unsqueeze(-1) gives [rows,1] # We have a CUDA kernel for row max, but not yet for cols. row_stats = get_row_absmax(A, threshold) if col_stats is None: # shape [cols]; unsqueeze(0) gives [1,cols] col_stats = absA.amax(dim=0, keepdim=False).float() return row_stats, col_stats, outlier_mask @deprecated("This function is deprecated and will be removed in a future release.", category=FutureWarning) def get_row_absmax(A: torch.Tensor, threshold=0.0): """Determine the quantization statistics for input matrix `A` in accordance to the `LLM.int8()` algorithm. For more information, see the [LLM.int8() paper](https://arxiv.org/abs/2208.07339). Args: A (`torch.Tensor` with dtype `torch.float16`): The input matrix. threshold (`float`, *optional*): An optional threshold for sparse decomposition of outlier features. No outliers are held back when 0.0. Defaults to 0.0. Returns: `torch.Tensor` with dtype `torch.float32`: The absolute maximum value for each row, with outliers ignored. """ assert A.dtype == torch.float16 rows = prod(A.shape[:-1]) cols = A.shape[-1] row_stats = torch.empty((rows,), dtype=torch.float32, device=A.device) is_on_gpu([A]) with _cuda_device_of(A): lib.cget_row_stats( get_ptr(A), get_ptr(row_stats), ct.c_float(threshold), ct.c_int32(rows), ct.c_int32(cols), _get_tensor_stream(A), ) return row_stats class COOSparseTensor: def __init__( self, rows: int, cols: int, nnz: int, rowidx: torch.Tensor, colidx: torch.Tensor, values: torch.Tensor ): assert rowidx.dtype == torch.int32 assert colidx.dtype == torch.int32 assert values.dtype == torch.float16 assert values.numel() == nnz assert rowidx.numel() == nnz assert colidx.numel() == nnz self.rows = rows self.cols = cols self.nnz = nnz self.rowidx = rowidx self.colidx = colidx self.values = values class CSRSparseTensor: def __init__(self, rows, cols, nnz, rowptr, colidx, values): assert rowptr.dtype == torch.int32 assert colidx.dtype == torch.int32 assert values.dtype == torch.float16 assert values.numel() == nnz assert colidx.numel() == nnz assert rowptr.numel() == rows + 1 self.rows = rows self.cols = cols self.nnz = nnz self.rowptr = rowptr self.colidx = colidx self.values = values class CSCSparseTensor: def __init__(self, rows, cols, nnz, colptr, rowidx, values): assert colptr.dtype == torch.int32 assert rowidx.dtype == torch.int32 assert values.dtype == torch.float16 assert values.numel() == nnz assert rowidx.numel() == nnz assert colptr.numel() == cols + 1 self.rows = rows self.cols = cols self.nnz = nnz self.colptr = colptr self.rowidx = rowidx self.values = values def coo2csr(cooA): values, counts = torch.unique(cooA.rowidx, return_counts=True) values.add_(1) rowptr = torch.zeros((cooA.rows + 1,), dtype=torch.int32, device=cooA.rowidx.device) rowptr.scatter_(index=values.long(), src=counts.int(), dim=0) rowptr.cumsum_(0) return CSRSparseTensor(cooA.rows, cooA.cols, cooA.nnz, rowptr, cooA.colidx, cooA.values) def coo2csc(cooA): val, col2rowidx = torch.sort(cooA.colidx) rowidx = cooA.rowidx[col2rowidx] values = cooA.values[col2rowidx] colvalues, counts = torch.unique(val, return_counts=True) colvalues.add_(1) colptr = torch.zeros((cooA.cols + 1,), dtype=torch.int32, device=cooA.colidx.device) colptr.scatter_(index=colvalues.long(), src=counts.int(), dim=0) colptr.cumsum_(0) return CSCSparseTensor(cooA.rows, cooA.cols, cooA.nnz, colptr, rowidx, values) def coo_zeros(rows, cols, nnz, device, dtype=torch.half): rowidx = torch.zeros((nnz,), dtype=torch.int32, device=device) colidx = torch.zeros((nnz,), dtype=torch.int32, device=device) values = torch.zeros((nnz,), dtype=dtype, device=device) return COOSparseTensor(rows, cols, nnz, rowidx, colidx, values) def int8_double_quant( A: torch.Tensor, col_stats: Optional[torch.Tensor] = None, row_stats: Optional[torch.Tensor] = None, out_col: Optional[torch.Tensor] = None, out_row: Optional[torch.Tensor] = None, threshold=0.0, ): """Determine the quantization statistics for input matrix `A` in accordance to the `LLM.int8()` algorithm. The statistics are determined both row-wise and column-wise (transposed). For more information, see the [LLM.int8() paper](https://arxiv.org/abs/2208.07339). This function is useful for training, but for inference it is advised to use [`int8_vectorwise_quant`] instead. This implementation performs additional column-wise transposed calculations which are not optimized. Args: A (`torch.Tensor` with dtype `torch.float16`): The input matrix. col_stats (`torch.Tensor`, *optional*): A pre-allocated tensor to hold the column-wise quantization scales. row_stats (`torch.Tensor`, *optional*): A pre-allocated tensor to hold the row-wise quantization scales. out_col (`torch.Tensor`, *optional*): A pre-allocated tensor to hold the column-wise quantized data. out_row (`torch.Tensor`, *optional*): A pre-allocated tensor to hold the row-wise quantized data. threshold (`float`, *optional*): An optional threshold for sparse decomposition of outlier features. No outliers are held back when 0.0. Defaults to 0.0. Returns: `Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]`: A tuple containing the quantized tensor and relevant statistics. - `torch.Tensor` with dtype `torch.int8`: The row-wise quantized data. - `torch.Tensor` with dtype `torch.int8`: The column-wise quantized data. - `torch.Tensor` with dtype `torch.float32`: The row-wise quantization scales. - `torch.Tensor` with dtype `torch.float32`: The column-wise quantization scales. - `torch.Tensor` with dtype `torch.int32`, *optional*: A list of column indices which contain outlier features. """ if row_stats is not None: raise ValueError("row_stats must be None. int8_double_quant() does not support pre-allocated row_stats.") if col_stats is not None: raise ValueError("col_stats must be None. int8_double_quant() does not support pre-allocated col_stats.") if out_col is not None: raise ValueError("out_col must be None. int8_double_quant() does not support pre-allocated out_col.") if out_row is not None: raise ValueError("out_row must be None. int8_double_quant() does not support pre-allocated out_row.") return torch.ops.bitsandbytes.int8_double_quant.default(A, threshold=threshold) def int8_vectorwise_dequant(A: torch.Tensor, stats: torch.Tensor): """Dequantizes a tensor with dtype `torch.int8` to `torch.float32`. Args: A (`torch.Tensor` with dtype `torch.int8`): The quantized int8 tensor. stats (`torch.Tensor` with dtype `torch.float32`): The row-wise quantization statistics. Returns: `torch.Tensor` with dtype `torch.float32`: The dequantized tensor. """ # To dequantize we divide by 127, or multiply by the reciprocal. return torch.ops.bitsandbytes.int8_vectorwise_dequant.default(A, stats) def int8_vectorwise_quant(A: torch.Tensor, threshold=0.0): """Quantizes a tensor with dtype `torch.float16` to `torch.int8` in accordance to the `LLM.int8()` algorithm. For more information, see the [LLM.int8() paper](https://arxiv.org/abs/2208.07339). Args: A (`torch.Tensor` with dtype `torch.float16`): The input tensor. threshold (`float`, *optional*): An optional threshold for sparse decomposition of outlier features. No outliers are held back when 0.0. Defaults to 0.0. Returns: `Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]`: A tuple containing the quantized tensor and relevant statistics. - `torch.Tensor` with dtype `torch.int8`: The quantized data. - `torch.Tensor` with dtype `torch.float32`: The quantization scales. - `torch.Tensor` with dtype `torch.int32`, *optional*: A list of column indices which contain outlier features. """ return torch.ops.bitsandbytes.int8_vectorwise_quant.default(A, threshold) def spmm_coo( cooA: Union[COOSparseTensor, torch.Tensor], B: torch.Tensor, out: Optional[torch.Tensor] = None, ): if not isinstance(cooA, COOSparseTensor): assert cooA.is_sparse and cooA.layout == torch.sparse_coo, ( "Tensor must be `COOSparseTensor or a PyTorch COO tensor." ) # Convert to custom COOSparseTensor cooA = COOSparseTensor( rows=cooA.shape[0], cols=cooA.shape[1], nnz=cooA._nnz(), rowidx=cooA.indices()[0].int(), colidx=cooA.indices()[1].int(), values=cooA.values(), ) if out is None: out = torch.empty((cooA.rows, B.shape[1]), device=B.device, dtype=B.dtype) nnz = cooA.nnz assert cooA.rowidx.numel() == nnz assert cooA.colidx.numel() == nnz assert cooA.values.numel() == nnz assert cooA.cols == B.shape[0] transposed_B = not B.is_contiguous() ldb = B.stride()[(1 if transposed_B else 0)] ldc = B.shape[1] ptr = Cusparse_Context.get_instance().context ptrRowidx = get_ptr(cooA.rowidx) ptrColidx = get_ptr(cooA.colidx) ptrValues = get_ptr(cooA.values) ptrB = get_ptr(B) ptrC = get_ptr(out) cnnz = ct.c_int32(cooA.nnz) crowsA = ct.c_int32(cooA.rows) ccolsA = ct.c_int32(cooA.cols) ccolsB = ct.c_int32(B.shape[1]) cldb = ct.c_int32(ldb) cldc = ct.c_int32(ldc) is_on_gpu([cooA.rowidx, cooA.colidx, cooA.values, B, out]) lib.cspmm_coo( ptr, ptrRowidx, ptrColidx, ptrValues, cnnz, crowsA, ccolsA, ccolsB, cldb, ptrB, cldc, ptrC, ct.c_bool(transposed_B), ) return out def spmm_coo_very_sparse(cooA, B, dequant_stats=None, out=None): if out is None: out = torch.zeros((cooA.rows, B.shape[1]), device=B.device, dtype=cooA.values.dtype) nnz = cooA.nnz assert cooA.rowidx.numel() == nnz assert cooA.colidx.numel() == nnz assert cooA.values.numel() == nnz assert cooA.cols == B.shape[0], f"{cooA.cols} vs {B.shape}" _, counts = torch.unique(cooA.rowidx, return_counts=True) offset = counts.cumsum(0).int() max_count, max_idx = torch.sort(counts, descending=True) max_idx = max_idx.int() max_count = max_count.int() assert max_count[0] <= 32, f"Current max count per row is 8 but found {max_count[0]}." assert B.dtype in [torch.float16, torch.int8] ptrOffset = get_ptr(offset) ptrMaxCount = get_ptr(max_count) ptrMaxIdx = get_ptr(max_idx) ptrRowidx = get_ptr(cooA.rowidx) ptrColidx = get_ptr(cooA.colidx) ptrValues = get_ptr(cooA.values) ptrB = get_ptr(B) ptrC = get_ptr(out) ptrDequantStats = get_ptr(dequant_stats) cnnz_rows = ct.c_int32(counts.numel()) cnnz = ct.c_int32(cooA.nnz) crowsA = ct.c_int32(cooA.rows) crowsB = ct.c_int32(B.shape[1]) ccolsB = ct.c_int32(B.shape[1]) with _cuda_device_of(B): is_on_gpu([cooA.rowidx, cooA.colidx, cooA.values, B, out, dequant_stats]) if B.dtype == torch.float16: lib.cspmm_coo_very_sparse_naive_fp16( ptrMaxCount, ptrMaxIdx, ptrOffset, ptrRowidx, ptrColidx, ptrValues, ptrB, ptrC, ptrDequantStats, cnnz_rows, cnnz, crowsA, crowsB, ccolsB, ) elif B.dtype == torch.int8: lib.cspmm_coo_very_sparse_naive_int8( ptrMaxCount, ptrMaxIdx, ptrOffset, ptrRowidx, ptrColidx, ptrValues, ptrB, ptrC, ptrDequantStats, cnnz_rows, cnnz, crowsA, crowsB, ccolsB, ) # else: assertion error return out C = 127.0 bitsandbytes/bitsandbytes/nn/000077500000000000000000000000001506752665600167015ustar00rootroot00000000000000bitsandbytes/bitsandbytes/nn/__init__.py000066400000000000000000000011251506752665600210110ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from .modules import ( Embedding, Embedding4bit, Embedding8bit, EmbeddingFP4, EmbeddingNF4, Int8Params, Linear4bit, Linear8bitLt, LinearFP4, LinearNF4, OutlierAwareLinear, Params4bit, StableEmbedding, SwitchBackLinearBnb, ) from .triton_based_modules import ( StandardLinear, SwitchBackLinear, SwitchBackLinearGlobal, SwitchBackLinearVectorwise, ) bitsandbytes/bitsandbytes/nn/modules.py000066400000000000000000001171461506752665600207350ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. import copy from typing import Any, Optional, TypeVar, Union, overload import warnings import torch from torch import Tensor, device, dtype, nn import torch.nn.functional as F import bitsandbytes as bnb from bitsandbytes.cextension import HIP_ENVIRONMENT from bitsandbytes.functional import QuantState from bitsandbytes.optim import GlobalOptimManager from bitsandbytes.utils import INVERSE_LINEAR_8BIT_WEIGHTS_FORMAT_MAPPING, OutlierTracer T = TypeVar("T", bound="torch.nn.Module") class StableEmbedding(torch.nn.Embedding): """ Custom embedding layer designed to improve stability during training for NLP tasks by using 32-bit optimizer states. It is designed to reduce gradient variations that can result from quantization. This embedding layer is initialized with Xavier uniform initialization followed by layer normalization. Example: ``` # Initialize StableEmbedding layer with vocabulary size 1000, embedding dimension 300 embedding_layer = StableEmbedding(num_embeddings=1000, embedding_dim=300) # Reset embedding parameters embedding_layer.reset_parameters() # Perform a forward pass with input tensor input_tensor = torch.tensor([1, 2, 3]) output_embedding = embedding_layer(input_tensor) ``` Attributes: norm (`torch.nn.LayerNorm`): Layer normalization applied after the embedding. Methods: reset_parameters(): Reset embedding parameters using Xavier uniform initialization. forward(input: Tensor) -> Tensor: Forward pass through the stable embedding layer. """ def __init__( self, num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None, max_norm: Optional[float] = None, norm_type: float = 2.0, scale_grad_by_freq: bool = False, sparse: bool = False, _weight: Optional[Tensor] = None, device=None, dtype=None, ) -> None: """ Args: num_embeddings (`int`): The number of unique embeddings (vocabulary size). embedding_dim (`int`): The dimensionality of the embedding. padding_idx (`Optional[int]`): Pads the output with zeros at the given index. max_norm (`Optional[float]`): Renormalizes embeddings to have a maximum L2 norm. norm_type (`float`, defaults to `2.0`): The p-norm to compute for the `max_norm` option. scale_grad_by_freq (`bool`, defaults to `False`): Scale gradient by frequency during backpropagation. sparse (`bool`, defaults to `False`): Computes dense gradients. Set to `True` to compute sparse gradients instead. _weight (`Optional[Tensor]`): Pretrained embeddings. """ super().__init__( num_embeddings, embedding_dim, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse, _weight, device, dtype, ) self.norm = torch.nn.LayerNorm(embedding_dim, device=device) GlobalOptimManager.get_instance().register_module_override(self, "weight", {"optim_bits": 32}) def reset_parameters(self) -> None: torch.nn.init.xavier_uniform_(self.weight) self._fill_padding_idx_with_zero() """ !!! This is a redefinition of _fill_padding_idx_with_zero in torch.nn.Embedding to make the Layer compatible with Pytorch < 1.9. This means that if this changes in future PyTorch releases this need to change too which is cumbersome. However, with this we can ensure compatibility with previous PyTorch releases. """ def _fill_padding_idx_with_zero(self) -> None: if self.padding_idx is not None: with torch.no_grad(): self.weight[self.padding_idx].fill_(0) def forward(self, input: Tensor) -> Tensor: emb = F.embedding( input, self.weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse, ) # always apply layer norm in full precision emb = emb.to(torch.get_default_dtype()) return self.norm(emb).to(self.weight.dtype) class Embedding(torch.nn.Embedding): """ Embedding class to store and retrieve word embeddings from their indices. """ def __init__( self, num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None, max_norm: Optional[float] = None, norm_type: float = 2.0, scale_grad_by_freq: bool = False, sparse: bool = False, _weight: Optional[Tensor] = None, device: Optional[device] = None, ) -> None: """ Args: num_embeddings (`int`): The number of unique embeddings (vocabulary size). embedding_dim (`int`): The dimensionality of the embedding. padding_idx (`Optional[int]`): Pads the output with zeros at the given index. max_norm (`Optional[float]`): Renormalizes embeddings to have a maximum L2 norm. norm_type (`float`, defaults to `2.0`): The p-norm to compute for the `max_norm` option. scale_grad_by_freq (`bool`, defaults to `False`): Scale gradient by frequency during backpropagation. sparse (`bool`, defaults to `False`): Computes dense gradients. Set to `True` to compute sparse gradients instead. _weight (`Optional[Tensor]`): Pretrained embeddings. """ super().__init__( num_embeddings, embedding_dim, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse, _weight, device=device, ) GlobalOptimManager.get_instance().register_module_override(self, "weight", {"optim_bits": 32}) def reset_parameters(self) -> None: torch.nn.init.xavier_uniform_(self.weight) self._fill_padding_idx_with_zero() """ !!! This is a redefinition of _fill_padding_idx_with_zero in torch.nn.Embedding to make the Layer compatible with Pytorch < 1.9. This means that if this changes in future PyTorch releases this need to change too which is cumbersome. However, with this we can ensure compatibility with previous PyTorch releases. """ def _fill_padding_idx_with_zero(self) -> None: if self.padding_idx is not None: with torch.no_grad(): self.weight[self.padding_idx].fill_(0) def forward(self, input: Tensor) -> Tensor: emb = F.embedding( input, self.weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse, ) return emb class Params4bit(torch.nn.Parameter): def __new__( cls, data: Optional[torch.Tensor] = None, requires_grad=False, # quantized weights should be frozen by default quant_state: Optional[QuantState] = None, blocksize: Optional[int] = None, compress_statistics: bool = True, quant_type: str = "fp4", quant_storage: torch.dtype = torch.uint8, module: Optional["Linear4bit"] = None, bnb_quantized: bool = False, ) -> "Params4bit": if data is None: data = torch.empty(0) if blocksize is None: blocksize = 64 if not HIP_ENVIRONMENT else 128 self = torch.Tensor._make_subclass(cls, data, requires_grad) self.blocksize = blocksize self.compress_statistics = compress_statistics self.quant_type = quant_type self.quant_state = quant_state self.quant_storage = quant_storage self.bnb_quantized = bnb_quantized self.data = data self.module = module return self def __getstate__(self): state = self.__dict__.copy() state["data"] = self.data state["requires_grad"] = self.requires_grad return state def __setstate__(self, state): self.requires_grad = state["requires_grad"] self.blocksize = state["blocksize"] self.compress_statistics = state["compress_statistics"] self.quant_type = state["quant_type"] self.quant_state = state["quant_state"] self.data = state["data"] self.quant_storage = state["quant_storage"] self.bnb_quantized = state["bnb_quantized"] self.module = state["module"] def __deepcopy__(self, memo): new_instance = type(self).__new__(type(self)) state = self.__getstate__() new_instance.__setstate__(state) new_instance.quant_state = copy.deepcopy(state["quant_state"]) new_instance.data = copy.deepcopy(state["data"]) return new_instance def __copy__(self): new_instance = type(self).__new__(type(self)) state = self.__getstate__() new_instance.__setstate__(state) return new_instance @classmethod def from_prequantized( cls, data: torch.Tensor, quantized_stats: dict[str, Any], requires_grad: bool = False, device="cuda", module: Optional["Linear4bit"] = None, **kwargs, ) -> "Params4bit": self = torch.Tensor._make_subclass(cls, data.to(device)) self.requires_grad = requires_grad self.quant_state = QuantState.from_dict(qs_dict=quantized_stats, device=device) self.blocksize = self.quant_state.blocksize self.compress_statistics = self.quant_state.nested self.quant_type = self.quant_state.quant_type self.bnb_quantized = True self.quant_storage = data.dtype self.module = module if self.module is not None: self.module.quant_state = self.quant_state return self def _quantize(self, device): w = self.data.contiguous().to(device) w_4bit, quant_state = bnb.functional.quantize_4bit( w, blocksize=self.blocksize, compress_statistics=self.compress_statistics, quant_type=self.quant_type, quant_storage=self.quant_storage, ) self.data = w_4bit self.quant_state = quant_state if self.module is not None: self.module.quant_state = quant_state self.bnb_quantized = True return self def cpu(self): return self.to(device="cpu") def cuda(self, device: Optional[Union[int, device, str]] = None, non_blocking: bool = False): return self.to(device="cuda" if device is None else device, non_blocking=non_blocking) def xpu(self, device: Optional[Union[int, device, str]] = None, non_blocking: bool = False): return self.to(device="xpu" if device is None else device, non_blocking=non_blocking) @overload def to( self: T, device: Optional[Union[int, device]] = ..., dtype: Optional[Union[dtype, str]] = ..., non_blocking: bool = ..., ) -> T: ... @overload def to(self: T, dtype: Union[dtype, str], non_blocking: bool = ...) -> T: ... @overload def to(self: T, tensor: Tensor, non_blocking: bool = ...) -> T: ... def to(self, *args, **kwargs): device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs) if device is not None and device.type != "meta" and not self.bnb_quantized: return self._quantize(device) else: if self.quant_state is not None: self.quant_state.to(device) new_param = Params4bit( super().to(device=device, dtype=dtype, non_blocking=non_blocking), requires_grad=self.requires_grad, quant_state=self.quant_state, blocksize=self.blocksize, compress_statistics=self.compress_statistics, quant_type=self.quant_type, quant_storage=self.quant_storage, bnb_quantized=self.bnb_quantized, ) return new_param @classmethod def __torch_function__(cls, func, types, args=(), kwargs=None): if kwargs is None: kwargs = {} if func in [torch.chunk, torch.split]: tensor = args[0] result = super().__torch_function__(func, types, args, kwargs) if isinstance(result, tuple): return tuple( cls( data=chunk, requires_grad=tensor.requires_grad, quant_state=tensor.quant_state, blocksize=tensor.blocksize, compress_statistics=tensor.compress_statistics, quant_type=tensor.quant_type, quant_storage=tensor.quant_storage, module=tensor.module, bnb_quantized=tensor.bnb_quantized, ) for chunk in result ) else: return cls( data=result, requires_grad=tensor.requires_grad, quant_state=tensor.quant_state, blocksize=tensor.blocksize, compress_statistics=tensor.compress_statistics, quant_type=tensor.quant_type, quant_storage=tensor.quant_storage, module=tensor.module, bnb_quantized=tensor.bnb_quantized, ) return super().__torch_function__(func, types, args, kwargs) def fix_4bit_weight_quant_state_from_module(module: Union["Embedding4bit", "Linear4bit"]): if getattr(module.weight, "quant_state", None) is not None: return if getattr(module, "quant_state", None) is None: warnings.warn( "FP4 quantization state not initialized. Please call .cuda() or .to(device) on the LinearFP4 layer first.", ) # the quant state got lost when the parameter got converted. This happens for example for fsdp # since we registered the module, we can recover the state here assert module.weight.shape[1] == 1 if not isinstance(module.weight, Params4bit): module.weight = Params4bit(module.weight, quant_storage=module.quant_storage, bnb_quantized=True) module.weight.quant_state = module.quant_state class Linear4bit(nn.Linear): """ This class is the base module for the 4-bit quantization algorithm presented in [QLoRA](https://arxiv.org/abs/2305.14314). QLoRA 4-bit linear layers uses blockwise k-bit quantization under the hood, with the possibility of selecting various compute datatypes such as FP4 and NF4. In order to quantize a linear layer one should first load the original fp16 / bf16 weights into the Linear4bit module, then call `quantized_module.to("cuda")` to quantize the fp16 / bf16 weights. Example: ```python import torch import torch.nn as nn import bitsandbytes as bnb from bnb.nn import Linear4bit fp16_model = nn.Sequential( nn.Linear(64, 64), nn.Linear(64, 64) ) quantized_model = nn.Sequential( Linear4bit(64, 64), Linear4bit(64, 64) ) quantized_model.load_state_dict(fp16_model.state_dict()) quantized_model = quantized_model.to(0) # Quantization happens here ``` """ def __init__( self, input_features, output_features, bias=True, compute_dtype=None, compress_statistics=True, quant_type="fp4", quant_storage=torch.uint8, device=None, ): """ Initialize Linear4bit class. Args: input_features (`str`): Number of input features of the linear layer. output_features (`str`): Number of output features of the linear layer. bias (`bool`, defaults to `True`): Whether the linear class uses the bias term as well. """ super().__init__(input_features, output_features, bias, device) self.weight = Params4bit( self.weight.data, requires_grad=False, compress_statistics=compress_statistics, quant_type=quant_type, quant_storage=quant_storage, module=self, ) # self.persistent_buffers = [] # TODO consider as way to save quant state self.compute_dtype = compute_dtype self.compute_type_is_set = compute_dtype is not None self.quant_state = None self.quant_storage = quant_storage def set_compute_type(self, x): if x.dtype in [torch.float32, torch.bfloat16]: # the input is in a dtype that is safe to compute in, we switch # to this type for speed and stability self.compute_dtype = x.dtype elif x.dtype == torch.float16: # we take the compoute dtype passed into the layer if self.compute_dtype in [None, torch.float32] and (x.numel() == x.shape[-1]): # single batch inference with input torch.float16 and compute_dtype float32 -> slow inference when it could be fast # warn the user about this warnings.warn( "Input type into Linear4bit is torch.float16, but bnb_4bit_compute_dtype=torch.float32 (default). This will lead to slow inference.", ) warnings.filterwarnings("ignore", message=".*inference.") if self.compute_dtype in [None, torch.float32] and (x.numel() != x.shape[-1]): warnings.warn( "Input type into Linear4bit is torch.float16, but bnb_4bit_compute_dtype=torch.float32 (default). This will lead to slow inference or training speed.", ) warnings.filterwarnings("ignore", message=".*inference or training") def _save_to_state_dict(self, destination, prefix, keep_vars): """ save weight and bias, then fill state_dict with components of quant_state """ super()._save_to_state_dict(destination, prefix, keep_vars) # saving weight and bias if getattr(self.weight, "quant_state", None) is not None: for k, v in self.weight.quant_state.as_dict(packed=True).items(): destination[prefix + "weight." + k] = v if keep_vars else v.detach() def forward(self, x: torch.Tensor): fix_4bit_weight_quant_state_from_module(self) # weights are cast automatically as Int8Params, but the bias has to be cast manually if self.bias is not None and self.bias.dtype != x.dtype: self.bias.data = self.bias.data.to(x.dtype) if not self.compute_type_is_set: self.set_compute_type(x) self.compute_type_is_set = True inp_dtype = x.dtype if self.compute_dtype is not None: x = x.to(self.compute_dtype) bias = None if self.bias is None else self.bias.to(self.compute_dtype) weight = self.weight.t() return bnb.matmul_4bit(x, weight, bias=bias, quant_state=self.weight.quant_state).to(inp_dtype) class LinearFP4(Linear4bit): """ Implements the FP4 data type. """ def __init__( self, input_features, output_features, bias=True, compute_dtype=None, compress_statistics=True, quant_storage=torch.uint8, device=None, ): """ Args: input_features (`str`): Number of input features of the linear layer. output_features (`str`): Number of output features of the linear layer. bias (`bool`, defaults to `True`): Whether the linear class uses the bias term as well. """ super().__init__( input_features, output_features, bias, compute_dtype, compress_statistics, "fp4", quant_storage, device, ) class LinearNF4(Linear4bit): """Implements the NF4 data type. Constructs a quantization data type where each bin has equal area under a standard normal distribution N(0, 1) that is normalized into the range [-1, 1]. For more information read the paper: QLoRA: Efficient Finetuning of Quantized LLMs (https://arxiv.org/abs/2305.14314) Implementation of the NF4 data type in bitsandbytes can be found in the `create_normal_map` function in the `functional.py` file: https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L236. """ def __init__( self, input_features, output_features, bias=True, compute_dtype=None, compress_statistics=True, quant_storage=torch.uint8, device=None, ): """ Args: input_features (`str`): Number of input features of the linear layer. output_features (`str`): Number of output features of the linear layer. bias (`bool`, defaults to `True`): Whether the linear class uses the bias term as well. """ super().__init__( input_features, output_features, bias, compute_dtype, compress_statistics, "nf4", quant_storage, device, ) class Int8Params(torch.nn.Parameter): def __new__( cls, data: Optional[torch.Tensor] = None, requires_grad=True, has_fp16_weights=False, CB: Optional[torch.Tensor] = None, SCB: Optional[torch.Tensor] = None, ): if data is None: data = torch.empty(0) obj = torch.Tensor._make_subclass(cls, data, requires_grad) obj.CB = CB obj.SCB = SCB obj.has_fp16_weights = has_fp16_weights return obj def _quantize(self, device): if self.has_fp16_weights: return super().to(device) # We quantize the weight and store in 8bit row-major B = self.data.contiguous().to(device=device, dtype=torch.float16) CB, SCB, _ = bnb.functional.int8_vectorwise_quant(B) self.data = CB self.CB = CB self.SCB = SCB return self def cpu(self): return self.to(device="cpu") def cuda(self, device: Optional[Union[int, device, str]] = None, non_blocking: bool = False): return self.to(device="cuda" if device is None else device, non_blocking=non_blocking) def xpu(self, device: Optional[Union[int, device, str]] = None, non_blocking: bool = False): return self.to(device="xpu" if device is None else device, non_blocking=non_blocking) def __deepcopy__(self, memo): # adjust this if new arguments are added to the constructor new_instance = type(self).__new__( type(self), data=copy.deepcopy(self.data, memo), requires_grad=self.requires_grad, has_fp16_weights=self.has_fp16_weights, CB=copy.deepcopy(self.CB, memo), SCB=copy.deepcopy(self.SCB, memo), ) return new_instance @overload def to( self: T, device: Optional[Union[int, device]] = ..., dtype: Optional[Union[dtype, str]] = ..., non_blocking: bool = ..., ) -> T: ... @overload def to(self: T, dtype: Union[dtype, str], non_blocking: bool = ...) -> T: ... @overload def to(self: T, tensor: Tensor, non_blocking: bool = ...) -> T: ... def to(self, *args, **kwargs): device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs) is_quantized = self.data.dtype == torch.int8 if not is_quantized and device is not None and device.type != "meta" and self.data.device.type == "cpu": # We're moving from a CPU device to a non-meta device. # In this circumstance, we want to quantize if we haven't already. return self._quantize(device) # Create a new parameter on the target device. new_param = Int8Params( super().to(device=device, dtype=dtype, non_blocking=non_blocking), requires_grad=self.requires_grad, has_fp16_weights=self.has_fp16_weights, ) # If we had already quantized, move the statistics appropriately. if is_quantized: new_param.CB = new_param.data if self.SCB is not None and device is not None: new_param.SCB = self.SCB.to(device) return new_param def maybe_rearrange_weight(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): weight = state_dict.get(f"{prefix}weight") if weight is None: # if the state dict has no weights for this layer (e.g., LoRA finetuning), do nothing return weight_format = state_dict.pop(f"{prefix}weight_format", "row") if isinstance(weight_format, torch.Tensor): weight_format = weight_format.item() # For new weights format storage type, we explicitly check # if weights_format is on the mapping if isinstance(weight_format, int) and weight_format not in INVERSE_LINEAR_8BIT_WEIGHTS_FORMAT_MAPPING: raise ValueError(f"Expected supported weight format - got {weight_format}") elif isinstance(weight_format, int) and weight_format in INVERSE_LINEAR_8BIT_WEIGHTS_FORMAT_MAPPING: weight_format = INVERSE_LINEAR_8BIT_WEIGHTS_FORMAT_MAPPING[weight_format] if weight_format != "row": raise ValueError(f"Only 'row' weight format is supported, got {weight_format}") class Embedding8bit(nn.Embedding): """ This class implements [LLM.int8()](https://arxiv.org/abs/2208.07339) algorithm for embedding layer Quantization API is similar to Linear8bitLt: ```python import torch import torch.nn as nn from bitsandbytes.nn import Embedding8bit fp16_module = nn.Embedding(128, 64) int8_module = Embedding8bit(128, 64) int8_module.load_state_dict(fp16_module.state_dict()) int8_module = int8_module.to(0) # Quantization happens here ``` """ def __init__(self, num_embeddings, embedding_dim, device=None, dtype=None): super().__init__(num_embeddings, embedding_dim, device=device, dtype=dtype) self.dtype = self.weight.data.dtype self.weight = Int8Params(self.weight.data, has_fp16_weights=False, requires_grad=False) def _save_to_state_dict(self, destination, prefix, keep_vars): raise NotImplementedError("Saving Embedding8bit module is not implemented") def forward(self, input: Tensor) -> Tensor: if not hasattr(self.weight, "SCB"): raise RuntimeError("Embedding layer is not quantized. Please call .cuda() or .to(device) first.") rows = self.weight.data row_stats = self.weight.SCB assert rows.shape == (self.num_embeddings, self.embedding_dim) assert row_stats.shape == (self.num_embeddings,) compressed_output = F.embedding(input, rows) compressed_output_stats = F.embedding(input, row_stats.view(self.num_embeddings, 1)) output = compressed_output * (compressed_output_stats / 127.0) return output.to(self.dtype) class Embedding4bit(nn.Embedding): """ This is the base class similar to Linear4bit. It implements the 4-bit quantization algorithm presented in [QLoRA](https://arxiv.org/abs/2305.14314) for embeddings. Quantization API is similar to Linear4bit: ```python import torch import torch.nn as nn from bitsandbytes.nn import Embedding4bit fp16_module = nn.Embedding(128, 64) quantized_module = Embedding4bit(128, 64) quantized_module.load_state_dict(fp16_module.state_dict()) quantized_module = quantized_module.to(0) # Quantization happens here ``` """ def __init__( self, num_embeddings, embedding_dim, dtype=None, quant_type="fp4", quant_storage=torch.uint8, device=None, ): super().__init__(num_embeddings, embedding_dim, device=device, dtype=dtype) self.dtype = self.weight.data.dtype self.weight = Params4bit( self.weight.data, requires_grad=False, compress_statistics=None, quant_type=quant_type, quant_storage=quant_storage, module=self, ) blocksize = self.weight.blocksize if embedding_dim % blocksize != 0: warnings.warn( f"Embedding size {embedding_dim} is not divisible by block size {blocksize}. " "This will lead to slow inference.", ) def _forward_with_partial_dequantize(self, input: Tensor): assert self.embedding_dim % self.weight.quant_state.blocksize == 0 w_4bit_uint8 = self.weight.data.view(torch.uint8).view(self.num_embeddings * self.embedding_dim // 2, 1) output_4bit = torch.nn.functional.embedding( weight=w_4bit_uint8.view(self.num_embeddings, self.embedding_dim // 2), input=input, ).view(-1, 1) assert output_4bit.shape == (input.numel() * self.embedding_dim // 2, 1) blocks_per_emb = self.embedding_dim // self.weight.blocksize absmax = self.weight.quant_state.absmax assert absmax.shape == (self.num_embeddings * blocks_per_emb,) output_absmax = torch.nn.functional.embedding( weight=absmax.view(self.num_embeddings, blocks_per_emb), input=input, ).view( -1, ) assert output_absmax.shape == (input.numel() * blocks_per_emb,) output_quant_state = copy.deepcopy(self.weight.quant_state) output_quant_state.absmax = output_absmax output_quant_state.shape = torch.Size((*input.shape, self.embedding_dim)) output = bnb.functional.dequantize_4bit(output_4bit, output_quant_state) assert output.shape == (*input.shape, self.embedding_dim) return output.to(self.dtype) def _save_to_state_dict(self, destination, prefix, keep_vars): raise NotImplementedError("Saving Embedding4bit module is not implemented") def forward(self, input: Tensor) -> Tensor: fix_4bit_weight_quant_state_from_module(self) if self.embedding_dim % self.weight.quant_state.blocksize == 0: return self._forward_with_partial_dequantize(input) dequantized_weight = bnb.functional.dequantize_4bit(self.weight.data, self.weight.quant_state) return torch.nn.functional.embedding( weight=dequantized_weight, input=input, ).to(self.dtype) class EmbeddingFP4(Embedding4bit): def __init__( self, num_embeddings, embedding_dim, dtype=None, quant_storage=torch.uint8, device=None, ): super().__init__( num_embeddings, embedding_dim, dtype=dtype, quant_type="fp4", quant_storage=quant_storage, device=device, ) class EmbeddingNF4(Embedding4bit): def __init__( self, num_embeddings, embedding_dim, dtype=None, quant_storage=torch.uint8, device=None, ): super().__init__( num_embeddings, embedding_dim, dtype=dtype, quant_type="nf4", quant_storage=quant_storage, device=device, ) class Linear8bitLt(nn.Linear): """ This class is the base module for the [LLM.int8()](https://arxiv.org/abs/2208.07339) algorithm. To read more about it, have a look at the paper. In order to quantize a linear layer one should first load the original fp16 / bf16 weights into the Linear8bitLt module, then call `int8_module.to("cuda")` to quantize the fp16 weights. Example: ```python import torch import torch.nn as nn import bitsandbytes as bnb from bnb.nn import Linear8bitLt fp16_model = nn.Sequential( nn.Linear(64, 64), nn.Linear(64, 64) ) int8_model = nn.Sequential( Linear8bitLt(64, 64, has_fp16_weights=False), Linear8bitLt(64, 64, has_fp16_weights=False) ) int8_model.load_state_dict(fp16_model.state_dict()) int8_model = int8_model.to(0) # Quantization happens here ``` """ def __init__( self, input_features: int, output_features: int, bias=True, has_fp16_weights=True, threshold=0.0, index=None, device=None, ): """ Initialize Linear8bitLt class. Args: input_features (`int`): Number of input features of the linear layer. output_features (`int`): Number of output features of the linear layer. bias (`bool`, defaults to `True`): Whether the linear class uses the bias term as well. """ super().__init__(input_features, output_features, bias, device) self.state = bnb.MatmulLtState() self.index = index self.state.threshold = threshold self.state.has_fp16_weights = has_fp16_weights if threshold > 0.0 and not has_fp16_weights: self.state.use_pool = True self.weight = Int8Params(self.weight.data, has_fp16_weights=has_fp16_weights, requires_grad=has_fp16_weights) self._register_load_state_dict_pre_hook(maybe_rearrange_weight) def _save_to_state_dict(self, destination, prefix, keep_vars): super()._save_to_state_dict(destination, prefix, keep_vars) # we only need to save SCB as extra data, because CB for quantized weights is already stored in weight.data scb_name = "SCB" # case 1: .cuda was called, SCB is in self.weight param_from_weight = getattr(self.weight, scb_name) # case 2: self.init_8bit_state was called, SCB is in self.state param_from_state = getattr(self.state, scb_name) key_name = prefix + f"{scb_name}" # We now only save in row-major. This format information is stored for backwards compatibility. format_name = prefix + "weight_format" if not self.state.has_fp16_weights: if param_from_weight is not None: destination[key_name] = param_from_weight if keep_vars else param_from_weight.detach() destination[format_name] = torch.tensor(0, dtype=torch.uint8) elif param_from_state is not None: destination[key_name] = param_from_state if keep_vars else param_from_state.detach() destination[format_name] = torch.tensor(0, dtype=torch.uint8) def _load_from_state_dict( self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs, ): super()._load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs, ) unexpected_copy = list(unexpected_keys) for key in unexpected_copy: input_name = key[len(prefix) :] if input_name == "SCB": if self.weight.SCB is None: # buffers not yet initialized, can't access them directly without quantizing first raise RuntimeError( "Loading a quantized checkpoint into non-quantized Linear8bitLt is " "not supported. Please call module.cuda() before module.load_state_dict()", ) input_param = state_dict[key] self.weight.SCB.copy_(input_param) if self.state.SCB is not None: self.state.SCB = self.weight.SCB unexpected_keys.remove(key) def init_8bit_state(self): self.state.CB = self.weight.CB self.state.SCB = self.weight.SCB self.weight.CB = None self.weight.SCB = None def to(self, *args, **kwargs): # Call the parent to() method to handle standard parameter/buffer movement result = super().to(*args, **kwargs) device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs) # Handle state tensors if needed. if device is not None: if result.state.CB is not None: result.state.CB = result.state.CB.to(device) if result.state.SCB is not None: result.state.SCB = result.state.SCB.to(device) return result def forward(self, x: torch.Tensor): self.state.is_training = self.training if self.weight.CB is not None: self.init_8bit_state() # weights are cast automatically as Int8Params, but the bias has to be cast manually if self.bias is not None and self.bias.dtype != x.dtype: self.bias.data = self.bias.data.to(x.dtype) out = bnb.matmul(x, self.weight, bias=self.bias, state=self.state) if not self.state.has_fp16_weights and self.state.CB is not None: self.weight.data = self.state.CB return out class OutlierAwareLinear(nn.Linear): def __init__(self, input_features, output_features, bias=True, device=None): super().__init__(input_features, output_features, bias, device) self.outlier_dim = None self.is_quantized = False def forward_with_outliers(self, x, outlier_idx): raise NotImplementedError("Please override the `forward_with_outliers(self, x, outlier_idx)` function") def quantize_weight(self, w, outlier_idx): raise NotImplementedError("Please override the `quantize_weights(self, w, outlier_idx)` function") def forward(self, x): if self.outlier_dim is None: tracer = OutlierTracer.get_instance() if not tracer.is_initialized(): print("Please use OutlierTracer.initialize(model) before using the OutlierAwareLinear layer") outlier_idx = tracer.get_outliers(self.weight) # print(outlier_idx, tracer.get_hvalue(self.weight)) self.outlier_dim = outlier_idx if not self.is_quantized: w = self.quantize_weight(self.weight, self.outlier_dim) self.weight.data.copy_(w) self.is_quantized = True class SwitchBackLinearBnb(nn.Linear): def __init__( self, input_features, output_features, bias=True, has_fp16_weights=True, memory_efficient_backward=False, threshold=0.0, index=None, device=None, ): super().__init__(input_features, output_features, bias, device) self.state = bnb.MatmulLtState() self.index = index self.state.threshold = threshold self.state.has_fp16_weights = has_fp16_weights self.state.memory_efficient_backward = memory_efficient_backward if threshold > 0.0 and not has_fp16_weights: self.state.use_pool = True self.weight = Int8Params(self.weight.data, has_fp16_weights=has_fp16_weights, requires_grad=has_fp16_weights) def init_8bit_state(self): self.state.CB = self.weight.CB self.state.SCB = self.weight.SCB self.weight.CB = None self.weight.SCB = None def forward(self, x): self.state.is_training = self.training if self.weight.CB is not None: self.init_8bit_state() return bnb.matmul_mixed(x.half(), self.weight.half(), bias=None, state=self.state) + self.bias bitsandbytes/bitsandbytes/nn/parametrize.py000066400000000000000000000162011506752665600215760ustar00rootroot00000000000000from functools import partial from typing import Any, Literal, Optional import torch import torch.nn as nn import torch.nn.utils.parametrize as P from .. import functional as F class Bnb4bitParametrization(nn.Module): """ A parametrization module that handles dequantization of a 4-bit quantized parameter. The parameter data is expected to be already quantized when this parametrization is applied. This module will dequantize the parameter data to its original floating-point representation when the forward method is called (i.e. when the parameter is accessed). Args: quant_state (`F.QuantState`): The quantization state containing the necessary information for dequantization. """ def __init__(self, quant_state: F.QuantState): super().__init__() self.quant_state = quant_state @torch.no_grad() def forward(self, quantized_param: torch.Tensor) -> torch.Tensor: """ Forward pass to dequantize the parameter. Args: quantized_param (`torch.Tensor`): The quantized parameter tensor (from .original) Returns: `torch.Tensor`: The dequantized parameter tensor in the original shape and dtype. """ return F.dequantize_4bit(quantized_param, self.quant_state) def replace_parameter_4bit_prequantized( module: nn.Module, param_name: str, qs_dict: dict[str, Any], device: torch.device ): if not hasattr(module, param_name): raise AttributeError(f"Module does not have parameter '{param_name}'") original_param = getattr(module, param_name) if not isinstance(original_param, nn.Parameter): raise TypeError(f"Parameter '{param_name}' is not an instance of nn.Parameter") quant_state = F.QuantState.from_dict(qs_dict, device=device) # Apply a parametrization to the module to handle dequantization. P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True) # Next, register hooks. _register_parametrization_hooks(module, param_name) def replace_parameter_4bit( module: nn.Module, param_name: str, compress_statistics: bool = False, quant_type: Literal["nf4", "fp4"] = "nf4", blocksize: Optional[int] = None, ): """ Replace a module parameter with a 4-bit quantized version using parametrization. This function quantizes an existing parameter in a PyTorch module to 4-bit precision and sets up parametrization to handle automatic dequantization during forward passes. The original parameter is replaced with quantized data, and a parametrization layer is registered to manage the quantization state and dequantization process. Additional, it registers a state dict post-hook to ensure that the quantization state is saved correctly when the model's state dict is saved. It is useful for MoE models or other scenarios where you want to quantize parameters outside of nn.Linear layers without changing the model's architecture. This feature is experimental and may change in future releases. Args: module (`nn.Module`): The PyTorch module containing the parameter to be quantized. param_name (`str`): The name of the parameter within the module to quantize. compress_statistics (`bool`, *optional*, defaults to `False`): Whether to compress quantization statistics to reduce memory usage. quant_type (`Literal["nf4", "fp4"]`, *optional*, defaults to `"nf4"`): The quantization format to use. blocksize (`int`, *optional*, defaults to `None`): The block size for quantization. If None, uses the default block size. Raises: AttributeError: If the module does not have the specified parameter. TypeError: If the specified attribute is not an instance of nn.Parameter. """ if not hasattr(module, param_name): raise AttributeError(f"Module does not have parameter '{param_name}'") original_param = getattr(module, param_name) if not isinstance(original_param, nn.Parameter): raise TypeError(f"Parameter '{param_name}' is not an instance of nn.Parameter") # Quantize the original parameter. quantized_data, quant_state = F.quantize_4bit( original_param.data, blocksize=blocksize, compress_statistics=compress_statistics, quant_type=quant_type, ) # Replace the parameter with the quantized data. setattr(module, param_name, nn.Parameter(quantized_data, requires_grad=False)) del original_param # Apply a parametrization to the module to handle dequantization. P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True) # Next, register hooks. _register_parametrization_hooks(module, param_name) def _disable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...], output: Any): P._cache_enabled -= 1 if not P._cache_enabled: P._cache = {} def _enable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...]): P._cache_enabled += 1 def _register_parametrization_hooks(module: nn.Module, param_name: str): # Register a state dict hook for saving. Note that this requires torch >= 2.5.0. if torch.__version__ >= (2, 5): module.register_state_dict_post_hook( partial( _parametrized_state_dict_post_hook, param_name=param_name, ) ) # Register hooks to enable caching for the dequantization parametrization. # This helps preserve time and memory when the same quantized parameter # is accessed multiple times in the forward computation. module.register_forward_pre_hook(_enable_parametrization_cache) module.register_forward_hook(_disable_parametrization_cache) def _parametrized_state_dict_post_hook( module: nn.Module, state_dict: dict[str, Any], prefix: str, local_metadata: Any, *, param_name: str = "weight", **kwargs: dict[str, Any], ) -> None: """ Hook to modify the state dict to include the quantization state. """ original_key = f"{prefix}parametrizations.{param_name}.original" if original_key in state_dict: # Create a clean entry. # The `parametrizations.{param_name}.original` key will have the quantized data, # but we would like it to keep it in the state_dict as `{param_name}`. clean_key = f"{prefix}{param_name}" state_dict[clean_key] = state_dict.pop(original_key) assert P.is_parametrized(module, param_name) # Find the parametrization, which should have the quantization state. parametrization: Bnb4bitParametrization = next( filter(lambda x: isinstance(x, Bnb4bitParametrization), module.parametrizations[param_name]), None ) assert parametrization is not None, "Parametrization not found for the parameter." quant_state = parametrization.quant_state # Next, we need to store the quantization state. if quant_state is not None: for k, v in quant_state.as_dict(packed=True).items(): state_dict[f"{prefix}{param_name}.{k}"] = v bitsandbytes/bitsandbytes/nn/triton_based_modules.py000066400000000000000000000231321506752665600234610ustar00rootroot00000000000000from functools import partial import torch import torch.nn as nn from bitsandbytes.triton.dequantize_rowwise import dequantize_rowwise from bitsandbytes.triton.int8_matmul_mixed_dequantize import ( int8_matmul_mixed_dequantize, ) from bitsandbytes.triton.int8_matmul_rowwise_dequantize import ( int8_matmul_rowwise_dequantize, ) from bitsandbytes.triton.quantize_columnwise_and_transpose import ( quantize_columnwise_and_transpose, ) from bitsandbytes.triton.quantize_global import ( quantize_global, quantize_global_transpose, ) from bitsandbytes.triton.quantize_rowwise import quantize_rowwise from bitsandbytes.triton.triton_utils import is_triton_available class _switchback_global(torch.autograd.Function): @staticmethod def forward(ctx, X_3D, W, bias): # reshape input to [N * L, D] X = X_3D.view(-1, X_3D.size(-1)) # rowwise quantize for X, global quantize for W X_int8, state_X = quantize_rowwise(X) W_int8, state_W = quantize_global(W) # save for backward. ctx.save_for_backward = X, W # matmult, fused dequant and add bias # call "mixed" because we are mixing rowwise quantized and global quantized return int8_matmul_mixed_dequantize(X_int8, W_int8.t(), state_X, state_W, bias).view(*X_3D.size()[:-1], -1) @staticmethod def backward(ctx, G_3D): # reshape input to [N_out * L, D] G = G_3D.reshape(-1, G_3D.size(-1)) grad_X = grad_W = grad_bias = None X, W = ctx.save_for_backward if ctx.needs_input_grad[0]: # rowwise quantize for G, global quantize for W # for W, we also fuse the transpose operation because only A @ B^T is supported # so we transpose once then call .t() in the matmul G_int8, state_G = quantize_rowwise(G) W_int8, state_W = quantize_global_transpose(W) grad_X = int8_matmul_mixed_dequantize(G_int8, W_int8.t(), state_G, state_W, None).view( *G_3D.size()[:-1], -1, ) if ctx.needs_input_grad[1]: # backward pass uses standard weight grad grad_W = torch.matmul(G.t(), X.to(G.dtype)) if ctx.needs_input_grad[2]: grad_bias = G.sum(dim=0) return grad_X, grad_W, grad_bias class _switchback_vectorrize(torch.autograd.Function): @staticmethod def forward(ctx, X_3D, W, bias): # reshape input to [N * L, D] X = X_3D.view(-1, X_3D.size(-1)) ctx.save_for_backward = X, W # rowwise quantize for X # columnwise quantize for W (first rowwise, transpose later) X_int8, state_X = quantize_rowwise(X) W_int8, state_W = quantize_rowwise(W) # matmult, fused dequant and add bias # call kernel which expects rowwise quantized X and W return int8_matmul_rowwise_dequantize(X_int8, W_int8.t(), state_X, state_W, bias).view(*X_3D.size()[:-1], -1) @staticmethod def backward(ctx, G_3D): X, W = ctx.save_for_backward G = G_3D.reshape(-1, G_3D.size(-1)) grad_X = grad_W = grad_bias = None if ctx.needs_input_grad[0]: # rowwise quantize for G, columnwise quantize for W and fused transpose # we call .t() for weight later because only A @ B^T is supported G_int8, state_G = quantize_rowwise(G) W_int8, state_W = quantize_columnwise_and_transpose(W) grad_X = int8_matmul_rowwise_dequantize(G_int8, W_int8.t(), state_G, state_W, None).view( *G_3D.size()[:-1], -1, ) if ctx.needs_input_grad[1]: # backward pass uses standard weight grad grad_W = torch.matmul(G.t(), X.to(G.dtype)) if ctx.needs_input_grad[2]: grad_bias = G.sum(dim=0) return grad_X, grad_W, grad_bias class _switchback_global_mem_efficient(torch.autograd.Function): @staticmethod def forward(ctx, X_3D, W, bias): # reshape input to [N * L, D] X = X_3D.view(-1, X_3D.size(-1)) X_3D_sz = X_3D.size() # rowwise quantize for X, global quantize for W X_int8, state_X = quantize_rowwise(X) del X W_int8, state_W = quantize_global(W) # save for backward. ctx.save_for_backward = X_int8, state_X, W_int8, state_W # matmult, fused dequant and add bias # call "mixed" because we are mixing rowwise quantized and global quantized return int8_matmul_mixed_dequantize(X_int8, W_int8.t(), state_X, state_W, bias).view(*X_3D_sz[:-1], -1) @staticmethod def backward(ctx, G_3D): # reshape input to [N_out * L, D] G = G_3D.reshape(-1, G_3D.size(-1)) G_3D_sz = G_3D.size() grad_X = grad_W = grad_bias = None X_int8, state_X, W_int8, state_W = ctx.save_for_backward if ctx.needs_input_grad[1]: real_X = dequantize_rowwise(X_int8, state_X) del X_int8 grad_W = torch.matmul(G.t(), real_X.to(G.dtype)) del real_X if ctx.needs_input_grad[2]: grad_bias = G.sum(dim=0) if ctx.needs_input_grad[0]: G_int8, state_G = quantize_rowwise(G) del G W_int8 = W_int8.t().contiguous() grad_X = int8_matmul_mixed_dequantize(G_int8, W_int8.t(), state_G, state_W, None).view(*G_3D_sz[:-1], -1) return grad_X, grad_W, grad_bias class SwitchBackLinear(nn.Linear): def __init__( self, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None, vector_wise_quantization: bool = False, mem_efficient: bool = False, ): super().__init__(in_features, out_features, bias, device, dtype) if not is_triton_available(): raise ImportError("""Could not import triton. Please install triton to use SwitchBackLinear. Alternatively, you can use bnb.nn.SwitchBackLinearBnb, but it will be slower""") # By default, we use the global quantization. self.vector_wise_quantization = vector_wise_quantization if self.vector_wise_quantization: self._fn = _switchback_vectorrize if mem_efficient: print("mem efficient is not supported for vector-wise quantization.") exit(1) else: if mem_efficient: self._fn = _switchback_global_mem_efficient else: self._fn = _switchback_global def prepare_for_eval(self): # If we just want to do eval, we can pre-quantize the weights instead of doing it on the forward pass. # Note this is experimental and not tested thoroughly. # Note this needs to be explicitly called with something like # def cond_prepare(m): # if hasattr(m, "prepare_for_eval"): # m.prepare_for_eval() # model.apply(cond_prepare) print("=> preparing for eval.") if self.vector_wise_quantization: W_int8, state_W = quantize_rowwise(self.weight) else: W_int8, state_W = quantize_global(self.weight) self.register_buffer("W_int8", W_int8) self.register_buffer("state_W", state_W) del self.weight def forward(self, x): if self.training: return self._fn.apply(x, self.weight, self.bias) else: # If it hasn't been "prepared for eval", run the standard forward pass. if not hasattr(self, "W_int8"): return self._fn.apply(x, self.weight, self.bias) # Otherwise, use pre-computed weights. X = x.view(-1, x.size(-1)) X_int8, state_X = quantize_rowwise(X) if self.vector_wise_quantization: return int8_matmul_rowwise_dequantize(X_int8, self.W_int8.t(), state_X, self.state_W, self.bias).view( *x.size()[:-1], -1, ) else: return int8_matmul_mixed_dequantize(X_int8, self.W_int8.t(), state_X, self.state_W, self.bias).view( *x.size()[:-1], -1, ) SwitchBackLinearGlobal = partial(SwitchBackLinear, vector_wise_quantization=False) SwitchBackLinearGlobalMemEfficient = partial(SwitchBackLinear, vector_wise_quantization=False, mem_efficient=True) SwitchBackLinearVectorwise = partial(SwitchBackLinear, vector_wise_quantization=True) # This is just the standard linear function. class StandardLinearFunction(torch.autograd.Function): @staticmethod def forward(ctx, input, weight, bias=None): X = input.view(-1, input.size(-1)) ctx.save_for_backward(X, weight, bias) output = input.matmul(weight.t()) if bias is not None: output += bias.unsqueeze(0).expand_as(output) return output.view(*input.size()[:-1], -1) @staticmethod def backward(ctx, grad_output_3D): input, weight, bias = ctx.saved_tensors grad_output = grad_output_3D.reshape(-1, grad_output_3D.size(-1)) grad_input = grad_weight = grad_bias = None if ctx.needs_input_grad[0]: grad_input = grad_output.matmul(weight.to(grad_output.dtype)).view(*grad_output_3D.size()[:-1], -1) if ctx.needs_input_grad[1]: grad_weight = grad_output.t().matmul(input.to(grad_output.dtype)) if bias is not None and ctx.needs_input_grad[2]: grad_bias = grad_output.sum(0) return grad_input, grad_weight, grad_bias class StandardLinear(nn.Linear): def forward(self, x): return StandardLinearFunction.apply(x, self.weight, self.bias) bitsandbytes/bitsandbytes/optim/000077500000000000000000000000001506752665600174165ustar00rootroot00000000000000bitsandbytes/bitsandbytes/optim/__init__.py000066400000000000000000000015611506752665600215320ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from .adagrad import Adagrad, Adagrad8bit, Adagrad32bit from .adam import Adam, Adam8bit, Adam32bit, PagedAdam, PagedAdam8bit, PagedAdam32bit from .adamw import ( AdamW, AdamW8bit, AdamW32bit, PagedAdamW, PagedAdamW8bit, PagedAdamW32bit, ) from .ademamix import AdEMAMix, AdEMAMix8bit, AdEMAMix32bit, PagedAdEMAMix, PagedAdEMAMix8bit, PagedAdEMAMix32bit from .lamb import LAMB, LAMB8bit, LAMB32bit from .lars import LARS, LARS8bit, LARS32bit, PytorchLARS from .lion import Lion, Lion8bit, Lion32bit, PagedLion, PagedLion8bit, PagedLion32bit from .optimizer import GlobalOptimManager from .rmsprop import RMSprop, RMSprop8bit, RMSprop32bit from .sgd import SGD, SGD8bit, SGD32bit bitsandbytes/bitsandbytes/optim/adagrad.py000066400000000000000000000173531506752665600213640ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from bitsandbytes.optim.optimizer import Optimizer1State class Adagrad(Optimizer1State): def __init__( self, params, lr=1e-2, lr_decay=0, weight_decay=0, initial_accumulator_value=0, eps=1e-10, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ Base Adagrad optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-2): The learning rate. lr_decay (`int`, defaults to 0): The learning rate decay. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. initial_accumulator_value (`int`, defaults to 0): The initial momemtum values. eps (`float`, defaults to 1e-10): The epsilon value prevents division by zero in the optimizer. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ if not 0.0 <= lr: raise ValueError(f"Invalid learning rate: {lr}") if not 0.0 <= weight_decay: raise ValueError(f"Invalid weight_decay value: {weight_decay}") if not 0.0 <= eps: raise ValueError(f"Invalid epsilon value: {eps}") if initial_accumulator_value != 0.0: raise ValueError("Initial accumulator value != 0.0 not supported!") if lr_decay != 0.0: raise ValueError("Lr Decay != 0.0 not supported!") super().__init__( "adagrad", params, lr, (0.0, 0.0), eps, weight_decay, optim_bits, args, min_8bit_size, percentile_clipping, block_wise, ) class Adagrad8bit(Optimizer1State): def __init__( self, params, lr=1e-2, lr_decay=0, weight_decay=0, initial_accumulator_value=0, eps=1e-10, optim_bits=8, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ 8-bit Adagrad optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-2): The learning rate. lr_decay (`int`, defaults to 0): The learning rate decay. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. initial_accumulator_value (`int`, defaults to 0): The initial momemtum values. eps (`float`, defaults to 1e-10): The epsilon value prevents division by zero in the optimizer. optim_bits (`int`, defaults to 8): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ if not 0.0 <= lr: raise ValueError(f"Invalid learning rate: {lr}") if not 0.0 <= weight_decay: raise ValueError(f"Invalid weight_decay value: {weight_decay}") if not 0.0 <= eps: raise ValueError(f"Invalid epsilon value: {eps}") if initial_accumulator_value != 0.0: raise ValueError("Initial accumulator value != 0.0 not supported!") if lr_decay != 0.0: raise ValueError("Lr Decay != 0.0 not supported!") assert block_wise super().__init__( "adagrad", params, lr, (0.0, 0.0), eps, weight_decay, 8, args, min_8bit_size, percentile_clipping, block_wise, ) class Adagrad32bit(Optimizer1State): def __init__( self, params, lr=1e-2, lr_decay=0, weight_decay=0, initial_accumulator_value=0, eps=1e-10, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ 32-bit Adagrad optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-2): The learning rate. lr_decay (`int`, defaults to 0): The learning rate decay. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. initial_accumulator_value (`int`, defaults to 0): The initial momemtum values. eps (`float`, defaults to 1e-10): The epsilon value prevents division by zero in the optimizer. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ if not 0.0 <= lr: raise ValueError(f"Invalid learning rate: {lr}") if not 0.0 <= weight_decay: raise ValueError(f"Invalid weight_decay value: {weight_decay}") if not 0.0 <= eps: raise ValueError(f"Invalid epsilon value: {eps}") if initial_accumulator_value != 0.0: raise ValueError("Initial accumulator value != 0.0 not supported!") if lr_decay != 0.0: raise ValueError("Lr Decay != 0.0 not supported!") super().__init__( "adagrad", params, lr, (0.0, 0.0), eps, weight_decay, 32, args, min_8bit_size, percentile_clipping, block_wise, ) bitsandbytes/bitsandbytes/optim/adam.py000066400000000000000000000370621506752665600207020ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from bitsandbytes.optim.optimizer import Optimizer2State class Adam(Optimizer2State): def __init__( self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, is_paged=False, ): """ Base Adam optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ super().__init__( "adam", params, lr, betas, eps, weight_decay, optim_bits, args, min_8bit_size, percentile_clipping, block_wise, is_paged=is_paged, ) class Adam8bit(Optimizer2State): def __init__( self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, is_paged=False, ): """ 8-bit Adam optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. Note: This parameter is not supported in Adam8bit and must be False. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. Note: This parameter is not used in Adam8bit as it always uses 8-bit optimization. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ # Validate unsupported parameters if amsgrad: raise ValueError("Adam8bit does not support amsgrad=True") if optim_bits != 32: # We allow the default value of 32 to maintain compatibility with the function signature, # but any other value is invalid since Adam8bit always uses 8-bit optimization raise ValueError("Adam8bit only supports optim_bits=32 (default value for compatibility)") super().__init__( "adam", params, lr, betas, eps, weight_decay, 8, # Hardcoded to 8 bits args, min_8bit_size, percentile_clipping, block_wise, is_paged=is_paged, ) class Adam32bit(Optimizer2State): def __init__( self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, is_paged=False, ): """ 32-bit Adam optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ super().__init__( "adam", params, lr, betas, eps, weight_decay, 32, args, min_8bit_size, percentile_clipping, block_wise, is_paged=is_paged, ) class PagedAdam(Optimizer2State): def __init__( self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, is_paged=False, ): """ Paged Adam optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ super().__init__( "adam", params, lr, betas, eps, weight_decay, optim_bits, args, min_8bit_size, percentile_clipping, block_wise, is_paged=True, ) class PagedAdam8bit(Optimizer2State): def __init__( self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, is_paged=False, ): """ 8-bit paged Adam optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. Note: This parameter is not supported in PagedAdam8bit and must be False. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. Note: This parameter is not used in PagedAdam8bit as it always uses 8-bit optimization. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ # Validate unsupported parameters if amsgrad: raise ValueError("PagedAdam8bit does not support amsgrad=True") if optim_bits != 32: # We allow the default value of 32 to maintain compatibility with the function signature, # but any other value is invalid since PagedAdam8bit always uses 8-bit optimization raise ValueError("PagedAdam8bit only supports optim_bits=32 (default value for compatibility)") super().__init__( "adam", params, lr, betas, eps, weight_decay, 8, # Hardcoded to 8 bits args, min_8bit_size, percentile_clipping, block_wise, is_paged=True, ) class PagedAdam32bit(Optimizer2State): def __init__( self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, is_paged=False, ): """ Paged 32-bit Adam optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ super().__init__( "adam", params, lr, betas, eps, weight_decay, 32, args, min_8bit_size, percentile_clipping, block_wise, is_paged=True, ) bitsandbytes/bitsandbytes/optim/adamw.py000066400000000000000000000362631506752665600210730ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from bitsandbytes.optim.optimizer import Optimizer2State class AdamW(Optimizer2State): def __init__( self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-2, amsgrad=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, is_paged=False, ): """ Base AdamW optimizer. Arguments: params (`torch.Tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ super().__init__( "adam", params, lr, betas, eps, weight_decay, optim_bits, args, min_8bit_size, percentile_clipping, block_wise, is_paged=is_paged, ) class AdamW8bit(Optimizer2State): def __init__( self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-2, amsgrad=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, is_paged=False, ): """ 8-bit AdamW optimizer. Arguments: params (`torch.Tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. Note: This parameter is not supported in AdamW8bit and must be False. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. Note: This parameter is not used in AdamW8bit as it always uses 8-bit optimization. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ # Validate unsupported parameters if amsgrad: raise ValueError("AdamW8bit does not support amsgrad=True") if optim_bits != 32: # We allow the default value of 32 to maintain compatibility with the function signature, # but any other value is invalid since AdamW8bit always uses 8-bit optimization raise ValueError("AdamW8bit only supports optim_bits=32 (default value for compatibility)") super().__init__( "adam", params, lr, betas, eps, weight_decay, 8, # Hardcoded to 8 bits args, min_8bit_size, percentile_clipping, block_wise, is_paged=is_paged, ) class AdamW32bit(Optimizer2State): def __init__( self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-2, amsgrad=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, is_paged=False, ): """ 32-bit AdamW optimizer. Arguments: params (`torch.Tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ super().__init__( "adam", params, lr, betas, eps, weight_decay, 32, args, min_8bit_size, percentile_clipping, block_wise, is_paged=is_paged, ) class PagedAdamW(Optimizer2State): def __init__( self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-2, amsgrad=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ Paged AdamW optimizer. Arguments: params (`torch.Tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ super().__init__( "adam", params, lr, betas, eps, weight_decay, optim_bits, args, min_8bit_size, percentile_clipping, block_wise, is_paged=True, ) class PagedAdamW8bit(Optimizer2State): def __init__( self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-2, amsgrad=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ Paged 8-bit AdamW optimizer. Arguments: params (`torch.Tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. Note: This parameter is not supported in PagedAdamW8bit and must be False. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. Note: This parameter is not used in PagedAdamW8bit as it always uses 8-bit optimization. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ # Validate unsupported parameters if amsgrad: raise ValueError("PagedAdamW8bit does not support amsgrad=True") if optim_bits != 32: # We allow the default value of 32 to maintain compatibility with the function signature, # but any other value is invalid since PagedAdamW8bit always uses 8-bit optimization raise ValueError("PagedAdamW8bit only supports optim_bits=32 (default value for compatibility)") super().__init__( "adam", params, lr, betas, eps, weight_decay, 8, # Hardcoded to 8 bits args, min_8bit_size, percentile_clipping, block_wise, is_paged=True, ) class PagedAdamW32bit(Optimizer2State): def __init__( self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-2, amsgrad=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ Paged 32-bit AdamW optimizer. Arguments: params (`torch.Tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ super().__init__( "adam", params, lr, betas, eps, weight_decay, 32, args, min_8bit_size, percentile_clipping, block_wise, is_paged=True, ) bitsandbytes/bitsandbytes/optim/ademamix.py000066400000000000000000000313471506752665600215650ustar00rootroot00000000000000from collections.abc import Iterable import math from typing import Literal, Optional import torch import bitsandbytes.functional as F from bitsandbytes.optim.optimizer import Optimizer2State class _ReferenceAdEMAMix(torch.optim.Optimizer): """ Reference: https://hf.co/papers/2409.03137 """ def __init__( self, params: Iterable[torch.nn.Parameter], lr: float = 1e-3, betas: tuple[float, float, float] = (0.9, 0.999, 0.9999), alpha: float = 5.0, eps: float = 1e-8, weight_decay: float = 1e-2, # default 0.0 or 1e-2? t_beta3: Optional[int] = None, t_alpha: Optional[int] = None, ): defaults = dict( lr=lr, betas=betas, alpha=alpha, eps=eps, weight_decay=weight_decay, t_beta3=t_beta3, t_alpha=t_alpha ) super().__init__(params, defaults) @torch.no_grad() def step(self, closure=None): loss = None if closure is not None: with torch.enable_grad(): loss = closure() for group in self.param_groups: if "step" in group: group["step"] += 1 else: group["step"] = 1 lr = group["lr"] eps = group["eps"] beta1, beta2, beta3 = group["betas"] alpha = group["alpha"] t_alpha = group["t_alpha"] t_beta3 = group["t_beta3"] weight_decay = group["weight_decay"] for p in group["params"]: if p.grad is None: continue grad = p.grad state = self.state[p] # State initialization if len(state) == 0: # For parity with bnb implementation we combine both fast # and slow EMA stats into one stacked tensor. state["m1_m2"] = p.new_zeros((2, *p.size())) state["nu"] = torch.zeros_like(p) # second moment estimate m1, m2, nu = state["m1_m2"][0], state["m1_m2"][1], state["nu"] bias_correction1 = 1 - beta1 ** group["step"] bias_correction2 = 1 - beta2 ** group["step"] # Apply scheduler for alpha if t_alpha is not None: alpha = min(group["step"] * alpha / t_alpha, alpha) # Apply scheduler for beta3 if t_beta3 is not None: ln_beta1 = math.log(beta1) ln_beta3 = math.log(beta3) step_scale = group["step"] / t_beta3 beta3 = min( math.exp((ln_beta1 * ln_beta3) / (((1 - step_scale) * ln_beta3) + (step_scale * ln_beta1))), beta3, ) # Update the EMAs m1.mul_(beta1).add_(grad, alpha=1 - beta1) m2.mul_(beta3).add_(grad, alpha=1 - beta3) nu.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) # Compute step denom = (nu.sqrt() / (bias_correction2**0.5)).add(eps) update = (m1.div(bias_correction1) + alpha * m2) / denom # Add weight decay update.add_(p, alpha=weight_decay) # Apply update scaled by learning rate p.add_(-lr * update) return loss class AdEMAMix(Optimizer2State): def __init__( self, params: Iterable[torch.nn.Parameter], lr: float = 1e-3, betas: tuple[float, float, float] = (0.9, 0.999, 0.9999), alpha: float = 5.0, t_alpha: Optional[int] = None, t_beta3: Optional[int] = None, eps: float = 1e-8, weight_decay: float = 1e-2, optim_bits: Literal[8, 32] = 32, min_8bit_size: int = 4096, is_paged: bool = False, ): super().__init__( "ademamix", params=params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, optim_bits=optim_bits, args=None, min_8bit_size=min_8bit_size, percentile_clipping=100, block_wise=True, is_paged=is_paged, alpha=alpha, t_alpha=t_alpha, t_beta3=t_beta3, ) @torch.no_grad() def init_state(self, group, p, gindex, pindex): # In our AdEMAMix implementation, we use `state` to hold # both the fast and slow EMAs. Here we override the base # `Optimizer2State` to allocate a buffer twice as large. # Additional consideration: we do not support block_wise=False, # percentile clipping, or max_unorm. config = self.get_config(gindex, pindex, group) if config["optim_bits"] == 32: dtype = torch.float32 elif config["optim_bits"] == 8: dtype = torch.uint8 else: raise NotImplementedError(f"Amount of optimizer bits not supported: {config['optim_bits']}") if p.numel() < config["min_8bit_size"]: dtype = torch.float32 state = self.state[p] state["step"] = 0 if dtype == torch.uint8: if "dynamic" not in self.name2qmap: self.fill_qmap() self.name2qmap["dynamic"] = state["qmap1"] = self.name2qmap["dynamic"].to(p.device) self.name2qmap["udynamic"] = state["qmap2"] = self.name2qmap["udynamic"].to(p.device) blocksize = 256 n = p.numel() blocks = (n // blocksize) + bool(n % blocksize) state["absmax1"] = torch.zeros((2, blocks), dtype=torch.float32, device=p.device) state["absmax2"] = torch.zeros((blocks,), dtype=torch.float32, device=p.device) state["state1"] = self._get_state_double_buffer(p, dtype=dtype) state["state2"] = self.get_state_buffer(p, dtype=dtype) @torch.no_grad() def update_step(self, group, p, gindex, pindex): config = self.get_config(gindex, pindex, group) if config["t_alpha"] is None and config["t_beta3"] is None: # Not using alpha/beta3 scheduler; we can fall through. super().update_step(group, p, gindex, pindex) return # Ensure contiguous memory layout p.data = p.data.contiguous() p.grad = p.grad.contiguous() state = self.state[p] grad = p.grad state["step"] += 1 step = state["step"] beta1, beta2, beta3 = config["betas"] alpha = config["alpha"] t_alpha = config["t_alpha"] t_beta3 = config["t_beta3"] # Apply scheduler for alpha if t_alpha is not None: alpha_t = min(step * alpha / t_alpha, alpha) else: alpha_t = alpha # Apply scheduler for beta3 if t_beta3 is not None: ln_beta1 = math.log(beta1) ln_beta3 = math.log(beta3) step_scale = step / t_beta3 beta3_t = min( math.exp((ln_beta1 * ln_beta3) / (((1 - step_scale) * ln_beta3) + (step_scale * ln_beta1))), beta3 ) else: beta3_t = beta3 # Apply updates if state["state1"].dtype == torch.float32: F.optimizer_update_32bit( self.optimizer_name, grad, p, state["state1"], beta1, config["eps"], step, config["lr"], state["state2"], beta2, beta3_t, alpha_t, config["weight_decay"], gnorm_scale=1.0, unorm_vec=state["unorm_vec"] if config["max_unorm"] > 0.0 else None, max_unorm=config["max_unorm"], skip_zeros=config["skip_zeros"], ) elif state["state1"].dtype == torch.uint8: F.optimizer_update_8bit_blockwise( self.optimizer_name, grad, p, state["state1"], state["state2"], config["betas"][0], config["betas"][1], beta3_t, alpha_t, config["eps"], step, config["lr"], state["qmap1"], state["qmap2"], state["absmax1"], state["absmax2"], config["weight_decay"], gnorm_scale=1.0, skip_zeros=config["skip_zeros"], ) def _get_state_double_buffer(self, p, dtype=torch.float32): if not self.is_paged or p.numel() < 1e5: return torch.zeros((2, *p.size()), dtype=dtype, device=p.device) else: buff = F.get_paged(*(2, *p.size()), dtype=dtype, device=p.device) F.fill(buff, 0) self.page_mng.paged_tensors.append(buff) return buff class AdEMAMix8bit(AdEMAMix): def __init__( self, params: Iterable[torch.nn.Parameter], lr: float = 1e-3, betas: tuple[float, float, float] = (0.9, 0.999, 0.9999), alpha: float = 5.0, t_alpha: Optional[int] = None, t_beta3: Optional[int] = None, eps: float = 1e-8, weight_decay: float = 1e-2, min_8bit_size: int = 4096, is_paged: bool = False, ): super().__init__( params, lr=lr, betas=betas, alpha=alpha, t_alpha=t_alpha, t_beta3=t_beta3, eps=eps, weight_decay=weight_decay, optim_bits=8, min_8bit_size=min_8bit_size, is_paged=is_paged, ) class PagedAdEMAMix8bit(AdEMAMix8bit): def __init__( self, params: Iterable[torch.nn.Parameter], lr: float = 1e-3, betas: tuple[float, float, float] = (0.9, 0.999, 0.9999), alpha: float = 5.0, t_alpha: Optional[int] = None, t_beta3: Optional[int] = None, eps: float = 1e-8, weight_decay: float = 1e-2, min_8bit_size: int = 4096, ): super().__init__( params, lr=lr, betas=betas, alpha=alpha, t_alpha=t_alpha, t_beta3=t_beta3, eps=eps, weight_decay=weight_decay, min_8bit_size=min_8bit_size, is_paged=True, ) class PagedAdEMAMix(AdEMAMix): def __init__( self, params: Iterable[torch.nn.Parameter], lr: float = 1e-3, betas: tuple[float, float, float] = (0.9, 0.999, 0.9999), alpha: float = 5.0, t_alpha: Optional[int] = None, t_beta3: Optional[int] = None, eps: float = 1e-8, weight_decay: float = 1e-2, optim_bits: Literal[8, 32] = 32, min_8bit_size: int = 4096, ): super().__init__( params, lr=lr, betas=betas, alpha=alpha, t_alpha=t_alpha, t_beta3=t_beta3, eps=eps, weight_decay=weight_decay, optim_bits=optim_bits, min_8bit_size=min_8bit_size, is_paged=True, ) class AdEMAMix32bit(Optimizer2State): def __init__( self, params: Iterable[torch.nn.Parameter], lr: float = 1e-3, betas: tuple[float, float, float] = (0.9, 0.999, 0.9999), alpha: float = 5.0, t_alpha: Optional[int] = None, t_beta3: Optional[int] = None, eps: float = 1e-8, weight_decay: float = 1e-2, min_8bit_size: int = 4096, is_paged: bool = False, ): super().__init__( "ademamix", params=params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, optim_bits=32, args=None, min_8bit_size=min_8bit_size, percentile_clipping=100, block_wise=True, is_paged=is_paged, alpha=alpha, t_alpha=t_alpha, t_beta3=t_beta3, ) class PagedAdEMAMix32bit(AdEMAMix32bit): def __init__( self, params: Iterable[torch.nn.Parameter], lr: float = 1e-3, betas: tuple[float, float, float] = (0.9, 0.999, 0.9999), alpha: float = 5.0, t_alpha: Optional[int] = None, t_beta3: Optional[int] = None, eps: float = 1e-8, weight_decay: float = 1e-2, min_8bit_size: int = 4096, ): super().__init__( params, lr=lr, betas=betas, alpha=alpha, t_alpha=t_alpha, t_beta3=t_beta3, eps=eps, weight_decay=weight_decay, min_8bit_size=min_8bit_size, is_paged=True, ) bitsandbytes/bitsandbytes/optim/lamb.py000066400000000000000000000174311506752665600207110ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from bitsandbytes.optim.optimizer import Optimizer2State class LAMB(Optimizer2State): def __init__( self, params, lr=1e-3, bias_correction=True, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False, adam_w_mode=True, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=False, max_unorm=1.0, ): """ Base LAMB optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. bias_correction (`bool`, defaults to `True`): Whether to apply bias correction to the first and second-order moments. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. adam_w_mode (`bool`, defaults to `True`): Whether to use the AdamW variant. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. max_unorm (`float`, defaults to 1.0): The maximum gradient norm. """ super().__init__( "lamb", params, lr, betas, eps, weight_decay, optim_bits, args, min_8bit_size, percentile_clipping, block_wise, max_unorm=1.0, ) class LAMB8bit(Optimizer2State): def __init__( self, params, lr=1e-3, bias_correction=True, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False, adam_w_mode=True, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=False, max_unorm=1.0, ): """ 8-bit LAMB optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. bias_correction (`bool`, defaults to `True`): Whether to apply bias correction to the first and second-order moments. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. adam_w_mode (`bool`, defaults to `True`): Whether to use the AdamW variant. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. max_unorm (`float`, defaults to 1.0): The maximum gradient norm. """ super().__init__( "lamb", params, lr, betas, eps, weight_decay, 8, args, min_8bit_size, percentile_clipping, block_wise, max_unorm=1.0, ) class LAMB32bit(Optimizer2State): def __init__( self, params, lr=1e-3, bias_correction=True, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False, adam_w_mode=True, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=False, max_unorm=1.0, ): """ 32-bit LAMB optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. bias_correction (`bool`, defaults to `True`): Whether to apply bias correction to the first and second-order moments. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. amsgrad (`bool`, defaults to `False`): Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead. adam_w_mode (`bool`, defaults to `True`): Whether to use the AdamW variant. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. max_unorm (`float`, defaults to 1.0): The maximum gradient norm. """ super().__init__( "lamb", params, lr, betas, eps, weight_decay, 32, args, min_8bit_size, percentile_clipping, block_wise, max_unorm=1.0, ) bitsandbytes/bitsandbytes/optim/lars.py000066400000000000000000000221141506752665600207310ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. import torch from torch.optim import Optimizer from bitsandbytes.optim.optimizer import Optimizer1State class LARS(Optimizer1State): def __init__( self, params, lr, momentum=0, dampening=0, weight_decay=0, nesterov=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, max_unorm=0.02, ): """ Base LARS optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`): The learning rate. momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. dampening (`float`, defaults to 0): The dampening value reduces the momentum of the optimizer. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. nesterov (`bool`, defaults to `False`): Whether to use Nesterov momentum. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. max_unorm (`float`, defaults to 0.02): The maximum gradient norm. """ if momentum == 0: raise NotImplementedError("LARS without momentum is not supported!") super().__init__( "lars", params, lr, (momentum, dampening), 0.0, weight_decay, optim_bits, args, min_8bit_size, percentile_clipping, max_unorm=max_unorm, block_wise=False, ) class LARS8bit(Optimizer1State): def __init__( self, params, lr, momentum=0, dampening=0, weight_decay=0, nesterov=False, args=None, min_8bit_size=4096, percentile_clipping=100, max_unorm=0.02, ): """ 8-bit LARS optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`): The learning rate. momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. dampening (`float`, defaults to 0): The dampening value reduces the momentum of the optimizer. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. nesterov (`bool`, defaults to `False`): Whether to use Nesterov momentum. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. max_unorm (`float`, defaults to 0.02): The maximum gradient norm. """ if momentum == 0: raise NotImplementedError("LARS without momentum is not supported!") super().__init__( "lars", params, lr, (momentum, dampening), 0.0, weight_decay, 8, args, min_8bit_size, percentile_clipping, max_unorm=max_unorm, block_wise=False, ) class LARS32bit(Optimizer1State): def __init__( self, params, lr, momentum=0, dampening=0, weight_decay=0, nesterov=False, args=None, min_8bit_size=4096, percentile_clipping=100, max_unorm=0.02, ): """ 32-bit LARS optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`): The learning rate. momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. dampening (`float`, defaults to 0): The dampening value reduces the momentum of the optimizer. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. nesterov (`bool`, defaults to `False`): Whether to use Nesterov momentum. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. max_unorm (`float`, defaults to 0.02): The maximum gradient norm. """ if momentum == 0: raise NotImplementedError("LARS without momentum is not supported!") super().__init__( "lars", params, lr, (momentum, dampening), 0.0, weight_decay, 32, args, min_8bit_size, percentile_clipping, max_unorm=max_unorm, block_wise=False, ) class PytorchLARS(Optimizer): def __init__( self, params, lr=0.01, momentum=0, dampening=0, weight_decay=0, nesterov=False, max_unorm=0.02, ): if lr < 0.0: raise ValueError(f"Invalid learning rate: {lr}") if momentum < 0.0: raise ValueError(f"Invalid momentum value: {momentum}") if weight_decay < 0.0: raise ValueError(f"Invalid weight_decay value: {weight_decay}") defaults = dict( lr=lr, momentum=momentum, dampening=dampening, weight_decay=weight_decay, nesterov=nesterov, max_unorm=max_unorm, ) if nesterov and (momentum <= 0 or dampening != 0): raise ValueError("Nesterov momentum requires a momentum and zero dampening") super().__init__(params, defaults) def __setstate__(self, state): super().__setstate__(state) for group in self.param_groups: group.setdefault("nesterov", False) @torch.no_grad() def step(self, closure=None): """Performs a single optimization step. Args: closure (callable, optional): A closure that reevaluates the model and returns the loss. """ loss = None if closure is not None: with torch.enable_grad(): loss = closure() for group in self.param_groups: weight_decay = group["weight_decay"] momentum = group["momentum"] dampening = group["dampening"] nesterov = group["nesterov"] max_unorm = group["max_unorm"] lr = group["lr"] for p in group["params"]: if p.grad is None: continue state = self.state[p] d_p = p.grad if weight_decay != 0: d_p = d_p.add(p, alpha=weight_decay) if momentum != 0: buf = state.get("momentum_buffer", None) if buf is None: buf = torch.clone(d_p).detach() state["momentum_buffer"] = buf else: buf.mul_(momentum).add_(d_p, alpha=1 - dampening) if nesterov: update = d_p + buf * momentum else: update = buf update_scale = 1.0 if max_unorm > 0.0: assert p.dtype == torch.float32 pnorm = torch.norm(p.detach()) unorm = torch.norm(update) if unorm > max_unorm * pnorm: update_scale = max_unorm * pnorm / unorm p.add_(update, alpha=-lr * update_scale) return loss bitsandbytes/bitsandbytes/optim/lion.py000066400000000000000000000265351506752665600207440ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from bitsandbytes.optim.optimizer import Optimizer1State class Lion(Optimizer1State): def __init__( self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, is_paged=False, ): """ Base Lion optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-4): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. weight_decay (`float`, defaults to 0): The weight decay value for the optimizer. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ super().__init__( "lion", params, lr, betas, 0.0, weight_decay, optim_bits, args, min_8bit_size, percentile_clipping, block_wise, is_paged=is_paged, ) class Lion8bit(Optimizer1State): def __init__( self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, is_paged=False, ): """ 8-bit Lion optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-4): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. weight_decay (`float`, defaults to 0): The weight decay value for the optimizer. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ super().__init__( "lion", params, lr, betas, 0.0, weight_decay, 8, args, min_8bit_size, percentile_clipping, block_wise, is_paged=is_paged, ) class Lion32bit(Optimizer1State): def __init__( self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, is_paged=False, ): """ 32-bit Lion optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-4): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. weight_decay (`float`, defaults to 0): The weight decay value for the optimizer. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ super().__init__( "lion", params, lr, betas, 0.0, weight_decay, 32, args, min_8bit_size, percentile_clipping, block_wise, is_paged=is_paged, ) class PagedLion(Optimizer1State): def __init__( self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ Paged Lion optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-4): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. weight_decay (`float`, defaults to 0): The weight decay value for the optimizer. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ super().__init__( "lion", params, lr, betas, 0.0, weight_decay, optim_bits, args, min_8bit_size, percentile_clipping, block_wise, is_paged=True, ) class PagedLion8bit(Optimizer1State): def __init__( self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ Paged 8-bit Lion optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-4): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. weight_decay (`float`, defaults to 0): The weight decay value for the optimizer. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ super().__init__( "lion", params, lr, betas, 0.0, weight_decay, 8, args, min_8bit_size, percentile_clipping, block_wise, is_paged=True, ) class PagedLion32bit(Optimizer1State): def __init__( self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ Paged 32-bit Lion optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-4): The learning rate. betas (`tuple(float, float)`, defaults to (0.9, 0.999)): The beta values are the decay rates of the first and second-order moment of the optimizer. weight_decay (`float`, defaults to 0): The weight decay value for the optimizer. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ super().__init__( "lion", params, lr, betas, 0.0, weight_decay, 32, args, min_8bit_size, percentile_clipping, block_wise, is_paged=True, ) bitsandbytes/bitsandbytes/optim/optimizer.py000066400000000000000000000730611506752665600220210ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from collections import abc as container_abcs, defaultdict from copy import deepcopy from itertools import chain from typing import Optional import torch import bitsandbytes.functional as F from bitsandbytes.utils import sync_gpu class MockArgs: def __init__(self, initial_data): for key in initial_data: setattr(self, key, initial_data[key]) class GlobalOptimManager: """ A global optimizer manager for enabling custom optimizer configs. """ _instance = None def __init__(self): raise RuntimeError("Call get_instance() instead") def initialize(self): self.pid2config = {} self.index2config = {} self.optimizer = None self.uses_config_override = False self.module_weight_config_triple = [] @classmethod def get_instance(cls): if cls._instance is None: cls._instance = cls.__new__(cls) cls._instance.initialize() return cls._instance def register_parameters(self, params): param_groups = list(params) if not isinstance(param_groups[0], dict): param_groups = [{"params": param_groups}] for group_index, group in enumerate(param_groups): for p_index, p in enumerate(group["params"]): if id(p) in self.pid2config: self.index2config[(group_index, p_index)] = self.pid2config[id(p)] def override_config(self, parameters, key=None, value=None, key_value_dict=None): """ Override initial optimizer config with specific hyperparameters. The key-values of the optimizer config for the input parameters are overridden This can be both, optimizer parameters like `betas` or `lr`, or it can be 8-bit specific parameters like `optim_bits` or `percentile_clipping`. Arguments: parameters (`torch.Tensor` or `list(torch.Tensors)`): The input parameters. key (`str`): The hyperparameter to override. value: The hyperparameter value. key_value_dict (`dict`): A dictionary with multiple key-values to override. Example: ```py import torch import bitsandbytes as bnb mng = bnb.optim.GlobalOptimManager.get_instance() model = MyModel() mng.register_parameters(model.parameters()) # 1. register parameters while still on CPU model = model.cuda() # use 8-bit optimizer states for all parameters adam = bnb.optim.Adam(model.parameters(), lr=0.001, optim_bits=8) # 2. override: the parameter model.fc1.weight now uses 32-bit Adam mng.override_config(model.fc1.weight, 'optim_bits', 32) ``` """ self.uses_config_override = True if isinstance(parameters, torch.nn.Parameter): parameters = [parameters] if isinstance(parameters, torch.Tensor): parameters = [parameters] if key is not None and value is not None: assert key_value_dict is None key_value_dict = {key: value} if key_value_dict is not None: for p in parameters: if id(p) in self.pid2config: self.pid2config[id(p)].update(key_value_dict) else: self.pid2config[id(p)] = key_value_dict def register_module_override(self, module, param_name, config): self.module_weight_config_triple.append((module, param_name, config)) class Optimizer8bit(torch.optim.Optimizer): def __init__(self, params, defaults, optim_bits=32, is_paged=False): """ Base 8-bit optimizer class. Arguments: params (`torch.Tensor`): The input parameters to optimize. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ super().__init__(params, defaults) self.initialized = False self.name2qmap = {} self.is_paged = is_paged self.page_mng = F.GlobalPageManager.get_instance() self.mng = GlobalOptimManager.get_instance() self.non_castable_tensor_keys = { "qmap1", "qmap2", "max1", "max2", "new_max1", "new_max2", "state1", "state2", "gnorm_vec", "absmax1", "absmax2", "unorm_vec", } if optim_bits == 8: self.fill_qmap() def fill_qmap(self): self.name2qmap["dynamic"] = F.create_dynamic_map(signed=True) self.name2qmap["udynamic"] = F.create_dynamic_map(signed=False) def __setstate__(self, state): super().__setstate__(state) def load_state_dict(self, state_dict, move_to_device=True): """Load an optimizer state. Arguments: state_dict (`dict`): An optimizer state (should be returned from a call to `state_dict`) to load. move_to_device (`bool`, defaults to `True`): Whether to move the optimizer's state to the device. """ # deepcopy, to be consistent with module API state_dict = deepcopy(state_dict) # Validate the state_dict groups = self.param_groups saved_groups = state_dict["param_groups"] if len(groups) != len(saved_groups): raise ValueError("loaded state dict has a different number of parameter groups") param_lens = (len(g["params"]) for g in groups) saved_lens = (len(g["params"]) for g in saved_groups) if any(p_len != s_len for p_len, s_len in zip(param_lens, saved_lens)): raise ValueError( "loaded state dict contains a parameter group that doesn't match the size of optimizer's group", ) # Update the state id_map = { old_id: p for old_id, p in zip( chain.from_iterable(g["params"] for g in saved_groups), chain.from_iterable(g["params"] for g in groups), ) } def cast(param, value): r"""Make a deep copy of value, casting all tensors to device of param.""" if isinstance(value, torch.Tensor): # Floating-point types are a bit special here. They are the only ones # that are assumed to always match the type of params. if param.is_floating_point() and value.dtype != torch.uint8: value = value.to(param.dtype) return value elif isinstance(value, dict): for k, v in value.items(): if k in self.non_castable_tensor_keys: if move_to_device: value[k] = v.to(param.device) else: value[k] = cast(param, v) return value elif isinstance(value, container_abcs.Iterable): return type(value)(cast(param, v) for v in value) else: return value # Copy state assigned to params (and cast tensors to appropriate types). # State that is not assigned to params is copied as is (needed for # backward compatibility). state = defaultdict(dict) for k, v in state_dict["state"].items(): if k in id_map: param = id_map[k] state[param] = cast(param, v) else: state[k] = v # Update parameter groups, setting their 'params' value def update_group(group, new_group): new_group["params"] = group["params"] return new_group param_groups = [update_group(g, ng) for g, ng in zip(groups, saved_groups)] self.__setstate__({"state": state, "param_groups": param_groups}) def to_gpu(self): for gindex, group in enumerate(self.param_groups): for pindex, p in enumerate(group["params"]): if p in self.state: values = self.state[p] for k, v in values.items(): if isinstance(v, torch.Tensor): is_paged = getattr(v, "is_paged", False) if not is_paged: self.state[p][k] = v.to(p.device) def check_overrides(self): for module, attr, config in self.mng.module_weight_config_triple: pmodule = getattr(module, attr) assert pmodule is not None assert isinstance(pmodule, torch.Tensor) or isinstance(pmodule, torch.Parameter) found = False for gindex, group in enumerate(self.param_groups): if found: break for pindex, p in enumerate(group["params"]): if found: break if id(p) == id(pmodule): # found the matching parameter # init override self.mng.pid2config[id(p)] = config self.mng.index2config[(gindex, pindex)] = self.mng.pid2config[id(p)] found = True @torch.no_grad() def step(self, closure=None): """Perform a single optimization step. Arguments: closure (`Callable`, *optional*, defaults to `None`): A closure that reevaluates the model and returns the loss. """ loss = None if closure is not None: with torch.enable_grad(): loss = closure() if not self.initialized: self.check_overrides() self.to_gpu() # needed for fairseq pure fp16 training self.initialized = True # if self.is_paged: self.page_mng.prefetch_all() p = None for gindex, group in enumerate(self.param_groups): for pindex, p in enumerate(group["params"]): if p.grad is None: continue state = self.state[p] if len(state) == 0: self.init_state(group, p, gindex, pindex) self.prefetch_state(p) self.update_step(group, p, gindex, pindex) sync_gpu(p) if self.is_paged and p is not None: # all paged operations are asynchronous, we need # to sync to make sure all tensors are in the right state sync_gpu(p) return loss def get_config(self, gindex, pindex, group): config = {} config["betas"] = group["betas"] config["eps"] = group["eps"] config["weight_decay"] = group["weight_decay"] config["lr"] = group["lr"] config["alpha"] = group.get("alpha", 0.0) config["t_alpha"] = group.get("t_alpha", 0) config["t_beta3"] = group.get("t_beta3", 0) config["optim_bits"] = self.args.optim_bits config["min_8bit_size"] = self.args.min_8bit_size config["percentile_clipping"] = self.args.percentile_clipping config["block_wise"] = self.args.block_wise config["max_unorm"] = self.args.max_unorm config["skip_zeros"] = self.args.skip_zeros if (gindex, pindex) in self.mng.index2config: config.update(self.mng.index2config[(gindex, pindex)]) return config def init_state(self, group, p, gindex, pindex): raise NotImplementedError("init_state method needs to be overridden") def update_step(self, group, p, gindex, pindex): raise NotImplementedError("The update_step method needs to be overridden") def get_state_buffer(self, p, dtype=torch.float32): if not self.is_paged or p.numel() < 1e5: return torch.zeros_like(p, dtype=dtype, device=p.device) else: # > 1 MB buff = F.get_paged(*p.shape, dtype=dtype, device=p.device) F.fill(buff, 0) self.page_mng.paged_tensors.append(buff) return buff def prefetch_state(self, p): if self.is_paged: state = self.state[p] s1 = state["state1"] is_paged = getattr(s1, "is_paged", False) if is_paged: F.prefetch_tensor(state["state1"]) if "state2" in state: F.prefetch_tensor(state["state2"]) class Optimizer2State(Optimizer8bit): def __init__( self, optimizer_name, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, max_unorm=0.0, skip_zeros=False, is_paged=False, alpha=0.0, t_alpha: Optional[int] = None, t_beta3: Optional[int] = None, ): """ Base 2-state update optimizer class. Arguments: optimizer_name (`str`): The name of the optimizer. params (`torch.Tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple`, defaults to (0.9, 0.999)): The beta values for the optimizer. eps (`float`, defaults to 1e-8): The epsilon value for the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. max_unorm (`float`, defaults to 0.0): The maximum value to normalize each block with. skip_zeros (`bool`, defaults to `False`): Whether to skip zero values for sparse gradients and models to ensure correct updates. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. alpha (`float`, defaults to 0.0): The alpha value for the AdEMAMix optimizer. t_alpha (`Optional[int]`, defaults to `None`): Number of iterations for alpha scheduling with AdEMAMix. t_beta3 (`Optional[int]`, defaults to `None`): Number of iterations for beta scheduling with AdEMAMix. """ if not 0.0 <= lr: raise ValueError(f"Invalid learning rate: {lr}") if not 0.0 <= eps: raise ValueError(f"Invalid epsilon value: {eps}") if isinstance(betas, str): # format: '(beta1, beta2)' betas = betas.replace("(", "").replace(")", "").strip().split(",") betas = [float(b) for b in betas] for i in range(len(betas)): if not 0.0 <= betas[i] < 1.0: raise ValueError(f"Invalid beta parameter at index {i}: {betas[i]}") if not 0.0 <= weight_decay: raise ValueError(f"Invalid weight_decay value: {weight_decay}") defaults = dict( lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, alpha=alpha, t_alpha=t_alpha, t_beta3=t_beta3 ) super().__init__(params, defaults, optim_bits, is_paged) if args is None: args = {} args["optim_bits"] = optim_bits args["min_8bit_size"] = min_8bit_size args["percentile_clipping"] = percentile_clipping args["block_wise"] = block_wise args["max_unorm"] = max_unorm args["skip_zeros"] = skip_zeros self.args = MockArgs(args) else: self.args = args self.optimizer_name = optimizer_name @torch.no_grad() def init_state(self, group, p, gindex, pindex): config = self.get_config(gindex, pindex, group) if config["optim_bits"] == 32: dtype = torch.float32 elif config["optim_bits"] == 8: dtype = torch.uint8 else: raise NotImplementedError(f"Amount of optimizer bits not supported: {config['optim_bits']}") if p.numel() < config["min_8bit_size"]: dtype = torch.float32 state = self.state[p] state["step"] = 0 if dtype == torch.float32: state["state1"] = self.get_state_buffer(p, dtype=torch.float32) state["state2"] = self.get_state_buffer(p, dtype=torch.float32) elif dtype == torch.uint8: if state["step"] == 0: if "dynamic" not in self.name2qmap: self.fill_qmap() self.name2qmap["dynamic"] = self.name2qmap["dynamic"].to(p.device) self.name2qmap["udynamic"] = self.name2qmap["udynamic"].to(p.device) state["state1"] = self.get_state_buffer(p, dtype=torch.uint8) state["qmap1"] = self.name2qmap["dynamic"] state["state2"] = self.get_state_buffer(p, dtype=torch.uint8) state["qmap2"] = self.name2qmap["udynamic"] if config["block_wise"]: blocksize = 256 n = p.numel() blocks = (n // blocksize) + bool(n % blocksize) state["absmax1"] = torch.zeros((blocks,), dtype=torch.float32, device=p.device) state["absmax2"] = torch.zeros((blocks,), dtype=torch.float32, device=p.device) else: state["max1"] = torch.zeros((1,), dtype=torch.float32, device=p.device) state["new_max1"] = torch.zeros((1,), dtype=torch.float32, device=p.device) state["max2"] = torch.zeros((1,), dtype=torch.float32, device=p.device) state["new_max2"] = torch.zeros((1,), dtype=torch.float32, device=p.device) if config["percentile_clipping"] < 100: state["gnorm_vec"] = torch.zeros((100,), device=p.device) if config["max_unorm"] > 0.0: state["unorm_vec"] = torch.zeros((1,), device=p.device) @torch.no_grad() def update_step(self, group, p, gindex, pindex): # avoid update error from non-contiguous memory layout p.data = p.data.contiguous() p.grad = p.grad.contiguous() state = self.state[p] grad = p.grad config = self.get_config(gindex, pindex, group) state["step"] += 1 step = state["step"] if config["percentile_clipping"] < 100: current_gnorm, clip_value, gnorm_scale = F.percentile_clipping( grad, state["gnorm_vec"], step, config["percentile_clipping"], ) else: gnorm_scale = 1.0 if state["state1"].dtype == torch.float: F.optimizer_update_32bit( self.optimizer_name, grad, p, state["state1"], config["betas"][0], config["eps"], step, config["lr"], state["state2"], config["betas"][1], config["betas"][2] if len(config["betas"]) >= 3 else 0.0, config.get("alpha", 0.0), config["weight_decay"], gnorm_scale, state["unorm_vec"] if config["max_unorm"] > 0.0 else None, max_unorm=config["max_unorm"], skip_zeros=config["skip_zeros"], ) elif state["state1"].dtype == torch.uint8 and not config["block_wise"]: F.optimizer_update_8bit( self.optimizer_name, grad, p, state["state1"], state["state2"], config["betas"][0], config["betas"][1], config["eps"], step, config["lr"], state["qmap1"], state["qmap2"], state["max1"], state["max2"], state["new_max1"], state["new_max2"], config["weight_decay"], gnorm_scale=gnorm_scale, unorm_vec=state["unorm_vec"] if config["max_unorm"] > 0.0 else None, max_unorm=config["max_unorm"], ) # swap maxes state["max1"], state["new_max1"] = state["new_max1"], state["max1"] state["max2"], state["new_max2"] = state["new_max2"], state["max2"] elif state["state1"].dtype == torch.uint8 and config["block_wise"]: F.optimizer_update_8bit_blockwise( self.optimizer_name, grad, p, state["state1"], state["state2"], config["betas"][0], config["betas"][1], config["betas"][2] if len(config["betas"]) >= 3 else 0.0, config.get("alpha", 0.0), config["eps"], step, config["lr"], state["qmap1"], state["qmap2"], state["absmax1"], state["absmax2"], config["weight_decay"], gnorm_scale=gnorm_scale, skip_zeros=config["skip_zeros"], ) class Optimizer1State(Optimizer8bit): def __init__( self, optimizer_name, params, lr=1e-3, betas=(0.9, 0.0), eps=1e-8, weight_decay=0.0, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, max_unorm=0.0, skip_zeros=False, is_paged=False, ): """ Base 1-state update optimizer class. Arguments: optimizer_name (`str`): The name of the optimizer. params (`torch.Tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-3): The learning rate. betas (`tuple`, defaults to (0.9, 0.0)): The beta values for the optimizer. eps (`float`, defaults to 1e-8): The epsilon value for the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. max_unorm (`float`, defaults to 0.0): The maximum value to normalize each block with. skip_zeros (`bool`, defaults to `False`): Whether to skip zero values for sparse gradients and models to ensure correct updates. is_paged (`bool`, defaults to `False`): Whether the optimizer is a paged optimizer or not. """ if not 0.0 <= lr: raise ValueError(f"Invalid learning rate: {lr}") if not 0.0 <= eps: raise ValueError(f"Invalid epsilon value: {eps}") for i in range(len(betas)): if not 0.0 <= betas[i] < 1.0: raise ValueError(f"Invalid beta parameter at index {i}: {betas[i]}") if not 0.0 <= weight_decay: raise ValueError(f"Invalid weight_decay value: {weight_decay}") defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay) super().__init__(params, defaults, optim_bits, is_paged) if args is None: args = {} args["optim_bits"] = optim_bits args["min_8bit_size"] = min_8bit_size args["percentile_clipping"] = percentile_clipping args["block_wise"] = block_wise args["max_unorm"] = max_unorm args["skip_zeros"] = skip_zeros self.args = MockArgs(args) else: self.args = args self.optimizer_name = optimizer_name @torch.no_grad() def init_state(self, group, p, gindex, pindex): config = self.get_config(gindex, pindex, group) if config["optim_bits"] == 32: dtype = torch.float32 elif config["optim_bits"] == 8: dtype = torch.uint8 else: raise NotImplementedError(f"Amount of optimizer bits not supported: {config['optim_bits']}") if p.numel() < config["min_8bit_size"]: dtype = torch.float32 state = self.state[p] state["step"] = 0 if dtype == torch.float32: state["state1"] = self.get_state_buffer(p, dtype=torch.float32) elif dtype == torch.uint8: if state["step"] == 0: if "dynamic" not in self.name2qmap: self.fill_qmap() self.name2qmap["dynamic"] = self.name2qmap["dynamic"].to(p.device) state["state1"] = self.get_state_buffer(p, dtype=torch.uint8) state["qmap1"] = self.name2qmap["dynamic"] if config["block_wise"]: blocksize = 256 n = p.numel() blocks = (n // blocksize) + bool(n % blocksize) state["absmax1"] = torch.zeros((blocks,), dtype=torch.float32, device=p.device) else: state["max1"] = torch.zeros((1,), dtype=torch.float32, device=p.device) state["new_max1"] = torch.zeros((1,), dtype=torch.float32, device=p.device) if config["percentile_clipping"] < 100: state["gnorm_vec"] = torch.zeros((100,), device=p.device) if config["max_unorm"] > 0.0: state["unorm_vec"] = torch.zeros((1,), device=p.device) @torch.no_grad() def update_step(self, group, p, gindex, pindex): # avoid update error from non-contiguous memory layout p.data = p.data.contiguous() p.grad = p.grad.contiguous() state = self.state[p] grad = p.grad config = self.get_config(gindex, pindex, group) state["step"] += 1 step = state["step"] if config["percentile_clipping"] < 100: current_gnorm, clip_value, gnorm_scale = F.percentile_clipping( grad, state["gnorm_vec"], step, config["percentile_clipping"], ) else: gnorm_scale = 1.0 if state["state1"].dtype == torch.float: F.optimizer_update_32bit( self.optimizer_name, grad, p, state["state1"], config["betas"][0], config["eps"], step, config["lr"], None, config["betas"][1], 0.0, 0.0, config["weight_decay"], gnorm_scale, state["unorm_vec"] if config["max_unorm"] > 0.0 else None, max_unorm=config["max_unorm"], skip_zeros=config["skip_zeros"], ) elif state["state1"].dtype == torch.uint8 and not config["block_wise"]: F.optimizer_update_8bit( self.optimizer_name, grad, p, state["state1"], None, config["betas"][0], config["betas"][1], config["eps"], step, config["lr"], state["qmap1"], None, state["max1"], None, state["new_max1"], None, config["weight_decay"], gnorm_scale, state["unorm_vec"] if config["max_unorm"] > 0.0 else None, max_unorm=config["max_unorm"], ) state["max1"], state["new_max1"] = state["new_max1"], state["max1"] elif state["state1"].dtype == torch.uint8 and config["block_wise"]: F.optimizer_update_8bit_blockwise( self.optimizer_name, grad, p, state["state1"], None, config["betas"][0], config["betas"][1], 0.0, 0.0, config["eps"], step, config["lr"], state["qmap1"], None, state["absmax1"], None, config["weight_decay"], gnorm_scale=gnorm_scale, skip_zeros=config["skip_zeros"], ) bitsandbytes/bitsandbytes/optim/rmsprop.py000066400000000000000000000171521506752665600215000ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from bitsandbytes.optim.optimizer import Optimizer1State class RMSprop(Optimizer1State): def __init__( self, params, lr=1e-2, alpha=0.99, eps=1e-8, weight_decay=0, momentum=0, centered=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ Base RMSprop optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-2): The learning rate. alpha (`float`, defaults to 0.99): The alpha value is the decay rate of the squared gradients of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. centered (`bool`, defaults to `False`): Whether the gradients are normalized by the variance. If `True`, it can help training at the expense of additional compute. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ if alpha == 0: raise NotImplementedError("RMSprop with alpha==0.0 is not supported!") if centered: raise NotImplementedError("Centered RMSprop is not supported!") super().__init__( "rmsprop", params, lr, (alpha, momentum), eps, weight_decay, optim_bits, args, min_8bit_size, percentile_clipping, block_wise, ) class RMSprop8bit(Optimizer1State): def __init__( self, params, lr=1e-2, alpha=0.99, eps=1e-8, weight_decay=0, momentum=0, centered=False, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ 8-bit RMSprop optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-2): The learning rate. alpha (`float`, defaults to 0.99): The alpha value is the decay rate of the squared gradients of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. centered (`bool`, defaults to `False`): Whether the gradients are normalized by the variance. If `True`, it can help training at the expense of additional compute. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ if alpha == 0: raise NotImplementedError("RMSprop with alpha==0.0 is not supported!") if centered: raise NotImplementedError("Centered RMSprop is not supported!") super().__init__( "rmsprop", params, lr, (alpha, momentum), eps, weight_decay, 8, args, min_8bit_size, percentile_clipping, block_wise, ) class RMSprop32bit(Optimizer1State): def __init__( self, params, lr=1e-2, alpha=0.99, eps=1e-8, weight_decay=0, momentum=0, centered=False, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ 32-bit RMSprop optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`, defaults to 1e-2): The learning rate. alpha (`float`, defaults to 0.99): The alpha value is the decay rate of the squared gradients of the optimizer. eps (`float`, defaults to 1e-8): The epsilon value prevents division by zero in the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. centered (`bool`, defaults to `False`): Whether the gradients are normalized by the variance. If `True`, it can help training at the expense of additional compute. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ if alpha == 0: raise NotImplementedError("RMSprop with alpha==0.0 is not supported!") if centered: raise NotImplementedError("Centered RMSprop is not supported!") super().__init__( "rmsprop", params, lr, (alpha, momentum), eps, weight_decay, 32, args, min_8bit_size, percentile_clipping, block_wise, ) bitsandbytes/bitsandbytes/optim/sgd.py000066400000000000000000000144751506752665600205600ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from bitsandbytes.optim.optimizer import Optimizer1State class SGD(Optimizer1State): def __init__( self, params, lr, momentum=0, dampening=0, weight_decay=0, nesterov=False, optim_bits=32, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ Base SGD optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`): The learning rate. momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. dampening (`float`, defaults to 0): The dampening value reduces the momentum of the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. nesterov (`bool`, defaults to `False`): Whether to use Nesterov momentum. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ if momentum == 0: raise NotImplementedError("SGD without momentum is not supported!") super().__init__( "momentum", params, lr, (momentum, dampening), 0.0, weight_decay, optim_bits, args, min_8bit_size, percentile_clipping, block_wise, ) class SGD8bit(Optimizer1State): def __init__( self, params, lr, momentum=0, dampening=0, weight_decay=0, nesterov=False, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ 8-bit SGD optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`): The learning rate. momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. dampening (`float`, defaults to 0): The dampening value reduces the momentum of the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. nesterov (`bool`, defaults to `False`): Whether to use Nesterov momentum. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ if momentum == 0: raise NotImplementedError("SGD without momentum is not supported!") super().__init__( "momentum", params, lr, (momentum, dampening), 0.0, weight_decay, 8, args, min_8bit_size, percentile_clipping, block_wise, ) class SGD32bit(Optimizer1State): def __init__( self, params, lr, momentum=0, dampening=0, weight_decay=0, nesterov=False, args=None, min_8bit_size=4096, percentile_clipping=100, block_wise=True, ): """ 32-bit SGD optimizer. Arguments: params (`torch.tensor`): The input parameters to optimize. lr (`float`): The learning rate. momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. dampening (`float`, defaults to 0): The dampening value reduces the momentum of the optimizer. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. nesterov (`bool`, defaults to `False`): Whether to use Nesterov momentum. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. percentile_clipping (`int`, defaults to 100): Adapts clipping threshold automatically by tracking the last 100 gradient norms and clipping the gradient at a certain percentile to improve stability. block_wise (`bool`, defaults to `True`): Whether to independently quantize each block of tensors to reduce outlier effects and improve stability. """ if momentum == 0: raise NotImplementedError("SGD without momentum is not supported!") super().__init__( "momentum", params, lr, (momentum, dampening), 0.0, weight_decay, 32, args, min_8bit_size, percentile_clipping, block_wise, ) bitsandbytes/bitsandbytes/py.typed000066400000000000000000000000001506752665600177530ustar00rootroot00000000000000bitsandbytes/bitsandbytes/research/000077500000000000000000000000001506752665600200625ustar00rootroot00000000000000bitsandbytes/bitsandbytes/research/__init__.py000066400000000000000000000001671506752665600221770ustar00rootroot00000000000000from . import nn from .autograd._functions import ( matmul_fp8_global, matmul_fp8_mixed, switchback_bnb, ) bitsandbytes/bitsandbytes/research/autograd/000077500000000000000000000000001506752665600216705ustar00rootroot00000000000000bitsandbytes/bitsandbytes/research/autograd/__init__.py000066400000000000000000000000001506752665600237670ustar00rootroot00000000000000bitsandbytes/bitsandbytes/research/autograd/_functions.py000066400000000000000000000337671506752665600244310ustar00rootroot00000000000000from functools import reduce # Required in Python 3 import operator from typing import Optional import warnings import torch from bitsandbytes.autograd._functions import GlobalOutlierPooler, MatmulLtState import bitsandbytes.functional as F # math.prod not compatible with python < 3.8 def prod(iterable): return reduce(operator.mul, iterable, 1) class MatMulFP8Mixed(torch.autograd.Function): # forward is the same, but we added the fallback for pre-turing GPUs # backward is mostly the same, but adds one extra clause (see "elif state.CxB is not None") @staticmethod def forward(ctx, A, B, out=None, fw_code=None, bw_code=None, bsz=1024, bsz2=1024): # default of pytorch behavior if inputs are empty ctx.is_empty = False if prod(A.shape) == 0: ctx.is_empty = True ctx.A = A ctx.B = B B_shape = B.shape if A.shape[-1] == B_shape[0]: return torch.empty(A.shape[:-1] + B_shape[1:], dtype=A.dtype, device=A.device) else: return torch.empty(A.shape[:-1] + B_shape[:1], dtype=A.dtype, device=A.device) # 1. Dequantize # 2. MatmulnN cA, state = F.quantize_blockwise(A, code=fw_code, blocksize=bsz) fp8A = F.dequantize_blockwise(cA, state, blocksize=bsz).to(A.dtype) cB, state = F.quantize(B.float(), code=fw_code) fp8B = F.dequantize(cB, state).to(B.dtype) output = torch.matmul(fp8A, fp8B) # output is half # 3. Save state ctx.fw_code = fw_code ctx.bw_code = bw_code ctx.bsz = bsz ctx.bsz2 = bsz2 ctx.dtype_A, ctx.dtype_B = A.dtype, B.dtype if any(ctx.needs_input_grad[:2]): # NOTE: we send back A, and re-quant. ctx.tensors = (A, fp8B) else: ctx.tensors = (None, None) return output @staticmethod def backward(ctx, grad_output): if ctx.is_empty: return torch.zeros_like(ctx.A), torch.zeros_like(ctx.B), None, None, None, None, None req_gradA, req_gradB, _, _, _, _, _ = ctx.needs_input_grad A, B = ctx.tensors grad_A, grad_B = None, None # TODO: Fix blocksize to be output_dim cgrad_out, state = F.quantize_blockwise(grad_output, code=ctx.bw_code, blocksize=ctx.bsz2) fp8out = F.dequantize_blockwise(cgrad_out, state, blocksize=ctx.bsz2).to(grad_output.dtype) # cgrad_output_2, state_2 = F.quantize(grad_output.float(), code=ctx.bw_code) # fp8out_2 = F.dequantize(cgrad_output_2, state_2).to(grad_output.dtype) # grad_output_reshape = grad_output.reshape(-1, grad_output.shape[-1]).contiguous() # fp8grad_transpose, stategrad_transpose = F.vectorwise_quant(grad_output_reshape, dim=0, quant_type='vector') # fp8out_transpose = (fp8grad_transpose / 7) * stategrad_transpose # fp8out_transpose = fp8out_transpose.view(grad_output.shape[0], grad_output.shape[1], grad_output.shape[2]) # not supported by PyTorch. TODO: create work-around if req_gradA: grad_A = torch.matmul(fp8out, B.t().to(fp8out.dtype)).to(A.dtype) if req_gradB: if len(A.shape) == 3: At = A.transpose(2, 1).contiguous() else: At = A.transpose(1, 0).contiguous() # cA, state = F.quantize(At.float(), code=ctx.fw_code) # fp8At = F.dequantize(cA, state).to(A.dtype) grad_B = torch.matmul(At.to(grad_output.dtype), grad_output).to(B.dtype) return grad_A, grad_B, None, None, None, None, None class MatMulFP8Global(torch.autograd.Function): # forward is the same, but we added the fallback for pre-turing GPUs # backward is mostly the same, but adds one extra clause (see "elif state.CxB is not None") @staticmethod def forward(ctx, A, B, out=None, fw_code=None, bw_code=None, bsz=1024, bsz2=1024): # default of pytorch behavior if inputs are empty ctx.is_empty = False if prod(A.shape) == 0: ctx.is_empty = True ctx.A = A ctx.B = B B_shape = B.shape if A.shape[-1] == B_shape[0]: return torch.empty(A.shape[:-1] + B_shape[1:], dtype=A.dtype, device=A.device) else: return torch.empty(A.shape[:-1] + B_shape[:1], dtype=A.dtype, device=A.device) # 1. Dequantize # 2. MatmulnN cA, state = F.quantize(A.float(), code=fw_code) fp8A = F.dequantize(cA, state).to(A.dtype) cB, state = F.quantize(B.float(), code=fw_code) fp8B = F.dequantize(cB, state).to(B.dtype) output = torch.matmul(fp8A, fp8B) # output is half # 3. Save state ctx.fw_code = fw_code ctx.bw_code = bw_code ctx.bsz = bsz ctx.bsz2 = bsz2 ctx.dtype_A, ctx.dtype_B = A.dtype, B.dtype if any(ctx.needs_input_grad[:2]): # NOTE: we send back A, and re-quant. ctx.tensors = (A, fp8B) else: ctx.tensors = (None, None) return output @staticmethod def backward(ctx, grad_output): if ctx.is_empty: return torch.zeros_like(ctx.A), torch.zeros_like(ctx.B), None, None, None, None, None req_gradA, req_gradB, _, _, _, _, _ = ctx.needs_input_grad A, B = ctx.tensors grad_A, grad_B = None, None # TODO: Fix blocksize to be output_dim cgrad_out, state = F.quantize(grad_output.float(), code=ctx.bw_code) fp8out = F.dequantize(cgrad_out, state).to(grad_output.dtype) # cgrad_output_2, state_2 = F.quantize(grad_output.float(), code=ctx.bw_code) # fp8out_2 = F.dequantize(cgrad_output_2, state_2).to(grad_output.dtype) # grad_output_reshape = grad_output.reshape(-1, grad_output.shape[-1]).contiguous() # fp8grad_transpose, stategrad_transpose = F.vectorwise_quant(grad_output_reshape, dim=0, quant_type='vector') # fp8out_transpose = (fp8grad_transpose / 7) * stategrad_transpose # fp8out_transpose = fp8out_transpose.view(grad_output.shape[0], grad_output.shape[1], grad_output.shape[2]) # not supported by PyTorch. TODO: create work-around if req_gradA: grad_A = torch.matmul(fp8out, B.t().to(fp8out.dtype)).to(A.dtype) if req_gradB: if len(A.shape) == 3: At = A.transpose(2, 1).contiguous() else: At = A.transpose(1, 0).contiguous() cA, state = F.quantize(At.float(), code=ctx.fw_code) fp8At = F.dequantize(cA, state).to(A.dtype) grad_B = torch.matmul(fp8At.to(fp8out.dtype), fp8out).to(B.dtype) return grad_A, grad_B, None, None, None, None, None class SwitchBackBnb(torch.autograd.Function): @staticmethod def forward(ctx, A, B, out=None, bias=None, state: Optional[MatmulLtState] = None): state = state or MatmulLtState() # default to pytorch behavior if inputs are empty ctx.is_empty = False if prod(A.shape) == 0: ctx.is_empty = True ctx.A = A ctx.B = B ctx.bias = bias if A.shape[-1] == B.shape[0]: return torch.empty(A.shape[:-1] + B.shape[1:], dtype=A.dtype, device=A.device) else: return torch.empty(A.shape[:-1] + B.shape[:1], dtype=A.dtype, device=A.device) # 1. Quantize A # 2. Quantize B # 3. Matmul # 4. Mixed-precision decomposition matmul # 5. Save state input_shape = A.shape if state.outlier_pool is None: state.outlier_pool = GlobalOutlierPooler.get_instance() # Cast A to fp16 if A.dtype != torch.float16: warnings.warn(f"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization") # 1. Quantize A if len(A.shape) == 3: A = A.view(-1, A.shape[-1]).contiguous() CA, CAt, SCA, SCAt, outlier_cols = F.int8_double_quant(A.to(torch.float16), threshold=state.threshold) if state.threshold > 0.0 and outlier_cols is not None: if state.has_fp16_weights: idx = outlier_cols CA[:, idx] = 0 subA = A[:, idx] state.subB = B[:, idx].t().contiguous() state.idx = idx else: if state.SB is None: state.SB = (state.CB.shape, "row") else: if not state.has_fp16_weights and state.SB is None: state.SB = (state.CB.shape, "row") subA = None # 2. Quantize B if state.has_fp16_weights: # print('B shape', B.shape) has_grad = getattr(B, "grad", None) is not None is_transposed = not B.is_contiguous() and B.shape[0] == B.stride(1) if is_transposed: B = B.contiguous() if (state.is_training and not has_grad) or state.SB is None: state.reset_grads() ( state.CB, state.CBt, state.SCB, state.SCBt, _, ) = F.int8_double_quant(B.to(torch.float16)) state.SB = (state.CB.shape, "row") else: has_grad = False if outlier_cols is not None and not state.has_fp16_weights: # extract outliers state.idx = outlier_cols outliers = state.CB[:, state.idx.long()].clone() state.subB = (outliers * state.SCB.view(-1, 1) / 127.0).t().contiguous().to(A.dtype) CA[:, state.idx.long()] = 0 subA = A[:, state.idx.long()] shapeB = state.SB[0] if len(input_shape) == 3: output_shape = (input_shape[0], input_shape[1], shapeB[0]) else: output_shape = (input_shape[0], shapeB[0]) # 3. Matmul out32 = F.int8_linear_matmul(CA, state.CB) # we apply the fused bias here if bias is None or bias.dtype == torch.float16: output = F.int8_mm_dequant(out32, SCA, state.SCB, bias=bias).to(A.dtype) else: # apply bias separately output = F.int8_mm_dequant(out32, SCA, state.SCB, bias=None).to(A.dtype) output.add_(bias) # 4. Mixed-precision decomposition matmul if outlier_cols is not None and subA is not None: output += torch.matmul(subA, state.subB) # 5. Save state ctx.state = state ctx.grad_shape = input_shape ctx.dtype_A, ctx.dtype_B, ctx.dtype_bias = A.dtype, B.dtype, None if bias is None else bias.dtype if any(ctx.needs_input_grad[:2]): ctx.tensors = (CAt, subA, A) ctx.tensor_states = (SCAt, state.idx) else: ctx.tensors = [None, None, None] ctx.tensor_states = (None, None) ctx.save_for_backward(None, None) clone_func = torch.clone if len(output_shape) == 3 else lambda x: x return clone_func(output.view(output_shape)) @staticmethod def backward(ctx, grad_output): if ctx.is_empty: bias_grad = None if ctx.bias is None else torch.zeros_like(ctx.bias) return torch.zeros_like(ctx.A), torch.zeros_like(ctx.B), None, bias_grad, None req_gradA, req_gradB, _, req_gradBias, _ = ctx.needs_input_grad CAt, subA, A = ctx.tensors SCAt, idx = ctx.tensor_states state = ctx.state grad_A = grad_B = grad_bias = None if req_gradBias: # compute grad_bias first before changing grad_output dtype grad_bias = grad_output.sum(0, dtype=ctx.dtype_bias) # Cast grad_output to fp16 if len(grad_output.shape) == 3: grad_output = grad_output.reshape(-1, grad_output.shape[-1]).contiguous() Cgrad, Cgradt, SCgrad, SCgradt, outlier_cols = F.int8_double_quant(grad_output.to(torch.float16)) if req_gradB: # print('back A shape', A.shape) # print('grad output t shape', grad_output.t().shape) grad_B = torch.matmul(grad_output.t(), A) if req_gradA: if state.CB is not None: CB = state.CB.to(ctx.dtype_A, copy=True).mul_(state.SCB.unsqueeze(1).mul(1.0 / 127.0)) grad_A = torch.matmul(grad_output, CB).view(ctx.grad_shape).to(ctx.dtype_A) else: raise Exception("State must contain either CBt or CB matrix for backward") return grad_A, grad_B, None, grad_bias, None def get_block_sizes(input_matrix, weight_matrix): input_features = input_matrix.shape[-1] output_features = weight_matrix.shape[0] if weight_matrix.shape[1] == input_features else weight_matrix.shape[1] array = [4096, 2048, 1024, 512, 256, 128, 64, 0] bsz, bsz2 = 1024, 1024 for i, k in enumerate(array): if input_features > array[i + 1]: bsz = k break for i, k in enumerate(array): if output_features > array[i + 1]: bsz2 = k break return bsz, bsz2 def matmul_fp8_global( A: torch.Tensor, B: torch.Tensor, fw_code: torch.Tensor, bw_code: torch.Tensor, out: Optional[torch.Tensor] = None, bsz: int = -1, bsz2: int = -1, ): if bsz == -1 or bsz2 == -1: bsz, bsz2 = get_block_sizes(A, B) return MatMulFP8Global.apply(A, B, out, fw_code, bw_code, bsz, bsz2) def matmul_fp8_mixed( A: torch.Tensor, B: torch.Tensor, fw_code: torch.Tensor, bw_code: torch.Tensor, out: Optional[torch.Tensor] = None, bsz: int = -1, bsz2: int = -1, ): if bsz == -1 or bsz2 == -1: bsz, bsz2 = get_block_sizes(A, B) return MatMulFP8Mixed.apply(A, B, out, fw_code, bw_code, bsz, bsz2) def switchback_bnb( A: torch.Tensor, B: torch.Tensor, out: Optional[torch.Tensor] = None, state: Optional[MatmulLtState] = None, threshold=0.0, bias=None, ): state = state or MatmulLtState() if threshold > 0.0: state.threshold = threshold return SwitchBackBnb.apply(A, B, out, bias, state) bitsandbytes/bitsandbytes/research/nn/000077500000000000000000000000001506752665600204755ustar00rootroot00000000000000bitsandbytes/bitsandbytes/research/nn/__init__.py000066400000000000000000000000651506752665600226070ustar00rootroot00000000000000from .modules import LinearFP8Global, LinearFP8Mixed bitsandbytes/bitsandbytes/research/nn/modules.py000066400000000000000000000044301506752665600225200ustar00rootroot00000000000000from typing import TypeVar import torch from torch import nn import bitsandbytes as bnb T = TypeVar("T", bound="torch.nn.Module") class LinearFP8Mixed(nn.Linear): def __init__(self, input_features, output_features, bias=True): super().__init__(input_features, output_features, bias) self.bw_code = None self.fw_code = None array = [4096, 2048, 1024, 512, 256, 128, 64, 0] for i, k in enumerate(array): if input_features > array[i + 1]: self.bsz = k break for i, k in enumerate(array): if output_features > array[i + 1]: self.bsz2 = k break def forward(self, x: torch.Tensor): if self.fw_code is None: self.bw_code = bnb.functional.create_fp8_map(True, 5, 2, 8).to(x.device) self.fw_code = bnb.functional.create_fp8_map(True, 4, 3, 8).to(x.device) out = bnb.research.matmul_fp8_mixed( x, self.weight.t(), fw_code=self.fw_code, bw_code=self.bw_code, bsz=self.bsz, bsz2=self.bsz2, ) if self.bias is not None: out += self.bias return out class LinearFP8Global(nn.Linear): def __init__(self, input_features, output_features, bias=True): super().__init__(input_features, output_features, bias) self.bw_code = None self.fw_code = None array = [4096, 2048, 1024, 512, 256, 128, 64, 0] for i, k in enumerate(array): if input_features > array[i + 1]: self.bsz = k break for i, k in enumerate(array): if output_features > array[i + 1]: self.bsz2 = k break def forward(self, x: torch.Tensor): if self.fw_code is None: self.bw_code = bnb.functional.create_fp8_map(True, 5, 2, 8).to(x.device) self.fw_code = bnb.functional.create_fp8_map(True, 4, 3, 8).to(x.device) out = bnb.matmul_fp8_global( x, self.weight.t(), fw_code=self.fw_code, bw_code=self.bw_code, bsz=self.bsz, bsz2=self.bsz2, ) if self.bias is not None: out += self.bias return out bitsandbytes/bitsandbytes/triton/000077500000000000000000000000001506752665600176055ustar00rootroot00000000000000bitsandbytes/bitsandbytes/triton/__init__.py000066400000000000000000000000001506752665600217040ustar00rootroot00000000000000bitsandbytes/bitsandbytes/triton/dequantize_rowwise.py000066400000000000000000000037701506752665600241160ustar00rootroot00000000000000import math import torch from bitsandbytes.triton.triton_utils import is_triton_available if not is_triton_available(): def dequantize_rowwise(x: torch.Tensor, state_x: torch.Tensor): return None else: import triton import triton.language as tl # rowwise quantize # TODO: autotune this better. @triton.autotune( configs=[ triton.Config({}, num_stages=1, num_warps=8), triton.Config({}, num_stages=2, num_warps=8), triton.Config({}, num_stages=4, num_warps=8), triton.Config({}, num_stages=8, num_warps=8), triton.Config({}, num_stages=1), triton.Config({}, num_stages=2), triton.Config({}, num_stages=4), triton.Config({}, num_stages=8), triton.Config({}, num_warps=1), triton.Config({}, num_warps=2), triton.Config({}, num_warps=4), triton.Config({}, num_warps=8), ], key=["n_elements"], ) @triton.jit def _dequantize_rowwise( x_ptr, state_x, output_ptr, inv_127, n_elements, BLOCK_SIZE: tl.constexpr, P2: tl.constexpr, ): pid = tl.program_id(axis=0) block_start = pid * BLOCK_SIZE arange = tl.arange(0, P2) offsets = block_start + arange row_mask = arange < BLOCK_SIZE x = tl.load(x_ptr + offsets, mask=row_mask) max_val = tl.load(state_x + pid) output = max_val * x * inv_127 tl.store(output_ptr + offsets, output, mask=row_mask) def dequantize_rowwise(x: torch.Tensor, state_x: torch.Tensor): output = torch.empty(*x.shape, device=x.device, dtype=torch.float16) P2 = int(2 ** (math.ceil(math.log2(x.shape[1])))) assert x.is_cuda and output.is_cuda n_elements = output.numel() grid = lambda meta: (x.shape[0],) _dequantize_rowwise[grid](x, state_x, output, 1.0 / 127, n_elements, BLOCK_SIZE=x.shape[1], P2=P2) return output bitsandbytes/bitsandbytes/triton/int8_matmul_mixed_dequantize.py000066400000000000000000000210731506752665600260420ustar00rootroot00000000000000import torch from bitsandbytes.triton.triton_utils import is_triton_available if not is_triton_available(): def int8_matmul_mixed_dequantize(a, b, state_x, state_w, bias): return None else: import triton import triton.language as tl from .matmul_perf_model import early_config_prune, estimate_matmul_time # This is a matmul kernel based on triton.ops.matmul # It is modified to support rowwise quantized input and global quantized weight # It's purpose is fused matmul then dequantize # It does support bias. def init_to_zero(name): return lambda nargs: nargs[name].zero_() def get_configs_io_bound(): configs = [] for num_stages in [2, 3, 4, 5, 6]: for block_m in [16, 32]: for block_k in [32, 64]: for block_n in [32, 64, 128, 256]: num_warps = 2 if block_n <= 64 else 4 configs.append( triton.Config( {"BLOCK_M": block_m, "BLOCK_N": block_n, "BLOCK_K": block_k, "SPLIT_K": 1}, num_stages=num_stages, num_warps=num_warps, ), ) # split_k for split_k in [2, 4, 8, 16]: configs.append( triton.Config( {"BLOCK_M": block_m, "BLOCK_N": block_n, "BLOCK_K": block_k, "SPLIT_K": split_k}, num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero("C"), ), ) return configs @triton.autotune( configs=[ # basic configs for compute-bound matmuls triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8), triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8), triton.Config({"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2), # good for int8 triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=3, num_warps=8), triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=3, num_warps=8), triton.Config({"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2), *get_configs_io_bound(), ], key=["M", "N", "K"], prune_configs_by={"early_config_prune": early_config_prune, "perf_model": estimate_matmul_time, "top_k": 10}, ) @triton.heuristics( { "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, }, ) @triton.jit def _int8_matmul_mixed_dequantize( A, B, C, bias, state_x_ptr, state_w_ptr, M, N, K, divfactor: tl.constexpr, has_bias: tl.constexpr, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr, SPLIT_K: tl.constexpr, EVEN_K: tl.constexpr, ACC_TYPE: tl.constexpr, ): # matrix multiplication pid = tl.program_id(0) pid_z = tl.program_id(1) grid_m = tl.cdiv(M, BLOCK_M) grid_n = tl.cdiv(N, BLOCK_N) # re-order program ID for better L2 performance width = GROUP_M * grid_n group_id = pid // width group_size = min(grid_m - group_id * GROUP_M, GROUP_M) pid_m = group_id * GROUP_M + (pid % group_size) pid_n = (pid % width) // (group_size) # do matrix multiplication rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K) # pointers A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) # rematerialize rm and rn to save registers rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) w_factor = tl.load(state_w_ptr) x_factor = tl.load(state_x_ptr + ram)[:, None] # acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.int32) for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)): if EVEN_K: a = tl.load(A) b = tl.load(B) else: k_remaining = K - k * (BLOCK_K * SPLIT_K) a = tl.load(A, mask=rk[None, :] < k_remaining, other=0.0) b = tl.load(B, mask=rk[:, None] < k_remaining, other=0.0) acc += tl.dot(a, b) A += BLOCK_K * SPLIT_K * stride_ak B += BLOCK_K * SPLIT_K * stride_bk acc = w_factor * (x_factor * (acc * divfactor)) acc = acc.to(C.dtype.element_ty) # conditionally add bias if has_bias: bias = tl.load(bias + rn).to(C.dtype.element_ty) acc = acc + bias[None, :] C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn) mask = (rm < M)[:, None] & (rn < N)[None, :] # handles write-back with reduction-splitting if SPLIT_K == 1: tl.store(C, acc, mask=mask) else: tl.atomic_add(C, acc, mask=mask) def int8_matmul_mixed_dequantize(a, b, state_x, state_w, bias): device = a.device divfactor = 1.0 / (127.0 * 127.0) has_bias = 0 if bias is None else 1 # handle non-contiguous inputs if necessary if a.stride(0) > 1 and a.stride(1) > 1: a = a.contiguous() if b.stride(0) > 1 and b.stride(1) > 1: b = b.contiguous() # checks constraints assert a.shape[1] == b.shape[0], "incompatible dimensions" M, K = a.shape _, N = b.shape # allocates output c = torch.empty((M, N), device=device, dtype=torch.float16) # accumulator types ACC_TYPE = tl.float32 # if a.dtype in [torch.float16, torch.bfloat16, torch.float32] else tl.int32 # launch int8_matmul_mixed_dequantize kernel grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]), META["SPLIT_K"]) _int8_matmul_mixed_dequantize[grid]( a, b, c, bias, state_x, state_w, M, N, K, divfactor, has_bias, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1), GROUP_M=8, ACC_TYPE=ACC_TYPE, ) return c bitsandbytes/bitsandbytes/triton/int8_matmul_rowwise_dequantize.py000066400000000000000000000210531506752665600264310ustar00rootroot00000000000000import torch from bitsandbytes.triton.triton_utils import is_triton_available if not is_triton_available(): def int8_matmul_rowwise_dequantize(a, b, state_x, state_w, bias): return None else: import triton import triton.language as tl from .matmul_perf_model import early_config_prune, estimate_matmul_time # This is a matmul kernel based on triton.ops.matmul # It is modified to support rowwise quantized input and columnwise quantized weight # It's purpose is fused matmul then dequantize # It does support bias. def init_to_zero(name): return lambda nargs: nargs[name].zero_() def get_configs_io_bound(): configs = [] for num_stages in [2, 3, 4, 5, 6]: for block_m in [16, 32]: for block_k in [32, 64]: for block_n in [32, 64, 128, 256]: num_warps = 2 if block_n <= 64 else 4 configs.append( triton.Config( {"BLOCK_M": block_m, "BLOCK_N": block_n, "BLOCK_K": block_k, "SPLIT_K": 1}, num_stages=num_stages, num_warps=num_warps, ), ) # split_k for split_k in [2, 4, 8, 16]: configs.append( triton.Config( {"BLOCK_M": block_m, "BLOCK_N": block_n, "BLOCK_K": block_k, "SPLIT_K": split_k}, num_stages=num_stages, num_warps=num_warps, pre_hook=init_to_zero("C"), ), ) return configs @triton.autotune( configs=[ # basic configs for compute-bound matmuls triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8), triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=3, num_warps=8), triton.Config({"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 32, "SPLIT_K": 1}, num_stages=5, num_warps=2), # good for int8 triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=3, num_warps=8), triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=3, num_warps=8), triton.Config({"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=4, num_warps=4), triton.Config({"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64, "SPLIT_K": 1}, num_stages=5, num_warps=2), *get_configs_io_bound(), ], key=["M", "N", "K"], prune_configs_by={"early_config_prune": early_config_prune, "perf_model": estimate_matmul_time, "top_k": 10}, ) @triton.heuristics( { "EVEN_K": lambda args: args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0, }, ) @triton.jit def _int8_matmul_rowwise_dequantize( A, B, C, bias, state_x_ptr, state_w_ptr, M, N, K, divfactor, has_bias: tl.constexpr, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr, SPLIT_K: tl.constexpr, EVEN_K: tl.constexpr, ACC_TYPE: tl.constexpr, ): # matrix multiplication pid = tl.program_id(0) pid_z = tl.program_id(1) grid_m = tl.cdiv(M, BLOCK_M) grid_n = tl.cdiv(N, BLOCK_N) # re-order program ID for better L2 performance width = GROUP_M * grid_n group_id = pid // width group_size = min(grid_m - group_id * GROUP_M, GROUP_M) pid_m = group_id * GROUP_M + (pid % group_size) pid_n = (pid % width) // (group_size) # do matrix multiplication rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M) rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N) rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K) # pointers A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak) B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn) # rematerialize rm and rn to save registers rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) w_factor = tl.load(state_w_ptr + rbn)[None, :] x_factor = tl.load(state_x_ptr + ram)[:, None] # acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.int32) for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)): if EVEN_K: a = tl.load(A) b = tl.load(B) else: k_remaining = K - k * (BLOCK_K * SPLIT_K) a = tl.load(A, mask=rk[None, :] < k_remaining, other=0.0) b = tl.load(B, mask=rk[:, None] < k_remaining, other=0.0) acc += tl.dot(a, b) A += BLOCK_K * SPLIT_K * stride_ak B += BLOCK_K * SPLIT_K * stride_bk acc = w_factor * (x_factor * (acc * divfactor)) acc = acc.to(C.dtype.element_ty) if has_bias: bias = tl.load(bias + rn).to(C.dtype.element_ty) acc = acc + bias[None, :] C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn) mask = (rm < M)[:, None] & (rn < N)[None, :] # handles write-back with reduction-splitting if SPLIT_K == 1: tl.store(C, acc, mask=mask) else: tl.atomic_add(C, acc, mask=mask) def int8_matmul_rowwise_dequantize(a, b, state_x, state_w, bias): divfactor = 1.0 / (127.0 * 127.0) has_bias = 0 if bias is None else 1 device = a.device # handle non-contiguous inputs if necessary if a.stride(0) > 1 and a.stride(1) > 1: a = a.contiguous() if b.stride(0) > 1 and b.stride(1) > 1: b = b.contiguous() # checks constraints assert a.shape[1] == b.shape[0], "incompatible dimensions" M, K = a.shape _, N = b.shape # allocates output c = torch.empty((M, N), device=device, dtype=torch.float16) # accumulator types ACC_TYPE = tl.float32 # if a.dtype in [torch.float16, torch.bfloat16, torch.float32] else tl.int32 # launch int8_matmul_rowwise_dequantize kernel grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]), META["SPLIT_K"]) _int8_matmul_rowwise_dequantize[grid]( a, b, c, bias, state_x, state_w, M, N, K, divfactor, has_bias, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1), GROUP_M=8, ACC_TYPE=ACC_TYPE, ) return c bitsandbytes/bitsandbytes/triton/matmul_perf_model.py000066400000000000000000000160661506752665600236630ustar00rootroot00000000000000# Adapted from https://github.com/triton-lang/kernels/blob/eeeebdd8be7d13629de22d600621e6234057eed3/kernels/matmul_perf_model.py # https://github.com/triton-lang/kernels is licensed under the MIT License. import functools import heapq import torch from triton import cdiv from triton.runtime import driver from triton.testing import ( get_dram_gbps, get_max_simd_tflops, get_max_tensorcore_tflops, nvsmi, ) @functools.lru_cache def get_clock_rate_in_khz(): try: return nvsmi(["clocks.max.sm"])[0] * 1e3 except FileNotFoundError: import pynvml pynvml.nvmlInit() handle = pynvml.nvmlDeviceGetHandleByIndex(0) return pynvml.nvmlDeviceGetMaxClockInfo(handle, pynvml.NVML_CLOCK_SM) * 1e3 def get_tensorcore_tflops(device, num_ctas, num_warps, dtype): """return compute throughput in TOPS""" total_warps = num_ctas * min(num_warps, 4) num_subcores = driver.active.utils.get_device_properties(device)["multiprocessor_count"] * 4 # on recent GPUs tflops = ( min(num_subcores, total_warps) / num_subcores * get_max_tensorcore_tflops(dtype, get_clock_rate_in_khz(), device) ) return tflops def get_simd_tflops(device, num_ctas, num_warps, dtype): """return compute throughput in TOPS""" total_warps = num_ctas * min(num_warps, 4) num_subcores = driver.active.utils.get_device_properties(device)["multiprocessor_count"] * 4 # on recent GPUs tflops = ( min(num_subcores, total_warps) / num_subcores * get_max_simd_tflops(dtype, get_clock_rate_in_khz(), device) ) return tflops def get_tflops(device, num_ctas, num_warps, dtype): capability = torch.cuda.get_device_capability(device) if capability[0] < 8 and dtype == torch.float32: return get_simd_tflops(device, num_ctas, num_warps, dtype) return get_tensorcore_tflops(device, num_ctas, num_warps, dtype) def estimate_matmul_time( # backend, device, num_warps, num_stages, # A, B, C, # M, N, K, # BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, # debug=False, **kwargs, # ): """return estimated running time in ms = max(compute, loading) + store""" device = torch.cuda.current_device() dtype = A.dtype dtsize = A.element_size() num_cta_m = cdiv(M, BLOCK_M) num_cta_n = cdiv(N, BLOCK_N) num_cta_k = SPLIT_K num_ctas = num_cta_m * num_cta_n * num_cta_k # If the input is smaller than the block size M, N = max(M, BLOCK_M), max(N, BLOCK_N) # time to compute total_ops = 2 * M * N * K / (1024 * 1024 * 1024) # GOPS tput = get_tflops(device, num_ctas, num_warps, dtype) compute_ms = total_ops / tput # time to load data num_sm = driver.active.utils.get_device_properties(device)["multiprocessor_count"] active_cta_ratio = min(1, num_ctas / num_sm) active_cta_ratio_bw1 = min(1, num_ctas / 32) # 32 active ctas are enough to saturate active_cta_ratio_bw2 = max(min(1, (num_ctas - 32) / (108 - 32)), 0) # 32-108, remaining 5% dram_bw = get_dram_gbps(device) * (active_cta_ratio_bw1 * 0.95 + active_cta_ratio_bw2 * 0.05) # in GB/s l2_bw = dram_bw * 4 # rough estimation (should be 4.7 for A100?) # assume 80% of (following) loads are in L2 cache load_a_dram = M * K * dtsize * (1 + 0.2 * (num_cta_n - 1)) load_a_l2 = M * K * dtsize * 0.8 * (num_cta_n - 1) load_b_dram = N * K * dtsize * (1 + 0.2 * (num_cta_m - 1)) load_b_l2 = N * K * dtsize * 0.8 * (num_cta_m - 1) # total total_dram = (load_a_dram + load_b_dram) / (1024 * 1024) # MB total_l2 = (load_a_l2 + load_b_l2) / (1024 * 1024) # loading time in ms load_ms = total_dram / dram_bw + total_l2 / l2_bw # estimate storing time store_bw = dram_bw * 0.6 # :o store_c_dram = M * N * dtsize * SPLIT_K / (1024 * 1024) # MB if SPLIT_K == 1: store_ms = store_c_dram / store_bw else: reduce_bw = store_bw store_ms = store_c_dram / reduce_bw # c.zero_() zero_ms = M * N * 2 / (1024 * 1024) / store_bw store_ms += zero_ms total_time_ms = max(compute_ms, load_ms) + store_ms if debug: print( f"Total time: {total_time_ms}ms, compute time: {compute_ms}ms, " f"loading time: {load_ms}ms, store time: {store_ms}ms, " f"Activate CTAs: {active_cta_ratio * 100}%" ) return total_time_ms def early_config_prune(configs, named_args, **kwargs): device = torch.cuda.current_device() capability = torch.cuda.get_device_capability() # BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, num_warps, num_stages dtsize = named_args["A"].element_size() dtype = named_args["A"].dtype # 1. make sure we have enough smem pruned_configs = [] for config in configs: kw = config.kwargs BLOCK_M, BLOCK_N, BLOCK_K, num_stages = ( kw["BLOCK_M"], kw["BLOCK_N"], kw["BLOCK_K"], config.num_stages, ) max_shared_memory = driver.active.utils.get_device_properties(device)["max_shared_mem"] required_shared_memory = (BLOCK_M + BLOCK_N) * BLOCK_K * num_stages * dtsize if required_shared_memory <= max_shared_memory: pruned_configs.append(config) configs = pruned_configs # Some dtypes do not allow atomic_add if dtype not in [torch.float16, torch.float32]: configs = [config for config in configs if config.kwargs["SPLIT_K"] == 1] # group configs by (BLOCK_M,_N,_K, SPLIT_K, num_warps) configs_map = {} for config in configs: kw = config.kwargs BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, num_warps, num_stages = ( kw["BLOCK_M"], kw["BLOCK_N"], kw["BLOCK_K"], kw["SPLIT_K"], config.num_warps, config.num_stages, ) key = (BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, num_warps) if key in configs_map: configs_map[key].append((config, num_stages)) else: configs_map[key] = [(config, num_stages)] pruned_configs = [] for k, v in configs_map.items(): BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, num_warps = k if capability[0] >= 8: # compute cycles (only works for ampere GPUs) mmas = BLOCK_M * BLOCK_N * BLOCK_K / (16 * 8 * 16) mma_cycles = mmas / min(4, num_warps) * 8 ldgsts_latency = 300 # Does this matter? optimal_num_stages = ldgsts_latency / mma_cycles # nearest stages, prefer large #stages nearest = heapq.nsmallest( 2, v, key=lambda x: ( 10 + abs(x[1] - optimal_num_stages) if (x[1] - optimal_num_stages) < 0 else x[1] - optimal_num_stages ), ) for n in nearest: pruned_configs.append(n[0]) else: # Volta & Turing only supports num_stages <= 2 random_config = v[0][0] random_config.num_stages = 2 pruned_configs.append(random_config) return pruned_configs bitsandbytes/bitsandbytes/triton/quantize_columnwise_and_transpose.py000066400000000000000000000050461506752665600272110ustar00rootroot00000000000000import math import torch from bitsandbytes.triton.triton_utils import is_triton_available if not is_triton_available(): def quantize_columnwise_and_transpose(x: torch.Tensor): return None else: import triton import triton.language as tl # This kernel does fused columnwise quantization and transpose. # TODO: autotune this better. @triton.autotune( configs=[ triton.Config({}, num_stages=1), triton.Config({}, num_stages=2), triton.Config({}, num_stages=4), triton.Config({}, num_stages=8), triton.Config({}, num_stages=16), triton.Config({}, num_stages=1, num_warps=8), triton.Config({}, num_stages=2, num_warps=8), triton.Config({}, num_stages=4, num_warps=8), triton.Config({}, num_stages=8, num_warps=8), triton.Config({}, num_stages=16, num_warps=8), triton.Config({}, num_warps=1), triton.Config({}, num_warps=2), triton.Config({}, num_warps=4), triton.Config({}, num_warps=8), ], key=["n_elements"], ) @triton.jit def _quantize_columnwise_and_transpose( x_ptr, output_ptr, output_maxs, n_elements, M: tl.constexpr, N: tl.constexpr, BLOCK_SIZE: tl.constexpr, P2: tl.constexpr, ): pid = tl.program_id(axis=0) block_start = pid p2_arange = tl.arange(0, P2) p2_arange_mask = p2_arange < M arange = p2_arange * N offsets = block_start + arange x = tl.load(x_ptr + offsets, mask=p2_arange_mask) abs_x = tl.abs(x) max_val = tl.max(tl.where(p2_arange_mask, abs_x, 0), axis=0) output = tl.libdevice.llrint(127.0 * (x / max_val)) new_start = pid * M new_offsets = new_start + p2_arange tl.store(output_ptr + new_offsets, output, mask=p2_arange_mask) tl.store(output_maxs + pid, max_val) def quantize_columnwise_and_transpose(x: torch.Tensor): M, N = x.shape output = torch.empty(N, M, device=x.device, dtype=torch.int8) output_maxs = torch.empty(x.shape[1], device=x.device, dtype=torch.float16) P2 = int(2 ** (math.ceil(math.log2(M)))) assert x.is_cuda and output.is_cuda n_elements = output.numel() grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) _quantize_columnwise_and_transpose[grid](x, output, output_maxs, n_elements, M, N, BLOCK_SIZE=M, P2=P2) return output, output_maxs bitsandbytes/bitsandbytes/triton/quantize_global.py000066400000000000000000000075211506752665600233440ustar00rootroot00000000000000import torch from bitsandbytes.triton.triton_utils import is_triton_available if not is_triton_available(): def quantize_global_transpose(input): return None def quantize_global(x: torch.Tensor): return None else: import triton import triton.language as tl # global quantize @triton.autotune( configs=[ triton.Config({"BLOCK_SIZE": 1024}, num_warps=4), triton.Config({"BLOCK_SIZE": 2048}, num_stages=1), ], key=["n_elements"], ) @triton.jit def _quantize_global( x_ptr, absmax_inv_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr, ): pid = tl.program_id(axis=0) block_start = pid * BLOCK_SIZE offsets = block_start + tl.arange(0, BLOCK_SIZE) mask = offsets < n_elements x = tl.load(x_ptr + offsets, mask=mask) absmax_inv = tl.load(absmax_inv_ptr) output = tl.libdevice.llrint(127.0 * (x * absmax_inv)) tl.store(output_ptr + offsets, output, mask=mask) def quantize_global(x: torch.Tensor): absmax = x.abs().max().unsqueeze(0) absmax_inv = 1.0 / absmax output = torch.empty(*x.shape, device="cuda", dtype=torch.int8) assert x.is_cuda and output.is_cuda n_elements = output.numel() grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) _quantize_global[grid](x, absmax_inv, output, n_elements) return output, absmax # global quantize and transpose @triton.autotune( configs=[ triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "GROUP_M": 8}, num_warps=4), triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "GROUP_M": 8}, num_warps=4), # ... ], key=["M", "N"], ) @triton.jit def _quantize_global_transpose( A, absmax_inv_ptr, B, stride_am, stride_an, stride_bn, stride_bm, M, N, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, GROUP_M: tl.constexpr, ): pid = tl.program_id(0) grid_m = (M + BLOCK_M - 1) // BLOCK_M grid_n = (N + BLOCK_N - 1) // BLOCK_N width = GROUP_M * grid_n group_id = pid // width group_size = min(grid_m - group_id * GROUP_M, GROUP_M) pid_m = group_id * GROUP_M + (pid % group_size) pid_n = (pid % width) // group_size rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) A = A + (rm[:, None] * stride_am + rn[None, :] * stride_an) mask = (rm < M)[:, None] & (rn < N)[None, :] a = tl.load(A, mask=mask) absmax_inv = tl.load(absmax_inv_ptr) # rematerialize to save registers rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) B = B + (rm[:, None] * stride_bm + rn[None, :] * stride_bn) mask = (rm < M)[:, None] & (rn < N)[None, :] output = tl.libdevice.llrint(127.0 * (a * absmax_inv)) tl.store(B, output, mask=mask) def quantize_global_transpose(input): absmax = input.abs().max().unsqueeze(0) absmax_inv = 1.0 / absmax M, N = input.shape out = torch.empty(N, M, device="cuda", dtype=torch.int8) assert out.size(0) == N and out.size(1) == M assert input.stride(0) == 1 or input.stride(1) == 1 assert out.stride(0) == 1 or out.stride(1) == 1 grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) _quantize_global_transpose[grid]( input, absmax_inv, out, input.stride(0), input.stride(1), out.stride(0), out.stride(1), M, N, ) return out, absmax bitsandbytes/bitsandbytes/triton/quantize_rowwise.py000066400000000000000000000041771506752665600236070ustar00rootroot00000000000000import math import torch from bitsandbytes.triton.triton_utils import is_triton_available if not is_triton_available(): def quantize_rowwise(x: torch.Tensor): return None else: import triton import triton.language as tl # rowwise quantize # TODO: autotune this better. @triton.autotune( configs=[ triton.Config({}, num_stages=1, num_warps=8), triton.Config({}, num_stages=2, num_warps=8), triton.Config({}, num_stages=4, num_warps=8), triton.Config({}, num_stages=8, num_warps=8), triton.Config({}, num_stages=1), triton.Config({}, num_stages=2), triton.Config({}, num_stages=4), triton.Config({}, num_stages=8), triton.Config({}, num_warps=1), triton.Config({}, num_warps=2), triton.Config({}, num_warps=4), triton.Config({}, num_warps=8), ], key=["n_elements"], ) @triton.jit def _quantize_rowwise( x_ptr, output_ptr, output_maxs, n_elements, BLOCK_SIZE: tl.constexpr, P2: tl.constexpr, ): pid = tl.program_id(axis=0) block_start = pid * BLOCK_SIZE arange = tl.arange(0, P2) offsets = block_start + arange row_mask = arange < BLOCK_SIZE x = tl.load(x_ptr + offsets, mask=row_mask) abs_x = tl.abs(x) max_val = tl.max(tl.where(row_mask, abs_x, 0), axis=0) output = tl.libdevice.llrint(127.0 * (x / max_val)) tl.store(output_ptr + offsets, output, mask=row_mask) tl.store(output_maxs + pid, max_val) def quantize_rowwise(x: torch.Tensor): output = torch.empty(*x.shape, device=x.device, dtype=torch.int8) output_maxs = torch.empty(x.shape[0], device=x.device, dtype=torch.float16) P2 = int(2 ** (math.ceil(math.log2(x.shape[1])))) assert x.is_cuda and output.is_cuda n_elements = output.numel() grid = lambda meta: (x.shape[0],) _quantize_rowwise[grid](x, output, output_maxs, n_elements, BLOCK_SIZE=x.shape[1], P2=P2) return output, output_maxs bitsandbytes/bitsandbytes/triton/triton_utils.py000066400000000000000000000003721506752665600227200ustar00rootroot00000000000000import functools @functools.lru_cache(None) def is_triton_available(): try: from torch.utils._triton import has_triton, has_triton_package return has_triton_package() and has_triton() except Exception: return False bitsandbytes/bitsandbytes/utils.py000066400000000000000000000153241506752665600200050ustar00rootroot00000000000000import json import shlex import subprocess import torch def outlier_hook(module, input): assert isinstance(module, torch.nn.Linear) tracer = OutlierTracer.get_instance() hvalue = tracer.get_hvalue(module.weight) if hvalue not in tracer.hvalue2outlier_idx: outlier_idx = find_outlier_dims(module.weight) tracer.outliers.append(outlier_idx) tracer.hvalues.append(hvalue) if len(tracer.outliers) > 1: # assign the current layer the outlier idx found from the weight # of the previous linear layer if tracer.outliers[-1].numel() > 0: assert tracer.outliers[-1].max() < module.weight.shape[1] tracer.hvalue2outlier_idx[hvalue] = tracer.outliers[-1] else: # first layer, we cannot use the weight for outlier detection # we follow a mixed approach: # (1) zscore test of std of hidden dimension # (2) magnitude > 6 test merged = input[0].view(-1, input[0].shape[-1]) # (1) zscore test of std of hidden dimension outlier_idx = find_outlier_dims(merged, reduction_dim=1, zscore=3) # (2) magnitude > 6 test dims = (torch.abs(input[0]) > 6).sum(dim=list(range(len(input[0].shape) - 1))) outlier_idx2 = torch.where(dims > 0)[0] outlier_idx = torch.cat([outlier_idx, outlier_idx2]).unique() tracer.hvalue2outlier_idx[hvalue] = outlier_idx else: for hook in tracer.hooks: hook.remove() class OutlierTracer: _instance = None def __init__(self): raise RuntimeError("Call get_instance() instead") def initialize(self, model): self.last_w = None self.current_outlier_dims = None self.hvalues = [] self.outliers = [] self.hvalue2outlier_idx = {} self.initialized = True self.hooks = [] for n, m in model.named_modules(): if isinstance(m, torch.nn.Linear): self.hooks.append(m.register_forward_pre_hook(outlier_hook)) def is_initialized(self): return getattr(self, "initialized", False) def get_hvalue(self, weight): return weight.data.storage().data_ptr() def get_outliers(self, weight): if not self.is_initialized(): print("Outlier tracer is not initialized...") return None hvalue = self.get_hvalue(weight) if hvalue in self.hvalue2outlier_idx: return self.hvalue2outlier_idx[hvalue] else: return None @classmethod def get_instance(cls): if cls._instance is None: cls._instance = cls.__new__(cls) return cls._instance def find_outlier_dims(weight, reduction_dim=0, zscore=4.0, topk=None, rdm=False): if rdm: return torch.randint(0, weight.shape[1], size=(topk,), device=weight.device).long() std = weight.std(reduction_dim) stdm = std.mean() stdstd = std.std() zstd = (std - stdm) / stdstd if topk is not None: val, idx = torch.topk(std.abs(), k=topk, dim=0) else: idx = torch.where(zstd > zscore)[0] return idx def execute_and_return(command_string: str) -> tuple[str, str]: def _decode(subprocess_err_out_tuple): return tuple(to_decode.decode("UTF-8").strip() for to_decode in subprocess_err_out_tuple) def execute_and_return_decoded_std_streams(command_string): return _decode( subprocess.Popen( shlex.split(command_string), stdout=subprocess.PIPE, stderr=subprocess.PIPE, ).communicate(), ) std_out, std_err = execute_and_return_decoded_std_streams(command_string) return std_out, std_err def replace_linear( model, linear_replacement, skip_modules=("lm_head",), copy_weights=False, post_processing_function=None, ): """ Replace linear modules with a new Linear module. Parameters: model (`torch.nn.Module`): Input model or `torch.nn.Module` as the function is run recursively. linear_replacement (`torch.nn.Module`): The linear module that replaces the old one. Only expects standard arguments. If other arguments need to be passed, use a lambda. skip_modules (`List[str]`, *optional*, defaults to `lm_head`): List of modules names not to convert. Defaults to `lm_head`. copy_weights (`bool`): Copy the weights from the old linear module to the new one post_processing_function (`str`): A function name of the replacement linear class that is called after processing. """ for name, module in model.named_children(): if len(list(module.children())) > 0: replace_linear(module, linear_replacement, skip_modules, copy_weights, post_processing_function) if isinstance(module, torch.nn.Linear) and name not in skip_modules: old_module = model._modules[name] model._modules[name] = linear_replacement( module.in_features, module.out_features, module.bias is not None, ) if copy_weights: model._modules[name].weight = old_module.weight model._modules[name].bias = old_module.bias if post_processing_function is not None: func = getattr(module, post_processing_function, None) if func is not None: func(module) return model def pack_dict_to_tensor(source_dict): """ Pack a dictionary into a torch tensor for storing quant_state items in state_dict. Parameters: - source_dict: The dictionary to be packed. Returns: A torch tensor containing the packed data. """ json_str = json.dumps(source_dict) json_bytes = json_str.encode("utf-8") tensor_data = torch.tensor(list(json_bytes), dtype=torch.uint8) return tensor_data def unpack_tensor_to_dict(tensor_data): """ Unpack a torch tensor into a Python dictionary. Parameters: - tensor_data: The torch tensor containing the packed data. Returns: A Python dictionary containing the unpacked data. """ json_bytes = bytes(tensor_data.cpu().numpy()) json_str = json_bytes.decode("utf-8") unpacked_dict = json.loads(json_str) return unpacked_dict LINEAR_8BIT_WEIGHTS_FORMAT_MAPPING = {"row": 0, "col32": 1, "col_turing": 2, "col_ampere": 3} INVERSE_LINEAR_8BIT_WEIGHTS_FORMAT_MAPPING = {val: name for (name, val) in LINEAR_8BIT_WEIGHTS_FORMAT_MAPPING.items()} def sync_gpu(t: torch.Tensor): if t.device.type == "cuda": torch.cuda.synchronize() elif t.device.type == "xpu": torch.xpu.synchronize() bitsandbytes/check_bnb_install.py000066400000000000000000000005151506752665600175720ustar00rootroot00000000000000import torch import bitsandbytes as bnb p = torch.nn.Parameter(torch.rand(10, 10).cuda()) a = torch.rand(10, 10).cuda() p1 = p.data.sum().item() adam = bnb.optim.Adam([p]) out = a * p loss = out.sum() loss.backward() adam.step() p2 = p.data.sum().item() assert p1 != p2 print("SUCCESS!") print("Installation was successful!") bitsandbytes/csrc/000077500000000000000000000000001506752665600145255ustar00rootroot00000000000000bitsandbytes/csrc/common.cpp000066400000000000000000000025141506752665600165230ustar00rootroot00000000000000#include #include void quantize_block(const quantize_block_args& args) { // 1. find absmax in block // 2. divide input value by absmax to normalize into [-1.0, 1.0] // 3. do binary search to find the closest value // 4. check minimal distance // 5. store index // 1. find absmax in block float absmax_block = -FLT_MAX; for (long long i = args.block_idx; i < args.block_end; i++) absmax_block = fmax(absmax_block, fabs(args.A[i])); args.absmax[args.block_idx / args.blocksize] = absmax_block; for (long long i = args.block_idx; i < args.block_end; i++) { // 2. divide input value by absmax to normalize into [-1.0, 1.0] // 3. do binary search to find the closest value float normed_value = args.A[i] / absmax_block; long long idx = args.bin_searcher->scalar(normed_value); // 4. check minimal distance // The binary search returns always the value to the left, which might not be the closest value if (idx < 255) { float dist_left = fabs(normed_value - (args.code[idx])); float dist_right = fabs(normed_value - (args.code[idx + 1])); if (dist_right < dist_left) { idx += 1; } } // 5. store index args.out[i] = (unsigned char)idx; } } bitsandbytes/csrc/common.cuh000066400000000000000000000032631506752665600165220ustar00rootroot00000000000000#pragma once // TODO: Let's make some of these constexpr and put in a namespace. #define BNB_CC_MAXWELL 500 #define BNB_CC_MAXWELL2 520 #define BNB_CC_MAXWELL2_X1 530 #define BNB_CC_PASCAL 600 #define BNB_CC_PASCAL_X2 620 #define BNB_CC_VOLTA 700 #define BNB_CC_VOLTA_XAVIER 720 #define BNB_CC_TURING 750 #define BNB_CC_AMPERE 800 #define BNB_CC_AMPERE2 860 #define BNB_CC_AMPERE2_ORIN 870 #define BNB_CC_ADA 890 #define BNB_CC_HOPPER 900 #define BNB_CC_BLACKWELL 1000 #define BNB_FP16_AVAILABLE (__CUDA_ARCH__ >= BNB_CC_MAXWELL2_X1) #define BNB_FP16_MMA_AVAILABLE (__CUDA_ARCH__ >= BNB_CC_VOLTA) #define BNB_INT8_MMA_AVAILABLE (__CUDA_ARCH__ >= BNB_CC_VOLTA_XAVIER) #define BNB_BF16_AVAILABLE (__CUDA_ARCH__ >= BNB_CC_AMPERE) #define BNB_FP8_AVAILABLE (__CUDA_ARCH__ >= BNB_CC_ADA) #define BNB_WARP_SIZE 32 // The maximum number of resident threads per SM varies by arch. // For A100/H100 and all prior to Turing, it is 2048, which allows // for 2 full blocks of 1024 threads per SM. // Reference: // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#features-and-technical-specifications-technical-specifications-per-compute-capability #if __CUDA_ARCH__ == 750 #define BNB_MAX_THREADS_PER_SM 1024 #elif __CUDA_ARCH__ >= 860 && __CUDA_ARCH__ <= 890 #define BNB_MAX_THREADS_PER_SM 1536 #else #define BNB_MAX_THREADS_PER_SM 2048 #endif // Maximum resident warps per SM is always directly related to the number of threads. #define BNB_MAX_WARPS_PER_SM ((BNB_MAX_THREADS_PER_SM) / (BNB_WARP_SIZE)) // Maximum resident blocks per SM may vary. #if __CUDA_ARCH__ == 860 || __CUDA_ARCH__ == 870 #define BNB_MAX_BLOCKS_PER_SM 16 #else #define BNB_MAX_BLOCKS_PER_SM ((BNB_MAX_WARPS_PER_SM) / 2) #endif bitsandbytes/csrc/common.h000066400000000000000000000006231506752665600161670ustar00rootroot00000000000000#include #ifndef common #define common using namespace BinSearch; struct quantize_block_args { BinAlgo* bin_searcher; float* code; float* A; float* absmax; unsigned char* out; long long block_end; long long block_idx; long long threadidx; long long blocksize; }; void quantize_block(const quantize_block_args& args); #endif bitsandbytes/csrc/common_hip.cuh000066400000000000000000000003271506752665600173600ustar00rootroot00000000000000#pragma once #define BNB_WARP_SIZE warpSize // These are set based on current BNB support for CDNA 2 & RDNA 3. Update as needed for future archs #define BNB_MAX_THREADS_PER_SM 2048 #define BNB_BF16_AVAILABLE true bitsandbytes/csrc/cpu_ops.cpp000066400000000000000000000047131506752665600167060ustar00rootroot00000000000000#include #include #include using namespace BinSearch; void dequantize_cpu(float* code, unsigned char* A, float* absmax, float* out, long long blocksize, long long n) { for (long long block_idx = 0; block_idx < n; block_idx += blocksize) { long long valid_items = n - block_idx >= blocksize ? blocksize : n - block_idx; long long block_end = block_idx + valid_items; for (long long i = block_idx; i < block_end; i++) out[i] = code[A[i]] * absmax[block_idx / blocksize]; } } void quantize_cpu(float* code, float* A, float* absmax, unsigned char* out, long long blocksize, long long n) { // the default code is has range [-0.993, 1.0] which can cause an error in the binary search algorithm used below code[0] = -1.0f; long long num_blocks = n / blocksize; num_blocks += n % blocksize == 0 ? 0 : 1; const uint32 elements_code = 256; BinAlgo bin_searcher(code, elements_code); int thread_wave_size = 256; // we chunk the threads into waves of 256 since the max limit is // between 16k and 64k on Linux (we reach this when running BLOOM-176B with a large batch size) for (long long offset = 0; offset < num_blocks; offset += thread_wave_size) { long long valid_chunks = num_blocks - offset >= thread_wave_size ? thread_wave_size : num_blocks - offset; std::vector threads(valid_chunks); std::vector args(valid_chunks); int chunks_processed = 0; for (long long block_idx = offset * blocksize; block_idx < n; block_idx += blocksize) { long long valid_items = n - block_idx >= blocksize ? blocksize : n - block_idx; long long block_end = block_idx + valid_items; struct quantize_block_args& arg = args[chunks_processed]; arg.bin_searcher = &bin_searcher; arg.code = code; arg.A = A; arg.absmax = absmax; arg.out = out; arg.block_end = block_end; arg.block_idx = block_idx; arg.threadidx = block_idx / blocksize; arg.blocksize = blocksize; threads[chunks_processed] = std::thread([arg] { quantize_block(arg); }); chunks_processed += 1; if (chunks_processed == valid_chunks) { break; } } for (int i = 0; i < valid_chunks; i++) threads[i].join(); } } bitsandbytes/csrc/cpu_ops.h000066400000000000000000000005171506752665600163510ustar00rootroot00000000000000#ifndef BITSANDBYTES_CPU_OPS_H #define BITSANDBYTES_CPU_OPS_H #include #include void quantize_cpu(float* code, float* A, float* absmax, unsigned char* out, long long blocksize, long long n); void dequantize_cpu(float* code, unsigned char* A, float* absmax, float* out, long long blocksize, long long n); #endif bitsandbytes/csrc/kernels.cu000066400000000000000000003734041506752665600165340ustar00rootroot00000000000000// Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include "common.cuh" #include "kernels.cuh" #include #include #include #include #include #include #include #include #include #include #if CCCL_VERSION >= 2008002 #include #define CUB_REDUCTIONOP_MAX \ cuda::maximum<> {} #else #define CUB_REDUCTIONOP_MAX cub::Max() #endif #define HLF_MAX 65504 #define TH 1024 #define NUM 4 #define NUM_BLOCK 4096 __device__ static float fp4_dequantization_lut[8] = { 0.0f, // 0b000 0.005208333333f, // 0b001 0.66666667f, // 0b010 1.0f, // 0b011 0.33333333f, // 0b100 0.5f, // 0b101 0.16666667f, // 0b110 0.25f // 0b111 }; __device__ static float nf4_dequantization_lut[16] = { -1.0f, // 0b0000 -0.6961928009986877f, // 0b0001 -0.5250730514526367f, // 0b0010 -0.39491748809814453f, // 0b0011 -0.28444138169288635f, // 0b0100 -0.18477343022823334f, // 0b0101 -0.09105003625154495f, // 0b0110 0.0f, // 0b0111 0.07958029955625534f, // 0b1000 0.16093020141124725f, // 0b1001 0.24611230194568634f, // 0b1010 0.33791524171829224f, // 0b1011 0.44070982933044434f, // 0b1100 0.5626170039176941f, // 0b1101 0.7229568362236023f, // 0b1110 1.0f // 0b1111 }; // source: https://stackoverflow.com/questions/17399119/how-do-i-use-atomicmax-on-floating-point-values-in-cuda __device__ float atomicMax(float* address, float val) { int* address_as_i = reinterpret_cast(address); int old = *address_as_i, assumed; do { assumed = old; old = atomicCAS(reinterpret_cast(address), assumed, __float_as_int(fmaxf(val, __int_as_float(assumed)))); } while (assumed != old); return __int_as_float(old); } __device__ __forceinline__ float dDequantizeFP4Tree(unsigned char val) { float sign = 1.0f - 2 * ((val & 0b1000) >> 3); return fp4_dequantization_lut[val & 0b111] * sign; } __device__ unsigned char dQuantizeFP4(float x) { // FP4 with bias of 3 // first bit is a sign // subnormals // 0b000 = 0 // 0b001 = 0.0625 // 0b110 = 2 // 0b111 = 3 // 0b100 = 4 // 0b101 = 6 // 0b010 = 8 // 0b011 = 12 // we do a binary search // the pivots are divided by 12 (the FP4 absmax) // since we assume input data is in [-1.0, 1.0] // !be careful here, its easy to make a mistake // that is difficult to notice if you add an extra // zero somewhere! int sign = x < 0 ? 0b1000 : 0b0000; x = fabsf(x); if (x > 0.29166667f) if (x > 0.583333f) if (x > 0.8333333f) return 0b0011 + sign; else return 0b0010 + sign; else if (x > 0.4166667f) return 0b101 + sign; else return 0b100 + sign; else if (x > 0.0859375f) if (x > 0.20833333f) return 0b0111 + sign; else return 0b0110 + sign; else if (x > 0.00260417f) return 0b0001 + sign; else return 0b0000 + sign; } __device__ __forceinline__ float dDequantizeNF4(unsigned char val) { return nf4_dequantization_lut[val & 0x0F]; } __device__ unsigned char dQuantizeNF4(float x) { // the values for this tree was generated by test_normal_map_tree // in the file tests/test_functional.py if (x > 0.03979014977812767f) if (x > 0.3893125355243683f) // 1 if (x > 0.6427869200706482f) // 11 if (x > 0.8614784181118011f) // 111 return 0b1111; else return 0b1110; else if (x > 0.5016634166240692f) // 110 return 0b1101; else return 0b1100; else if (x > 0.2035212516784668f) // 10 if (x > 0.2920137718319893f) // 101 return 0b1011; else return 0b1010; else if (x > 0.1202552504837513f) // 100 return 0b1001; else return 0b1000; else if (x > -0.33967943489551544f) // 0 if (x > -0.13791173323988914f) // 01 if (x > -0.045525018125772476f) // 011 return 0b0111; else return 0b0110; else if (x > -0.23460740596055984f) // 010 return 0b0101; else return 0b0100; else if (x > -0.6106329262256622f) // 00 if (x > -0.4599952697753906f) // 001 return 0b0011; else return 0b0010; else if (x > -0.8480964004993439f) // 000 return 0b0001; else return 0b0000; } // sign function for lion // taken from https://stackoverflow.com/a/4609795, but not sure if there's a proper way to do this in CUDA template __device__ int sgn(T val) { return (T(0) < val) - (val < T(0)); } template __device__ unsigned char dQuantize(float* smem_code, const float rand, float x) { int pivot = 127; int upper_pivot = 255; int lower_pivot = 0; float lower = -1.0f; float upper = 1.0f; float val = smem_code[pivot]; // i>>=1 = {32, 16, 8, 4, 2, 1} for (int i = 64; i > 0; i >>= 1) { if (x > val) { lower_pivot = pivot; lower = val; pivot += i; } else { upper_pivot = pivot; upper = val; pivot -= i; } val = smem_code[pivot]; } if (upper_pivot == 255) upper = smem_code[upper_pivot]; if (lower_pivot == 0) lower = smem_code[lower_pivot]; if (!STOCHASTIC) { if (x > val) { float midpoint = (upper + val) * 0.5f; if (x > midpoint) { return upper_pivot; } else return pivot; } else { float midpoint = (lower + val) * 0.5f; if (x < midpoint) return lower_pivot; else return pivot; } } else { if (x > val) { float dist_to_upper = fabsf(upper - x); float dist_full = upper - val; if (rand >= dist_to_upper / dist_full) return upper_pivot; else return pivot; } else { float dist_to_lower = fabsf(lower - x); float dist_full = val - lower; if (rand >= dist_to_lower / dist_full) return lower_pivot; else return pivot; } } } template __device__ __forceinline__ unsigned char quantize_2D(float* __restrict__ quadrants, float* __restrict__ const smem_code, float x) { int pivot = 127; int upper_pivot = 255; int lower_pivot = 0; float lower = SIGNED ? -1.0f : 0.0f; float upper = 1.0f; float midpoint; float val = quadrants[1]; int local_pivot = 1; int offset = 1; // i>>=1 = {32, 16, 8, 4, 2, 1} for (int i = 64; i > 0; i >>= 1) { if (x > val) { lower_pivot = pivot; lower = val; pivot += i; // val = i == 64 ? quadrants[2] : smem_code[pivot]; local_pivot += offset; } else { upper_pivot = pivot; upper = val; pivot -= i; // val = i == 64 ? quadrants[0] : smem_code[pivot]; local_pivot -= offset; } val = i >= 64 ? quadrants[local_pivot] : smem_code[pivot]; offset -= 1; } if (x > val) { midpoint = (upper + val) * 0.5f; if (x > midpoint) return upper_pivot; else return pivot; } else { midpoint = (lower + val) * 0.5f; if (x < midpoint) return lower_pivot; else return pivot; } } __launch_bounds__(TH, 4) __global__ void kQuantize(float* code, float* __restrict__ const A, unsigned char* out, const int n) { const int n_full = (NUM_BLOCK * (n / NUM_BLOCK)) + (n % NUM_BLOCK == 0 ? 0 : NUM_BLOCK); int valid_items = (blockIdx.x + 1 == gridDim.x) ? n - (blockIdx.x * NUM_BLOCK) : NUM_BLOCK; const int base_idx = (blockIdx.x * NUM_BLOCK); float vals[NUM]; unsigned char qvals[NUM]; // const int lane_id = threadIdx.x % 2; typedef cub::BlockLoad LoadFloat; typedef cub::BlockStore StoreChar; __shared__ typename LoadFloat::TempStorage loadf; __shared__ typename StoreChar::TempStorage storec; __shared__ float smem_code[256]; //__shared__ float smem_code[2][257]; if (threadIdx.x < 256) { smem_code[threadIdx.x] = code[threadIdx.x]; // smem_code[0][threadIdx.x] = code[threadIdx.x]; // smem_code[1][threadIdx.x] = smem_code[0][threadIdx.x]; } for (unsigned int i = base_idx; i < n_full; i += gridDim.x * NUM_BLOCK) { // number of values already processed in blocks + // number of values already processed in this block + // rand_offset % mod value valid_items = n - i > NUM_BLOCK ? NUM_BLOCK : n - i; __syncthreads(); LoadFloat(loadf).Load(&(A[i]), vals, valid_items); #pragma unroll 4 for (int j = 0; j < NUM; j++) qvals[j] = dQuantize<0>(smem_code, 0.0f, vals[j]); __syncthreads(); StoreChar(storec).Store(&(out[i]), qvals, valid_items); } } template //__launch_bounds__(TH, 4) __global__ void kQuantizeBlockwise( float* code, T* __restrict__ const A, float* absmax, unsigned char* out, float* __restrict__ const rand, const int rand_offset, const int n ) { const int n_full = gridDim.x * BLOCK_SIZE; int valid_items = 0; const int base_idx = (blockIdx.x * BLOCK_SIZE); T vals[NUM_PER_TH]; float rand_vals[NUM_PER_TH]; unsigned char qvals[(DATA_TYPE > 0) ? NUM_PER_TH / 2 : NUM_PER_TH]; // float local_abs_max = -FLT_MAX; float local_abs_max = 0.0f; int local_rand_idx = 0; typedef cub::BlockLoad LoadT; typedef cub::BlockStore< unsigned char, BLOCK_SIZE / NUM_PER_TH, (DATA_TYPE > 0) ? NUM_PER_TH / 2 : NUM_PER_TH, cub::BLOCK_STORE_WARP_TRANSPOSE> StoreChar; typedef cub::BlockReduce BlockReduce; typedef cub::BlockLoad LoadFloat; __shared__ typename LoadT::TempStorage loadt; __shared__ typename LoadFloat::TempStorage loadf; __shared__ typename StoreChar::TempStorage storec; __shared__ typename BlockReduce::TempStorage reduce; __shared__ float smem_code[256]; __shared__ float smem_absmax_value[1]; if (DATA_TYPE == General8bit) for (int i = threadIdx.x; i < 256; i += blockDim.x) smem_code[i] = code[i]; for (int i = base_idx; i < n_full; i += gridDim.x * BLOCK_SIZE) { valid_items = n - i > BLOCK_SIZE ? BLOCK_SIZE : n - i; local_abs_max = -FLT_MAX; __syncthreads(); LoadT(loadt).Load(&(A[i]), vals, valid_items, (T)0.0f); // 1. compute local max // 2. broadcast local max // 3. normalize inputs and quantize #pragma unroll NUM_PER_TH for (int j = 0; j < NUM_PER_TH; j++) local_abs_max = fmaxf(local_abs_max, fabsf((float)vals[j])); local_abs_max = BlockReduce(reduce).Reduce(local_abs_max, CUB_REDUCTIONOP_MAX, valid_items); if (threadIdx.x == 0) { smem_absmax_value[0] = 1.0f / local_abs_max; absmax[i / BLOCK_SIZE] = local_abs_max; } __syncthreads(); local_abs_max = smem_absmax_value[0]; if (STOCHASTIC) { local_rand_idx = ((blockIdx.x * NUM_BLOCK) + (threadIdx.x * NUM) + rand_offset) % (1024 - 4); LoadFloat(loadf).Load(&rand[local_rand_idx], rand_vals, BLOCK_SIZE, 0); } switch (DATA_TYPE) { case General8bit: #pragma unroll NUM_PER_TH for (int j = 0; j < NUM_PER_TH; j++) { if (!STOCHASTIC) qvals[j] = dQuantize<0>(smem_code, 0.0f, ((float)vals[j]) * local_abs_max); else qvals[j] = dQuantize<1>(smem_code, rand_vals[j], ((float)vals[j]) * local_abs_max); } break; case FP4: #pragma unroll NUM_PER_TH for (int j = 0; j < NUM_PER_TH / 2; j++) { qvals[j] = dQuantizeFP4(((float)vals[2 * j]) * local_abs_max) << 4; qvals[j] |= dQuantizeFP4(((float)vals[2 * j + 1]) * local_abs_max); } break; case NF4: #pragma unroll NUM_PER_TH for (int j = 0; j < NUM_PER_TH / 2; j++) { qvals[j] = dQuantizeNF4(((float)vals[2 * j]) * local_abs_max) << 4; qvals[j] |= dQuantizeNF4(((float)vals[2 * j + 1]) * local_abs_max); } break; } __syncthreads(); StoreChar(storec).Store( &(out[(DATA_TYPE > 0) ? i / 2 : i]), qvals, (DATA_TYPE > 0) ? (valid_items + 1) / 2 : valid_items ); } } template __global__ void kDequantizeBlockwise(float* code, unsigned char* A, float* absmax, T* out, const int blocksize, const int n) { const int n_load = (gridDim.x * TILE_SIZE); int valid_items_load = 0; int valid_items_store = 0; const int base_idx = (blockIdx.x * TILE_SIZE); T vals[NUM_PER_TH * ((DATA_TYPE > 0) ? 2 : 1)]; unsigned char qvals[NUM_PER_TH]; float local_abs_max = -FLT_MAX; typedef cub::BlockLoad LoadChar; typedef cub::BlockStore 0) ? 2 : 1), cub::BLOCK_STORE_WARP_TRANSPOSE> StoreT; __shared__ typename LoadChar::TempStorage loadchar; __shared__ typename StoreT::TempStorage storet; for (int i = base_idx; i < n_load; i += gridDim.x * TILE_SIZE) { if (DATA_TYPE > 0) { valid_items_load = min(TILE_SIZE, (n + 1) / 2 - i); valid_items_store = min(TILE_SIZE * 2, n - i * 2); } else { valid_items_load = min(TILE_SIZE, n - i); valid_items_store = valid_items_load; } // Since blocksize will always be a power-of-2, we avoid more expensive // division by the blocksize and instead use a shift operation. // This is equivalent to (i+threadId.x*NUM_PER_TH)/blocksize. local_abs_max = __ldg(&absmax[(i + threadIdx.x * NUM_PER_TH) >> (31 - __clz(blocksize))]); __syncthreads(); LoadChar(loadchar).Load(&(A[i]), qvals, valid_items_load, 128); switch (DATA_TYPE) { case General8bit: // load code through read-only cache via __ldg #pragma unroll NUM_PER_TH for (int j = 0; j < NUM_PER_TH; j++) vals[j] = __ldg(&code[qvals[j]]) * local_abs_max; break; case FP4: #pragma unroll NUM_PER_TH for (int j = 0; j < NUM_PER_TH; j++) { vals[j * 2] = dDequantizeFP4Tree(qvals[j] >> 4) * local_abs_max; vals[j * 2 + 1] = dDequantizeFP4Tree(qvals[j] & 0x0F) * local_abs_max; } break; case NF4: #pragma unroll NUM_PER_TH for (int j = 0; j < NUM_PER_TH; j++) { vals[j * 2] = dDequantizeNF4(qvals[j] >> 4) * local_abs_max; vals[j * 2 + 1] = dDequantizeNF4(qvals[j] & 0x0F) * local_abs_max; } break; } __syncthreads(); StoreT(storet).Store(&(out[(DATA_TYPE > 0) ? i * 2 : i]), vals, valid_items_store); } } __global__ void kDequantize(float* code, unsigned char* A, float* out, const int n) { const unsigned int numThreads = blockDim.x * gridDim.x; const int idx = (blockIdx.x * blockDim.x) + threadIdx.x; __shared__ float smem_code[256]; if (threadIdx.x < 256) { smem_code[threadIdx.x] = code[threadIdx.x]; } __syncthreads(); for (int i = idx; i < n; i += numThreads) { out[i] = smem_code[A[i]]; } } template __launch_bounds__(BLOCK_SIZE / NUM_VALS, 1) __global__ void kPreconditionOptimizer32bit2State( T* g, T* p, float* state1, float* state2, float* unorm, const float beta1, const float beta2, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const int n ) { const int n_full = (BLOCK_SIZE * (n / BLOCK_SIZE)) + (n % BLOCK_SIZE == 0 ? 0 : BLOCK_SIZE); const int base_idx = (blockIdx.x * blockDim.x * NUM_VALS); int valid_items = 0; T g_vals[NUM_VALS]; float s1_vals[NUM_VALS]; float s2_vals[NUM_VALS]; const float correction1 = 1.0f / (1.0f - powf(beta1, step)); const float correction2 = 1.0f / (1.0f - powf(beta2, step)); typedef cub::BlockLoad Load; typedef cub::BlockLoad LoadFloat; typedef cub::BlockReduce BlockReduce; __shared__ union { typename Load::TempStorage load; typename LoadFloat::TempStorage loadf; typename BlockReduce::TempStorage reduce; } temp_storage; for (unsigned int i = base_idx; i < n_full; i += gridDim.x * BLOCK_SIZE) { valid_items = n - i >= (BLOCK_SIZE) ? (BLOCK_SIZE) : n - i; __syncthreads(); Load(temp_storage.load).Load(&(g[i]), g_vals, valid_items, 0.0f); __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state1[i]), s1_vals, valid_items, 0.0f); __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state2[i]), s2_vals, valid_items, 0.0f); #pragma unroll NUM_VALS for (unsigned int j = 0; j < NUM_VALS; j++) g_vals[j] = gnorm_scale * ((float)g_vals[j]); #pragma unroll NUM_VALS for (unsigned int j = 0; j < NUM_VALS; j++) { switch (OPTIMIZER) { case ADAM: s1_vals[j] = s1_vals[j] * beta1 + ((1.0f - beta1) * ((float)g_vals[j])); s2_vals[j] = s2_vals[j] * beta2 + ((1.0f - beta2) * (((float)g_vals[j]) * ((float)g_vals[j]))); s1_vals[j] *= correction1; s2_vals[j] *= correction2; s1_vals[j] = s1_vals[j] / (sqrtf(s2_vals[j]) + eps); // update s1_vals[j] *= s1_vals[j]; // update l2 norm (update*update) break; } } #pragma unroll NUM_VALS - 1 for (unsigned int j = 1; j < NUM_VALS; j++) s1_vals[0] += s1_vals[j]; __syncthreads(); s1_vals[0] = BlockReduce(temp_storage.reduce).Sum(s1_vals[0]); if (threadIdx.x == 0) atomicAdd(&unorm[0], s1_vals[0]); __syncwarp(); } } #define NUM_PER_THREAD 4 template __launch_bounds__(TH, 1) __global__ void kOptimizer32bit2State( T* g, T* p, float* state1, float* state2, float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n ) { const int n_full = ((TH * NUM_PER_THREAD) * (n / (TH * NUM_PER_THREAD))) + (n % (TH * NUM_PER_THREAD) == 0 ? 0 : (TH * NUM_PER_THREAD)); const int base_idx = (blockIdx.x * blockDim.x * NUM_PER_THREAD); int valid_items = 0; float update_scale = 0.0f; T g_vals[NUM_PER_THREAD]; T p_vals[NUM_PER_THREAD]; float s1_vals[NUM_PER_THREAD]; float s2_vals[NUM_PER_THREAD]; // AdEMAMix has an additional state buffer, which we packed // into state1. We need thread-local storage here for these. // TODO: Mark with [[maybe_unused]] after upgrade to min compiler. float s3_vals[NUM_PER_THREAD]; const float correction1 = 1.0f - powf(beta1, step); const float correction2 = sqrtf(1.0f - powf(beta2, step)); const float step_size = -lr * correction2 / correction1; if (max_unorm > 0.0f) { update_scale = max_unorm > 0.0f ? sqrtf(unorm[0]) : 1.0f; if (update_scale > max_unorm * param_norm) { update_scale = (max_unorm * param_norm) / update_scale; } else { update_scale = 1.0f; } } else { update_scale = 1.0f; } typedef cub::BlockLoad Load; typedef cub::BlockStore Store; typedef cub::BlockLoad LoadFloat; typedef cub::BlockStore StoreFloat; __shared__ union { typename Load::TempStorage load; typename Store::TempStorage store; typename LoadFloat::TempStorage loadf; typename StoreFloat::TempStorage storef; } temp_storage; for (unsigned int i = base_idx; i < n_full; i += gridDim.x * TH * NUM_PER_THREAD) { valid_items = n - i >= (TH * NUM_PER_THREAD) ? (TH * NUM_PER_THREAD) : n - i; __syncthreads(); Load(temp_storage.load).Load(&(g[i]), g_vals, valid_items); __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state1[i]), s1_vals, valid_items); __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state2[i]), s2_vals, valid_items); __syncthreads(); Load(temp_storage.load).Load(&(p[i]), p_vals, valid_items); // Load additional state1 data for AdEMAMix // TODO: Make constexpr after updating min compiler if (OPTIMIZER == ADEMAMIX) { __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state1[n + i]), s3_vals, valid_items); } #pragma unroll 4 for (unsigned int j = 0; j < NUM_PER_THREAD; j++) g_vals[j] = gnorm_scale * ((float)g_vals[j]); #pragma unroll 4 for (unsigned int j = 0; j < NUM_PER_THREAD; j++) { switch (OPTIMIZER) { case ADEMAMIX: // m1 update: m1 = beta1 * m1 + (1-beta1) * g s1_vals[j] = (s1_vals[j] * beta1) + ((1.0f - beta1) * (float)g_vals[j]); // m2 update: m2 = m2 * beta3 + (1-beta3) * g s3_vals[j] = (s3_vals[j] * beta3) + ((1.0f - beta3) * (float)g_vals[j]); // nu update: nu = beta2 * nu + (1-beta2) * g^2 s2_vals[j] = (s2_vals[j] * beta2) + ((1.0f - beta2) * (float)g_vals[j] * (float)g_vals[j]); p_vals[j] = (float)p_vals[j] - lr * (((s1_vals[j] / correction1) + (alpha * s3_vals[j])) / ((sqrtf(s2_vals[j]) / correction2) + eps)); if (weight_decay > 0.0f) p_vals[j] = ((float)p_vals[j]) * (1.0f - (lr * weight_decay)); break; case ADAM: if (!skip_zeros || (skip_zeros && ((float)g_vals[j] != 0.0f))) { s1_vals[j] = s1_vals[j] * beta1 + ((1.0f - beta1) * ((float)g_vals[j])); s2_vals[j] = s2_vals[j] * beta2 + ((1.0f - beta2) * (((float)g_vals[j]) * ((float)g_vals[j]))); p_vals[j] = ((float)p_vals[j]) + (update_scale * step_size * (s1_vals[j] / (sqrtf(s2_vals[j]) + (eps * correction2)))); if (weight_decay > 0.0f) p_vals[j] = ((float)p_vals[j]) * (1.0f - (lr * weight_decay)); } break; } } __syncthreads(); Store(temp_storage.store).Store(&(p[i]), p_vals, valid_items); __syncthreads(); StoreFloat(temp_storage.storef).Store(&(state1[i]), s1_vals, valid_items); __syncthreads(); StoreFloat(temp_storage.storef).Store(&(state2[i]), s2_vals, valid_items); if (OPTIMIZER == ADEMAMIX) { __syncthreads(); StoreFloat(temp_storage.storef).Store(&(state1[n + i]), s3_vals, valid_items); } } } template __launch_bounds__(BLOCK_SIZE / NUM_VALS, 1) __global__ void kPreconditionOptimizer32bit1State( T* g, T* p, float* state1, float* unorm, const float beta1, const float beta2, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const int n ) { const int n_full = (BLOCK_SIZE * (n / BLOCK_SIZE)) + (n % BLOCK_SIZE == 0 ? 0 : BLOCK_SIZE); const int base_idx = (blockIdx.x * blockDim.x * NUM_VALS); int valid_items = 0; T g_vals[NUM_VALS]; float s1_vals[NUM_VALS]; typedef cub::BlockLoad Load; typedef cub::BlockLoad LoadFloat; typedef cub::BlockReduce BlockReduce; __shared__ union { typename Load::TempStorage load; typename LoadFloat::TempStorage loadf; typename BlockReduce::TempStorage reduce; } temp_storage; for (unsigned int i = base_idx; i < n_full; i += gridDim.x * BLOCK_SIZE) { valid_items = n - i >= (BLOCK_SIZE) ? (BLOCK_SIZE) : n - i; __syncthreads(); Load(temp_storage.load).Load(&(g[i]), g_vals, valid_items, 0.0f); __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state1[i]), s1_vals, valid_items, 0.0f); #pragma unroll NUM_VALS for (unsigned int j = 0; j < NUM_VALS; j++) g_vals[j] = gnorm_scale * ((float)g_vals[j]); #pragma unroll NUM_VALS for (unsigned int j = 0; j < NUM_VALS; j++) { switch (OPTIMIZER) { case MOMENTUM: if (step == 1) s1_vals[j] = (float)g_vals[j]; // state update else s1_vals[j] = s1_vals[j] * beta1 + ((float)g_vals[j]); // state update s1_vals[j] = s1_vals[j] * s1_vals[j]; // update norm break; case LION: s1_vals[j] = s1_vals[j] * beta2 + ((1.0f - beta2) * (float)g_vals[j]); // state update break; case RMSPROP: s1_vals[j] = s1_vals[j] * beta1 + ((1.0f - beta1) * ((float)g_vals[j]) * ((float)g_vals[j])); // state update s1_vals[j] = __fdividef((float)g_vals[j], sqrtf(s1_vals[j]) + eps); // update value s1_vals[j] = s1_vals[j] * s1_vals[j]; // update norm break; case ADAGRAD: s1_vals[j] = s1_vals[j] + ((float)g_vals[j]) * ((float)g_vals[j]); // state update s1_vals[j] = __fdividef((float)g_vals[j], sqrtf(s1_vals[j]) + eps); // update value s1_vals[j] = s1_vals[j] * s1_vals[j]; // update norm break; } } #pragma unroll for (unsigned int j = 1; j < NUM_VALS; j++) s1_vals[0] += s1_vals[j]; __syncthreads(); s1_vals[0] = BlockReduce(temp_storage.reduce).Sum(s1_vals[0], valid_items); if (threadIdx.x == 0) atomicAdd(&unorm[0], s1_vals[0]); __syncwarp(); } } template __launch_bounds__(TH, 1) __global__ void kOptimizer32bit1State( T* g, T* p, float* state1, float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n ) { const int n_full = ((TH * NUM_PER_THREAD) * (n / (TH * NUM_PER_THREAD))) + (n % (TH * NUM_PER_THREAD) == 0 ? 0 : (TH * NUM_PER_THREAD)); const int base_idx = (blockIdx.x * blockDim.x * NUM_PER_THREAD); int valid_items = 0; float update_scale = 0.0f; if (max_unorm > 0.0f) { update_scale = max_unorm > 0.0f ? sqrtf(unorm[0]) : 1.0f; if (update_scale > max_unorm * param_norm + eps) { update_scale = (max_unorm * param_norm + eps) / update_scale; } else { update_scale = 1.0f; } } else { update_scale = 1.0f; } T g_vals[NUM_PER_THREAD]; T p_vals[NUM_PER_THREAD]; float s1_vals[NUM_PER_THREAD]; typedef cub::BlockLoad Load; typedef cub::BlockStore Store; typedef cub::BlockLoad LoadFloat; typedef cub::BlockStore StoreFloat; __shared__ union { typename Load::TempStorage load; typename Store::TempStorage store; typename LoadFloat::TempStorage loadf; typename StoreFloat::TempStorage storef; } temp_storage; for (unsigned int i = base_idx; i < n_full; i += gridDim.x * TH * NUM_PER_THREAD) { valid_items = n - i >= (TH * NUM_PER_THREAD) ? (TH * NUM_PER_THREAD) : n - i; __syncthreads(); Load(temp_storage.load).Load(&(g[i]), g_vals, valid_items); __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state1[i]), s1_vals, valid_items); __syncthreads(); Load(temp_storage.load).Load(&(p[i]), p_vals, valid_items); #pragma unroll 4 for (unsigned int j = 0; j < NUM_PER_THREAD; j++) { g_vals[j] = gnorm_scale * ((float)g_vals[j]); if (weight_decay > 0.0f) g_vals[j] = (float)g_vals[j] + (((float)p_vals[j]) * weight_decay); } #pragma unroll 4 for (unsigned int j = 0; j < NUM_PER_THREAD; j++) { if (!skip_zeros || (skip_zeros && ((float)g_vals[j] != 0.0f))) { switch (OPTIMIZER) { case MOMENTUM: if (step == 1) s1_vals[j] = (float)g_vals[j]; else s1_vals[j] = s1_vals[j] * beta1 + ((float)g_vals[j]); p_vals[j] = ((float)p_vals[j]) + update_scale * (-lr * (s1_vals[j])); break; case LION: p_vals[j] = ((float)p_vals[j]) - update_scale * (lr * sgn(((float)s1_vals[j]) * beta1 + ((1.0f - beta1) * ((float)g_vals[j])))); s1_vals[j] = s1_vals[j] * beta2 + ((1.0f - beta2) * ((float)g_vals[j])); break; case RMSPROP: s1_vals[j] = s1_vals[j] * beta1 + ((1.0f - beta1) * ((float)g_vals[j]) * ((float)g_vals[j])); p_vals[j] = ((float)p_vals[j]) - update_scale * (lr * __fdividef((float)g_vals[j], sqrtf((float)s1_vals[j]) + eps)); break; case ADAGRAD: s1_vals[j] = s1_vals[j] + ((float)g_vals[j]) * ((float)g_vals[j]); p_vals[j] = ((float)p_vals[j]) - lr * __fdividef((float)g_vals[j], sqrtf((float)s1_vals[j]) + eps); break; } } } __syncthreads(); Store(temp_storage.store).Store(&(p[i]), p_vals, valid_items); __syncthreads(); StoreFloat(temp_storage.storef).Store(&(state1[i]), s1_vals, valid_items); } } #define NUM8BIT 16 #define NUM_THREADS 256 #define NUM_PER_BLOCK 4096 template __global__ void __launch_bounds__(NUM_THREADS, 2) kPreconditionOptimizerStatic8bit2State( T* p, T* __restrict__ const g, unsigned char* __restrict__ const state1, unsigned char* __restrict__ const state2, float* unorm, const float beta1, const float beta2, const float eps, const int step, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, const float gnorm_scale, const int n ) { const int n_full = gridDim.x * NUM_PER_BLOCK; const int base_idx = (blockIdx.x * blockDim.x * NUM_PER_THREAD); int valid_items = n - (blockIdx.x * NUM_PER_BLOCK) > NUM_PER_BLOCK ? NUM_PER_BLOCK : n - (blockIdx.x * NUM_PER_BLOCK); float g_val = 0.0f; float local_max_s1 = -FLT_MAX; float local_max_s2 = -FLT_MAX; float local_unorm = 0.0f; float s2_vals[NUM8BIT]; float s1_vals[NUM8BIT]; T g_vals[NUM8BIT]; unsigned char m_c1[NUM8BIT]; unsigned char r_c2[NUM8BIT]; typedef cub::BlockLoad LoadT; typedef cub::BlockLoad LoadUInt8; typedef cub::BlockReduce BlockReduce; __shared__ union { typename LoadT::TempStorage loadh; typename LoadUInt8::TempStorage loadc; typename BlockReduce::TempStorage reduce; } temp_storage; __shared__ float smem_quantiles1[256]; __shared__ float smem_quantiles2[256]; if (threadIdx.x < 256) { smem_quantiles1[threadIdx.x] = quantiles1[threadIdx.x]; smem_quantiles2[threadIdx.x] = quantiles2[threadIdx.x]; } __syncthreads(); for (unsigned int i = base_idx; i < n_full; i += NUM_THREADS * gridDim.x * NUM8BIT) { valid_items = n - i >= (TH * NUM_PER_THREAD) ? (TH * NUM_PER_THREAD) : n - i; LoadT(temp_storage.loadh).Load(&(g[i]), g_vals, valid_items, (T)0.0f); __syncthreads(); LoadUInt8(temp_storage.loadc).Load(&(state1[i]), m_c1, valid_items, 128); __syncthreads(); LoadUInt8(temp_storage.loadc).Load(&(state2[i]), r_c2, valid_items, 128); __syncthreads(); #pragma unroll 16 for (int j = 0; j < NUM8BIT; j++) { g_val = g_vals[j]; g_val *= gnorm_scale; s1_vals[j] = smem_quantiles1[m_c1[j]] * max1[0] * beta1; s1_vals[j] += (1.0f - beta1) * g_val; local_max_s1 = fmaxf(local_max_s1, fabsf(s1_vals[j])); } #pragma unroll 16 for (int j = 0; j < NUM8BIT; j++) { g_val = g_vals[j]; g_val *= gnorm_scale; s2_vals[j] = smem_quantiles2[r_c2[j]] * max2[0] * beta2; s2_vals[j] += (1.0f - beta2) * g_val * g_val; local_max_s2 = fmaxf(local_max_s2, fabsf(s2_vals[j])); } if (unorm != NULL) { #pragma unroll 16 for (int j = 0; j < NUM8BIT; j++) { float correction1 = __fdividef(1.0f, 1.0f - powf(beta1, step)); float correction2 = __fdividef(1.0f, 1.0f - powf(beta2, step)); s1_vals[j] *= correction1; s2_vals[j] *= correction2; float update_val = s1_vals[j] / (sqrtf(s2_vals[j]) + eps); // update local_unorm += update_val * update_val; } } } __syncthreads(); local_max_s1 = BlockReduce(temp_storage.reduce).Reduce(local_max_s1, CUB_REDUCTIONOP_MAX, valid_items); __syncthreads(); local_max_s2 = BlockReduce(temp_storage.reduce).Reduce(local_max_s2, CUB_REDUCTIONOP_MAX, valid_items); if (unorm != NULL) { __syncthreads(); local_unorm = BlockReduce(temp_storage.reduce).Sum(local_unorm, valid_items); } if (threadIdx.x == 0) { atomicMax(&new_max1[0], local_max_s1); atomicMax(&new_max2[0], local_max_s2); if (unorm != NULL) { atomicAdd(&unorm[0], local_unorm); } } } #define NUM_PER_THREAD2 4 #define NUM_THREADS2 1024 #define NUM_PER_BLOCK2 4096 template __global__ void __launch_bounds__(NUM_THREADS2, 1) kOptimizerStatic8bit2State( T* p, T* const g, unsigned char* state1, unsigned char* state2, const float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, float weight_decay, const float gnorm_scale, const int n ) { const int n_full = (blockDim.x * gridDim.x) * NUM_PER_THREAD2; const int base_idx = (blockIdx.x * blockDim.x * NUM_PER_THREAD2); int valid_items = 0; float g_val = 0.0f; float s1_vals[NUM_PER_THREAD2]; float s2_vals[NUM_PER_THREAD2]; const float correction1 = 1.0f - powf(beta1, step); const float correction2 = sqrtf(1.0f - powf(beta2, step)); const float step_size = -lr * correction2 / correction1; // const float step_size = -lr*correction2/correction1; float new_max_val1 = 1.0f / new_max1[0]; float new_max_val2 = 1.0f / new_max2[0]; float update_scale = 1.0f; if (max_unorm > 0.0f) { update_scale = max_unorm > 0.0f ? sqrtf(unorm[0]) : 1.0f; if (update_scale > max_unorm * param_norm) { update_scale = (max_unorm * param_norm) / update_scale; } else { update_scale = 1.0f; } } else { update_scale = 1.0f; } unsigned char c1s[NUM_PER_THREAD2]; unsigned char c2s[NUM_PER_THREAD2]; T p_vals[NUM_PER_THREAD2]; T g_vals[NUM_PER_THREAD2]; typedef cub::BlockLoad LoadT; typedef cub::BlockLoad LoadChar; typedef cub::BlockStore StoreChar; typedef cub::BlockStore StoreT; __shared__ float smem_quantiles1[256]; __shared__ float smem_quantiles2[256]; __shared__ union { typename LoadT::TempStorage loadh; typename LoadChar::TempStorage loadc; typename StoreChar::TempStorage storec; typename StoreT::TempStorage storeh; } temp_storage; if (threadIdx.x < 512) { if (threadIdx.x < 256) smem_quantiles1[threadIdx.x] = quantiles1[threadIdx.x]; else smem_quantiles2[threadIdx.x - 256] = quantiles2[threadIdx.x - 256]; } __syncthreads(); for (unsigned int i = base_idx; i < n_full; i += gridDim.x * NUM_THREADS2 * NUM_PER_THREAD2) { valid_items = n - i >= (TH * NUM_PER_THREAD) ? (TH * NUM_PER_THREAD) : n - i; LoadT(temp_storage.loadh).Load(&(g[i]), g_vals, valid_items, (T)0.0f); __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state1[i]), c1s, valid_items, 128); __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state2[i]), c2s, valid_items, 0); __syncthreads(); LoadT(temp_storage.loadh).Load(&(p[i]), p_vals, valid_items); if ((i + (threadIdx.x * NUM_PER_THREAD2) + NUM_PER_THREAD2) > n) { continue; } #pragma unroll 4 for (unsigned int j = 0; j < NUM_PER_THREAD2; j++) { g_val = float(g_vals[j]); g_val *= gnorm_scale; s1_vals[j] = smem_quantiles1[c1s[j]]; s1_vals[j] = s1_vals[j] * max1[0]; s1_vals[j] = (s1_vals[j] * beta1) + (((1.0f - beta1) * g_val)); c1s[j] = dQuantize<0>(smem_quantiles1, 0.0f, s1_vals[j] * new_max_val1); // make sure state1 term has still the same sign after quantization // (not needed for state2 term which has only positive values) if (signbit(smem_quantiles1[c1s[j]]) != signbit(s1_vals[j])) { if (s1_vals[j] > 0.0f) c1s[j] += 1; else c1s[j] -= 1; } s2_vals[j] = smem_quantiles2[c2s[j]]; s2_vals[j] = s2_vals[j] * max2[0]; s2_vals[j] = (s2_vals[j] * beta2) + (((1.0f - beta2) * g_val * g_val)); c2s[j] = dQuantize<0>(smem_quantiles2, 0.0f, s2_vals[j] * new_max_val2); } #pragma unroll 4 for (unsigned int j = 0; j < NUM_PER_THREAD2; j++) { p_vals[j] = (T)(((float)p_vals[j]) + ((update_scale * step_size * (s1_vals[j] / (sqrtf(s2_vals[j]) + (correction2 * eps)))))); if (weight_decay > 0.0f) p_vals[j] = update_scale * ((float)p_vals[j]) * (1.0f - (lr * weight_decay)); } StoreT(temp_storage.storeh).Store(&(p[i]), p_vals, valid_items); __syncthreads(); StoreChar(temp_storage.storec).Store(&(state1[i]), c1s, valid_items); __syncthreads(); StoreChar(temp_storage.storec).Store(&(state2[i]), c2s, valid_items); __syncthreads(); } } template __global__ void __launch_bounds__(NUM_THREADS, 2) kPreconditionOptimizerStatic8bit1State( T* p, T* __restrict__ const g, unsigned char* __restrict__ const state1, float* unorm, const float beta1, const float beta2, const float eps, const int step, float* __restrict__ const quantiles1, float* max1, float* new_max1, const float weight_decay, const float gnorm_scale, const int n ) { const int n_full = gridDim.x * NUM_PER_BLOCK; const int base_idx = (blockIdx.x * blockDim.x * NUM_PER_THREAD); int valid_items = n - (blockIdx.x * NUM_PER_BLOCK) > NUM_PER_BLOCK ? NUM_PER_BLOCK : n - (blockIdx.x * NUM_PER_BLOCK); float g_val = 0.0f; float local_max_s1 = -FLT_MAX; float local_unorm = 0.0f; float s1_vals[NUM8BIT]; T g_vals[NUM8BIT]; unsigned char m_c1[NUM8BIT]; typedef cub::BlockLoad LoadT; typedef cub::BlockLoad LoadUInt8; typedef cub::BlockReduce BlockReduce; __shared__ union { typename LoadT::TempStorage loadh; typename LoadUInt8::TempStorage loadc; typename BlockReduce::TempStorage reduce; } temp_storage; __shared__ float smem_quantiles1[256]; if (threadIdx.x < 256) smem_quantiles1[threadIdx.x] = quantiles1[threadIdx.x]; __syncthreads(); for (unsigned int i = base_idx; i < n_full; i += gridDim.x * NUM_THREADS * NUM8BIT) { valid_items = n - i >= (TH * NUM_PER_THREAD) ? (TH * NUM_PER_THREAD) : n - i; __syncthreads(); LoadT(temp_storage.loadh).Load(&(g[i]), g_vals, valid_items, (T)0.0f); __syncthreads(); LoadUInt8(temp_storage.loadc).Load(&(state1[i]), m_c1, valid_items, 128); #pragma unroll 16 for (int j = 0; j < NUM8BIT; j++) { g_val = g_vals[j]; g_val *= gnorm_scale; s1_vals[j] = smem_quantiles1[m_c1[j]] * max1[0]; switch (OPTIMIZER) { case ADAGRAD: case MOMENTUM: if (step == 1) s1_vals[j] = (float)g_vals[j]; else s1_vals[j] = s1_vals[j] * beta1 + ((float)g_vals[j]); if (unorm != NULL) local_unorm += s1_vals[j] * s1_vals[j]; break; case LION: s1_vals[j] = s1_vals[j] * beta2 + ((1.0f - beta2) * g_val); break; case RMSPROP: s1_vals[j] = s1_vals[j] * beta1 + ((1.0f - beta1) * (g_val * g_val)); break; } local_max_s1 = fmaxf(local_max_s1, fabsf(s1_vals[j])); } } __syncthreads(); local_max_s1 = BlockReduce(temp_storage.reduce).Reduce(local_max_s1, CUB_REDUCTIONOP_MAX, valid_items); if (threadIdx.x == 0) { atomicMax(&new_max1[0], local_max_s1); } if (unorm != NULL) { __syncthreads(); local_unorm = BlockReduce(temp_storage.reduce).Sum(local_unorm, valid_items); if (threadIdx.x == 0) { atomicAdd(&unorm[0], local_unorm); } } } template __global__ void __launch_bounds__(1024, 1) kOptimizerStatic8bit1State( T* p, T* const g, unsigned char* state1, const float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* max1, float* new_max1, float weight_decay, const float gnorm_scale, const int n ) { const int n_full = (blockDim.x * gridDim.x) * NUM_PER_THREAD2; const int base_idx = (blockIdx.x * blockDim.x * NUM_PER_THREAD2); int valid_items = 0; float g_val = 0.0f; float s1_vals[NUM_PER_THREAD2]; float new_max_val1 = 1.0f / new_max1[0]; float update_scale = 1.0f; if (max_unorm > 0.0f) { update_scale = max_unorm > 0.0f ? sqrtf(unorm[0]) : 1.0f; if (update_scale > max_unorm * param_norm) { update_scale = (max_unorm * param_norm) / update_scale; } else { update_scale = 1.0f; } } else { update_scale = 1.0f; } unsigned char c1s[NUM_PER_THREAD2]; T p_vals[NUM_PER_THREAD2]; T g_vals[NUM_PER_THREAD2]; typedef cub::BlockLoad LoadT; typedef cub::BlockLoad LoadChar; typedef cub::BlockStore StoreChar; typedef cub::BlockStore StoreT; __shared__ float smem_quantiles1[256]; __shared__ union { typename LoadT::TempStorage loadh; typename LoadChar::TempStorage loadc; typename StoreChar::TempStorage storec; typename StoreT::TempStorage storeh; } temp_storage; if (threadIdx.x < 256) smem_quantiles1[threadIdx.x] = quantiles1[threadIdx.x]; __syncthreads(); for (unsigned int i = base_idx; i < n_full; i += gridDim.x * NUM_THREADS2 * NUM_PER_THREAD2) { valid_items = n - i >= (TH * NUM_PER_THREAD) ? (TH * NUM_PER_THREAD) : n - i; LoadT(temp_storage.loadh).Load(&(g[i]), g_vals, valid_items, (T)0.0f); __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state1[i]), c1s, valid_items, 128); __syncthreads(); LoadT(temp_storage.loadh).Load(&(p[i]), p_vals, valid_items); if ((i + (threadIdx.x * NUM_PER_THREAD2) + NUM_PER_THREAD2) > n) { continue; } #pragma unroll 4 for (unsigned int j = 0; j < NUM_PER_THREAD2; j++) { g_val = float(g_vals[j]); g_val *= gnorm_scale; if (weight_decay > 0.0f) { switch (OPTIMIZER) { case ADAGRAD: case MOMENTUM: case RMSPROP: g_val += ((float)p_vals[j]) * weight_decay; break; case LION: p_vals[j] = ((float)p_vals[j]) * (1.0f - lr * weight_decay); break; } } s1_vals[j] = smem_quantiles1[c1s[j]] * max1[0]; switch (OPTIMIZER) { case ADAGRAD: case MOMENTUM: if (step == 1) s1_vals[j] = g_vals[j]; else s1_vals[j] = s1_vals[j] * beta1 + ((float)g_vals[j]); p_vals[j] = ((float)p_vals[j]) + (-lr * update_scale * (s1_vals[j])); break; case LION: p_vals[j] = ((float)p_vals[j]) - (lr * sgn(((float)s1_vals[j]) * beta1 + ((1.0f - beta1) * ((float)g_val)))); s1_vals[j] = s1_vals[j] * beta2 + ((1.0f - beta2) * g_val); break; case RMSPROP: s1_vals[j] = s1_vals[j] * beta1 + ((1.0f - beta1) * (g_val * g_val)); p_vals[j] = ((float)p_vals[j]) - (lr * __fdividef(g_val, sqrtf(s1_vals[j]) + eps)); break; } c1s[j] = dQuantize<0>(smem_quantiles1, 0.0f, s1_vals[j] * new_max_val1); // make sure state1 term has still the same sign after quantization if (signbit(smem_quantiles1[c1s[j]]) != signbit(s1_vals[j])) { if (s1_vals[j] > 0.0f) c1s[j] += 1; else c1s[j] -= 1; } } StoreT(temp_storage.storeh).Store(&(p[i]), p_vals, valid_items); __syncthreads(); StoreChar(temp_storage.storec).Store(&(state1[i]), c1s, valid_items); __syncthreads(); } } template __global__ void kPercentileClipping(T* __restrict__ g, float* gnorm_vec, int step, const int n) { const int n_full = (BLOCK_SIZE * (n / BLOCK_SIZE)) + (n % BLOCK_SIZE == 0 ? 0 : BLOCK_SIZE); int valid_items = 0; typedef cub::BlockReduce BlockReduce; typedef cub::BlockLoad LoadT; __shared__ typename BlockReduce::TempStorage reduce; __shared__ typename LoadT::TempStorage loadT; T vals[NUM_VALS]; float local_sum = 0.0f; for (unsigned int i = (blockIdx.x * BLOCK_SIZE); i < n_full; i += gridDim.x * BLOCK_SIZE) { valid_items = n - i > BLOCK_SIZE ? BLOCK_SIZE : n - i; local_sum = 0.0f; __syncthreads(); LoadT(loadT).Load(&(g[i]), vals, valid_items, (T)0.0f); #pragma unroll NUM_VALS for (int j = 0; j < NUM_VALS; j++) local_sum += ((float)vals[j]) * ((float)vals[j]); local_sum = BlockReduce(reduce).Sum(local_sum, valid_items); if (threadIdx.x == 0) { if (step == 1) { // initialize with the same norm for all positions // #pragma unroll 10 for (int j = 0; j < 100; j++) atomicAdd(&gnorm_vec[j], local_sum); } else atomicAdd(&gnorm_vec[step % 100], local_sum); } } } #define LANES 2 #define QUAD 3 template __launch_bounds__(256, 3) __global__ void kOptimizerStatic8bit2StateBlockwise( T* p, T* __restrict__ const g, unsigned char* state1, unsigned char* state2, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* absmax1, float* absmax2, float weight_decay, const float gnorm_scale, const bool skip_zeros, const int n ) { // const int n_full = n + (n%BLOCK_SIZE); const int n_full = gridDim.x * BLOCK_SIZE; const int base_idx = (blockIdx.x * BLOCK_SIZE); int valid_items = 0; float g_val = 0.0f; float s1_vals[N_PER_TH]; float s2_vals[N_PER_TH]; float s3_vals[N_PER_TH]; // 2-5% const float correction1 = 1.0f - __powf(beta1, step); const float correction2 = sqrtf(1.0f - __powf(beta2, step)); const float step_size = __fdividef(-lr * correction2, correction1); const int lane_id = threadIdx.x % LANES; float new_local_abs_max1 = -FLT_MAX; float new_local_abs_max2 = -FLT_MAX; float new_local_abs_max3 = -FLT_MAX; float quadrants1[QUAD]; float quadrants2[QUAD]; unsigned char c1s[N_PER_TH]; unsigned char c2s[N_PER_TH]; unsigned char c3s[N_PER_TH]; T g_vals[N_PER_TH]; T p_vals[N_PER_TH]; typedef cub::BlockLoad LoadT; typedef cub::BlockLoad LoadChar; typedef cub::BlockStore StoreChar; typedef cub::BlockStore StoreT; __shared__ float smem_quantiles1[LANES][257]; __shared__ float smem_quantiles2[LANES][257]; typedef cub::BlockReduce BlockReduce1; typedef cub::BlockReduce BlockReduce2; typedef cub::BlockReduce BlockReduce3; __shared__ typename BlockReduce1::TempStorage reduce1; __shared__ typename BlockReduce2::TempStorage reduce2; __shared__ typename BlockReduce2::TempStorage reduce3; __shared__ float smem_exchange1[1]; __shared__ float smem_exchange2[1]; __shared__ float smem_exchange3[1]; // [[maybe_unused]] __shared__ union { typename LoadT::TempStorage loadh; typename LoadChar::TempStorage loadc; typename StoreChar::TempStorage storec; typename StoreT::TempStorage storeh; } temp_storage; // init: 0.2 -> 0.23 // 0.23 -> 0.23 smem_quantiles1[0][threadIdx.x] = quantiles1[threadIdx.x]; smem_quantiles2[0][threadIdx.x] = quantiles2[threadIdx.x]; #pragma unroll for (unsigned int j = 1; j < LANES; j++) { smem_quantiles1[j][threadIdx.x] = smem_quantiles1[0][threadIdx.x]; smem_quantiles2[j][threadIdx.x] = smem_quantiles2[0][threadIdx.x]; } __syncthreads(); #pragma unroll for (int k = 0; k < QUAD; k++) { quadrants1[k] = smem_quantiles1[lane_id][(k * 256 / (QUAD + 1)) + (256 / (QUAD + 1) - 1)]; quadrants2[k] = smem_quantiles2[lane_id][(k * 256 / (QUAD + 1)) + (256 / (QUAD + 1) - 1)]; } for (unsigned int i = base_idx; i < n_full; i += gridDim.x * BLOCK_SIZE) { // loads: 0.23 -> 0.85/1.44 valid_items = n - i >= BLOCK_SIZE ? BLOCK_SIZE : n - i; __syncthreads(); LoadT(temp_storage.loadh).Load(&(g[i]), g_vals, valid_items, (T)0.0f); __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state1[i]), c1s, valid_items, 128); __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state2[i]), c2s, valid_items, 0); // AdEMAMix has an additional state packed into state1. if (OPTIMIZER == ADEMAMIX) { __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state1[n + i]), c3s, valid_items, 128); } new_local_abs_max1 = -FLT_MAX; new_local_abs_max2 = -FLT_MAX; new_local_abs_max3 = -FLT_MAX; // update: 2.48/1.57 -> 2.51/1.60 #pragma unroll N_PER_TH for (unsigned int j = 0; j < N_PER_TH; j++) { if (!isnan((float)g_vals[j]) && !isinf((float)g_vals[j])) { s2_vals[j] = smem_quantiles2[lane_id][c2s[j]] * absmax2[i / BLOCK_SIZE]; g_val = g_vals[j]; // float ratio = (g_val*g_val)/fmaxf(s2_vals[j], eps*eps); // g_val = ratio > 2.0f ? 2.0f*g_val/ratio : g_val; g_val *= gnorm_scale; s2_vals[j] = (s2_vals[j] * beta2) + (((1.0f - beta2) * g_val * g_val)); s1_vals[j] = smem_quantiles1[lane_id][c1s[j]] * absmax1[i / BLOCK_SIZE]; s1_vals[j] = (s1_vals[j] * beta1) + (((1.0f - beta1) * g_val)); if (OPTIMIZER == ADEMAMIX) { // The absmax for the third state is appended to absmax1 s3_vals[j] = smem_quantiles1[lane_id][c3s[j]] * absmax1[(n + i) / BLOCK_SIZE]; s3_vals[j] = (s3_vals[j] * beta3) + (((1.0f - beta3) * g_val)); } } else { s1_vals[j] = 0.0f; s2_vals[j] = 0.0f; if (OPTIMIZER == ADEMAMIX) { s3_vals[j] = 0.0f; } } new_local_abs_max1 = fmaxf(new_local_abs_max1, fabsf(s1_vals[j])); new_local_abs_max2 = fmaxf(new_local_abs_max2, fabsf(s2_vals[j])); if (OPTIMIZER == ADEMAMIX) { new_local_abs_max3 = fmaxf(new_local_abs_max3, fabsf(s3_vals[j])); } } // reduce: 2.51/1.60 -> 2.67/1.69 new_local_abs_max1 = BlockReduce1(reduce1).Reduce(new_local_abs_max1, CUB_REDUCTIONOP_MAX); new_local_abs_max2 = BlockReduce2(reduce2).Reduce(new_local_abs_max2, CUB_REDUCTIONOP_MAX); if (OPTIMIZER == ADEMAMIX) { new_local_abs_max3 = BlockReduce3(reduce3).Reduce(new_local_abs_max3, CUB_REDUCTIONOP_MAX); } if (threadIdx.x == 0) { smem_exchange1[0] = new_local_abs_max1; smem_exchange2[0] = new_local_abs_max2; if (OPTIMIZER == ADEMAMIX) { smem_exchange3[0] = new_local_abs_max3; } } __syncthreads(); if (threadIdx.x == 0) { absmax1[i / BLOCK_SIZE] = new_local_abs_max1; absmax2[i / BLOCK_SIZE] = new_local_abs_max2; if (OPTIMIZER == ADEMAMIX) { absmax1[(n + i) / BLOCK_SIZE] = new_local_abs_max3; } } else { new_local_abs_max1 = smem_exchange1[0]; new_local_abs_max2 = smem_exchange2[0]; if (OPTIMIZER == ADEMAMIX) { new_local_abs_max3 = smem_exchange3[0]; } } __syncthreads(); LoadT(temp_storage.loadh).Load(&(p[i]), p_vals, valid_items, (T)0.0f); // reduce: 2.67/1.69 -> 2.67/1.70 #pragma unroll N_PER_TH for (unsigned int j = 0; j < N_PER_TH; j++) { // if(!skip_zeros || (skip_zeros && ((float)g_vals[j] != 0.0f))) if (!isnan((float)g_vals[j]) && !isinf((float)g_vals[j])) { if (OPTIMIZER == ADEMAMIX) { p_vals[j] = T((float)p_vals[j] - lr * (((s1_vals[j] / correction1) + (alpha * s3_vals[j])) / ((sqrtf(s2_vals[j]) / correction2) + eps))); } else { p_vals[j] = (T)(((float)p_vals[j]) + ((step_size * (__fdividef(s1_vals[j], (sqrtf(s2_vals[j]) + (correction2 * eps))))))); } if (weight_decay > 0.0f) p_vals[j] = ((float)p_vals[j]) * (1.0f - (lr * weight_decay)); } } // store: 0.85/1.44 -> 2.48/1.57 __syncthreads(); StoreT(temp_storage.storeh).Store(&(p[i]), p_vals, valid_items); // quantizaztion: 2.67/1.70 -> 3.4/3.3 #pragma unroll N_PER_TH for (unsigned int j = 0; j < N_PER_TH; j++) { c1s[j] = quantize_2D<1>(quadrants1, smem_quantiles1[lane_id], __fdividef(s1_vals[j], new_local_abs_max1)); c2s[j] = quantize_2D<0>(quadrants2, smem_quantiles2[lane_id], __fdividef(s2_vals[j], new_local_abs_max2)); // make sure state1 term has still the same sign after quantization // (not needed for state2 term which has only positive values) if (signbit(smem_quantiles1[lane_id][c1s[j]]) != signbit(s1_vals[j])) { if (s1_vals[j] > 0.0f) c1s[j] += 1; else c1s[j] -= 1; } if (OPTIMIZER == ADEMAMIX) { c3s[j] = quantize_2D<1>(quadrants1, smem_quantiles1[lane_id], __fdividef(s3_vals[j], new_local_abs_max3)); if (signbit(smem_quantiles1[lane_id][c3s[j]]) != signbit(s3_vals[j])) { c3s[j] += (s3_vals[j] > 0.0f) ? 1 : -1; } } } __syncthreads(); StoreChar(temp_storage.storec).Store(&(state1[i]), c1s, valid_items); __syncthreads(); StoreChar(temp_storage.storec).Store(&(state2[i]), c2s, valid_items); if (OPTIMIZER == ADEMAMIX) { __syncthreads(); StoreChar(temp_storage.storec).Store(&(state1[n + i]), c3s, valid_items); } } } #define LANES 2 #define QUAD 3 template __launch_bounds__(256, 3) __global__ void kOptimizerStatic8bit1StateBlockwise( T* p, T* __restrict__ const g, unsigned char* state1, const float beta1, const float beta2, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* absmax1, float weight_decay, const float gnorm_scale, const bool skip_zeros, const int n ) { // const int n_full = n + (n%BLOCK_SIZE); const int n_full = gridDim.x * BLOCK_SIZE; const int base_idx = (blockIdx.x * BLOCK_SIZE); int valid_items = 0; float g_val = 0.0f; float s1_vals[N_PER_TH]; // 2-5% const int lane_id = threadIdx.x % LANES; float new_local_abs_max1 = -FLT_MAX; float quadrants1[QUAD]; unsigned char c1s[N_PER_TH]; T g_vals[N_PER_TH]; T p_vals[N_PER_TH]; typedef cub::BlockLoad LoadT; typedef cub::BlockLoad LoadChar; typedef cub::BlockStore StoreChar; typedef cub::BlockStore StoreT; __shared__ float smem_quantiles1[LANES][257]; typedef cub::BlockReduce BlockReduce1; __shared__ typename BlockReduce1::TempStorage reduce1; __shared__ float smem_exchange1[1]; __shared__ union { typename LoadT::TempStorage loadh; typename LoadChar::TempStorage loadc; typename StoreChar::TempStorage storec; typename StoreT::TempStorage storeh; } temp_storage; // init: 0.2 -> 0.23 // 0.23 -> 0.23 smem_quantiles1[0][threadIdx.x] = quantiles1[threadIdx.x]; #pragma unroll for (unsigned int j = 1; j < LANES; j++) smem_quantiles1[j][threadIdx.x] = smem_quantiles1[0][threadIdx.x]; __syncthreads(); #pragma unroll for (int k = 0; k < QUAD; k++) quadrants1[k] = smem_quantiles1[lane_id][(k * 256 / (QUAD + 1)) + (256 / (QUAD + 1) - 1)]; for (unsigned int i = base_idx; i < n_full; i += gridDim.x * BLOCK_SIZE) { // loads: 0.23 -> 0.85/1.44 valid_items = n - i >= BLOCK_SIZE ? BLOCK_SIZE : n - i; __syncthreads(); LoadT(temp_storage.loadh).Load(&(g[i]), g_vals, valid_items, (T)0.0f); __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state1[i]), c1s, valid_items, 128); __syncthreads(); LoadT(temp_storage.loadh).Load(&(p[i]), p_vals, valid_items, (T)0.0f); new_local_abs_max1 = -FLT_MAX; // update: 2.48/1.57 -> 2.51/1.60 #pragma unroll N_PER_TH for (unsigned int j = 0; j < N_PER_TH; j++) { g_val = float(g_vals[j]); g_val *= gnorm_scale; if (!skip_zeros || (skip_zeros && ((float)g_vals[j] != 0.0f))) { if (weight_decay > 0.0f) { switch (OPTIMIZER) { case MOMENTUM: case ADAGRAD: case RMSPROP: g_val += ((float)p_vals[j]) * weight_decay; break; case LION: p_vals[j] = ((float)p_vals[j]) * (1.0f - lr * weight_decay); break; } } s1_vals[j] = smem_quantiles1[lane_id][c1s[j]] * absmax1[i / BLOCK_SIZE]; switch (OPTIMIZER) { case MOMENTUM: if (step == 1) s1_vals[j] = g_val; else s1_vals[j] = (s1_vals[j] * beta1) + g_val; break; case LION: // here, using gvals[j] to store the gradient smoothed by beta1 for the following parameter update, // before the momentum is updated by beta2 g_vals[j] = lr * sgn(((float)s1_vals[j]) * beta1 + ((1.0f - beta1) * g_val)); s1_vals[j] = s1_vals[j] * beta2 + ((1.0f - beta2) * g_val); break; case RMSPROP: s1_vals[j] = s1_vals[j] * beta1 + ((1.0f - beta1) * (g_val * g_val)); break; case ADAGRAD: s1_vals[j] = s1_vals[j] + (g_val * g_val); break; } } new_local_abs_max1 = fmaxf(new_local_abs_max1, fabsf(s1_vals[j])); } // reduce: 2.51/1.60 -> 2.67/1.69 new_local_abs_max1 = BlockReduce1(reduce1).Reduce(new_local_abs_max1, CUB_REDUCTIONOP_MAX); if (threadIdx.x == 0) smem_exchange1[0] = new_local_abs_max1; __syncthreads(); if (threadIdx.x == 0) absmax1[i / BLOCK_SIZE] = new_local_abs_max1; else new_local_abs_max1 = smem_exchange1[0]; // reduce: 2.67/1.69 -> 2.67/1.70 #pragma unroll N_PER_TH for (unsigned int j = 0; j < N_PER_TH; j++) { if (!skip_zeros || (skip_zeros && ((float)g_vals[j] != 0.0f))) { switch (OPTIMIZER) { case MOMENTUM: p_vals[j] = ((float)p_vals[j]) - lr * (s1_vals[j]); break; case LION: p_vals[j] = ((float)p_vals[j]) - ((float)g_vals[j]); break; case RMSPROP: g_val = g_vals[j]; p_vals[j] = ((float)p_vals[j]) - lr * (__fdividef(g_val, sqrtf(s1_vals[j]) + eps)); break; case ADAGRAD: g_val = g_vals[j]; p_vals[j] = ((float)p_vals[j]) - lr * (__fdividef(g_val, sqrtf(s1_vals[j]) + eps)); break; } } } // store: 0.85/1.44 -> 2.48/1.57 __syncthreads(); StoreT(temp_storage.storeh).Store(&(p[i]), p_vals, valid_items); // quantizaztion: 2.67/1.70 -> 3.4/3.3 #pragma unroll N_PER_TH for (unsigned int j = 0; j < N_PER_TH; j++) { c1s[j] = quantize_2D<1>(quadrants1, smem_quantiles1[lane_id], __fdividef(s1_vals[j], new_local_abs_max1)); // make sure state1 term has still the same sign after quantization // (not needed for state2 term which has only positive values) if (signbit(smem_quantiles1[lane_id][c1s[j]]) != signbit(s1_vals[j])) { if (s1_vals[j] > 0.0f) c1s[j] += 1; else c1s[j] -= 1; } } __syncthreads(); StoreChar(temp_storage.storec).Store(&(state1[i]), c1s, valid_items); } } // Inputs: // A [rows, cols] // Outputs: // rowStats [rows] // out [rows, cols] template __launch_bounds__(1024, BNB_MAX_THREADS_PER_SM / 1024) __global__ void kInt8VectorQuant(T* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols) { // For sm50/sm52 and CUDA < 12.2 we need to do the reduction in fp32. // Otherwise `T` is `fp16`. This can be removed when Maxwell is dropped. #if (__CUDACC_VER_MAJOR__ >= 12 && __CUDACC_VER_MINOR >= 2) || BNB_FP16_AVAILABLE using TReduction = T; #else using TReduction = float; #endif using BlockReduceT = cub::BlockReduce; // One block per row. // Threads load column values in a striped arrangement. // e.g. t0 reads row[0], row[0+nthreads], .. // and t1 reads row[1], row[1+nthreads], .. // Each thread will determine its local absmax. // We then do a blockwise reduction to determine the row's absmax. __shared__ typename BlockReduceT::TempStorage temp_storage; __shared__ TReduction smem_row_absmax; const int row_id = blockIdx.x; const T* row_data = A + (row_id * cols); // Threads will read the row values in a striped access pattern and find a local absmax. TReduction row_local_absmax = -FLT_MIN; for (int i = threadIdx.x; i < cols; i += THREADS) { const TReduction absval = fabsf(__ldcs(&(row_data[i]))); // For sparse decomposition, values outside of the threshold are not to be // included when calculating the row's absmax. if constexpr (SPARSE_DECOMP) { row_local_absmax = fmaxf(row_local_absmax, absval < TReduction(threshold) ? absval : row_local_absmax); } else { row_local_absmax = fmaxf(row_local_absmax, absval); } } // Reduce thread-local absmax across the block. const TReduction row_absmax = BlockReduceT(temp_storage).Reduce(row_local_absmax, CUB_REDUCTIONOP_MAX, cols); if (threadIdx.x == 0) { // Save our block's absmax to shared memory for the quantization step. rowStats[row_id] = smem_row_absmax = row_absmax; } __syncthreads(); // Quantize row-wise. const float scale = __fdividef(127.0f, smem_row_absmax); for (int i = threadIdx.x; i < cols; i += THREADS) { float val = row_data[i]; if constexpr (SPARSE_DECOMP) { // For sparse decomposition, we do not want to quantize the outliers. // Instead they're zeroed out. out[row_id * cols + i] = fabs(val) < threshold ? __float2int_rn(val * scale) : 0; } else { out[row_id * cols + i] = __float2int_rn(val * scale); } } } template __launch_bounds__(1024, BNB_MAX_THREADS_PER_SM / 1024) __global__ void kgetRowStats(T* __restrict__ A, float* rowStats, float threshold, int rows, int cols) { using BlockReduceT = cub::BlockReduce; // One block per row. // Threads load column values in a striped arrangement. // e.g. t0 reads row[0], row[0+nthreads], .. // and t1 reads row[1], row[1+nthreads], .. // Each thread will determine its local absmax. // We then do a blockwise reduction to determine the row's absmax. __shared__ typename BlockReduceT::TempStorage temp_storage; const int row_id = blockIdx.x; const T* __restrict__ row_data = A + (row_id * cols); // Threads will read the row values in a striped access pattern and find a local absmax. float row_local_absmax = -FLT_MIN; for (int i = threadIdx.x; i < cols; i += THREADS) { const float absval = fabsf(row_data[i]); // For sparse decomposition, values outside of the threshold are not to be // included when calculating the row's absmax. if constexpr (SPARSE_DECOMP) { row_local_absmax = fmaxf(row_local_absmax, absval < threshold ? absval : row_local_absmax); } else { row_local_absmax = fmaxf(row_local_absmax, absval); } } // Reduce thread-local absmax across the block. // TODO: Consider algorithm BLOCK_REDUCE_RAKING_COMMUTATIVE_ONLY const float row_absmax = BlockReduceT(temp_storage).Reduce(row_local_absmax, CUB_REDUCTIONOP_MAX, cols); if (threadIdx.x == 0) { // Save our block's absmax to shared memory for the quantization step. rowStats[row_id] = row_absmax; } } template __global__ void kgetRowStats(half* __restrict__ A, float* rowStats, float threshold, int rows, int cols); template __global__ void kgetRowStats(half* __restrict__ A, float* rowStats, float threshold, int rows, int cols); template __global__ void kInt8VectorQuant( half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols ); template __global__ void kInt8VectorQuant( half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols ); #define MM_DEQUANT_CONST 6.200012e-05f // 1.0f/(127.0f*127.0f) template __global__ void kdequant_mm_int32_fp16( int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, half* out, half* __restrict__ const bias, const int numRows, const int numCols, const int n ) { const int n_out = numRows * numCols; int block_offset = blockIdx.x * THREADS * ITEMS_PER_THREAD; int thread_offset = threadIdx.x * ITEMS_PER_THREAD; int local_values[ITEMS_PER_THREAD]; half local_output[ITEMS_PER_THREAD]; float local_rowStats[ITEMS_PER_THREAD]; float local_colStats[ITEMS_PER_THREAD]; float local_biasValue[ITEMS_PER_THREAD]; typedef cub::BlockLoad LoadInt32; __shared__ typename LoadInt32::TempStorage loadint32; int row_idx, col_idx; #pragma unroll ITEMS_PER_THREAD for (int j = 0; j < ITEMS_PER_THREAD; ++j) { row_idx = (block_offset + thread_offset + j) / numCols; col_idx = (block_offset + thread_offset + j) % numCols; local_colStats[j] = col_idx >= numCols ? 0.0f : __ldg(&colStats[col_idx]); local_rowStats[j] = row_idx >= numRows ? 0.0f : __ldg(&rowStats[row_idx]); local_biasValue[j] = ((bias == nullptr) || col_idx >= numCols) ? 0.0f : __half2float(bias[col_idx]); } // Each block loads THREADS * ITEMS_PER_THREAD values from A int valid_items = block_offset + THREADS * ITEMS_PER_THREAD < n_out ? THREADS * ITEMS_PER_THREAD : n_out - block_offset; LoadInt32(loadint32).Load(&(A[block_offset]), local_values, valid_items, 0); #pragma unroll ITEMS_PER_THREAD for (int j = 0; j < ITEMS_PER_THREAD; ++j) { local_output[j] = __float2half( fmaf(local_values[j] * local_rowStats[j] * local_colStats[j], MM_DEQUANT_CONST, local_biasValue[j]) ); } #pragma unroll ITEMS_PER_THREAD for (int j = 0; j < ITEMS_PER_THREAD; j++) { int outIdx = block_offset + thread_offset + j; if (outIdx < n_out) { out[outIdx] = local_output[j]; } } } #define DENORM 1.0f / 127.0f #define MAX_SPARSE_COUNT 32 #define SMEM_SIZE 8 * 256 template __global__ void kspmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, T* B, half* out, float* __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB ) { // 0. load balancing: We process rows with most columns first (count_vec)and we process one row per block // If a block finishes, the next one is scheduled. Since the last blocks like have fewer // elements they finish faster "fillin up" the gaps left by larger blocks // without tensor cores // 1. use rowidx_length to find what to load (as many blocks as there are rows) // 2. Load A into registers // 3. each warp loads all required rows of B but each warp is offset by k // 4. Do mma operations that accumulate into registers // 5. Each warp stores its output row into matrix C const int count = max_count[blockIdx.x]; const int local_max_idx = max_idx[blockIdx.x]; const int offset = local_max_idx == 0 ? 0 : offset_rowidx[local_max_idx - 1]; const int local_row_idx = rowidx[offset]; const int warp_id = threadIdx.x / 32; const int warp_idx = threadIdx.x % 32; const int warp_offset = (warp_id * 32) * SPMM_ITEMS; const int num_items = BITS == 8 ? 8 : 8; int idx_col_B = warp_offset; int local_idx_col_B_offset = 0; half local_valA[MAX_SPARSE_COUNT]; int local_colidxA[MAX_SPARSE_COUNT]; half local_valC[SPMM_ITEMS]; T local_valsB[num_items]; half local_valOut[num_items]; // 128 byte loads per warp == 4 bytes per thread // 2. Load A into registers for (int j = 0; j < MAX_SPARSE_COUNT; j++) { local_valA[j] = j < count ? values[offset + j] : __float2half(0.0f); local_colidxA[j] = j < count ? colidx[offset + j] : 0; } // each thread processes SPMM_ITEMS=32 per iteration. We have 256 threads. 32*256=x192 // we expect each warp to be SPMM_ITEMS*32 apart // we have a total of 128 bytes for the bank with a bank size of 4 bytes // added 3 bytes = 6 values between warps should reduce bank conflicts __shared__ half smem_dequant_stats[SMEM_SIZE]; while (idx_col_B < colsB) { if (dequant_stats != NULL) { for (int i = threadIdx.x; i < SMEM_SIZE; i += blockDim.x) if ((idx_col_B + i - local_idx_col_B_offset) < colsB) smem_dequant_stats[i] = dequant_stats[idx_col_B + i - local_idx_col_B_offset]; __syncthreads(); } #pragma unroll SPMM_ITEMS for (int j = 0; j < SPMM_ITEMS; j++) local_valC[j] = 0.0f; #pragma unroll for (int i = 0; i < count; i++) { // 3. each warp loads all required rows of B but each warp is offset by k int row_offset = colsB * local_colidxA[i]; #pragma unroll SPMM_ITEMS for (int j = 0; j < SPMM_ITEMS; j += num_items) { // 4. Multiply the tile -> accumulate outputs in shared memory until 128 bytes it reached int idx = idx_col_B + (warp_idx * SPMM_ITEMS) + j; if (idx >= colsB) { break; } if ((idx + num_items < colsB)) { if (BITS == 8) reinterpret_cast(local_valsB)[0] = reinterpret_cast(B)[(row_offset + idx) / num_items]; else reinterpret_cast(local_valsB)[0] = reinterpret_cast(B)[(row_offset + idx) / num_items]; } else { #pragma unroll num_items for (int k = 0; k < num_items; k++) if (idx + k < colsB) local_valsB[k] = B[row_offset + idx + k]; else local_valsB[k] = 0.0f; } #pragma unroll num_items for (int k = 0; k < num_items; k++) { if (BITS == 8 && dequant_stats != NULL) // we do texture cache reads (__ldg) on dequant_stats which should be super fast { float valB = local_valsB[k]; float valA = local_valA[i]; if (valB != 0.0 && valA != 0.0) local_valC[j + k] = (float)local_valC[j + k] + ((float)smem_dequant_stats[idx + k - local_idx_col_B_offset]) * DENORM * valB * valA; } else local_valC[j + k] = (float)local_valC[j + k] + (float)local_valsB[k] * (float)local_valA[i]; } } } int idx_row_C = (colsB * local_row_idx); #pragma unroll SPMM_ITEMS for (int j = 0; j < SPMM_ITEMS; j += num_items) { // int idx_col_C = idx_col_B + (32*j) + warp_idx; int idx_col_C = idx_col_B + warp_idx * SPMM_ITEMS + j; int idx_val = idx_col_C + idx_row_C; if (idx_col_C + num_items < colsB) { // load outputs to do inplace addition reinterpret_cast(local_valOut)[0] = reinterpret_cast(out)[idx_val / num_items]; #pragma unroll num_items for (int k = 0; k < num_items; k++) local_valC[(j / num_items) + k] = (float)local_valC[(j / num_items) + k] + (float)local_valOut[k]; reinterpret_cast(out)[idx_val / num_items] = reinterpret_cast(local_valC)[j / num_items]; } else { #pragma unroll num_items for (int k = 0; k < num_items; k++) if (idx_col_C + k < colsB) out[idx_val + k] = (float)out[idx_val + k] + (float)local_valC[j + k]; } } idx_col_B += blockDim.x * SPMM_ITEMS; local_idx_col_B_offset += blockDim.x * SPMM_ITEMS; } } #define WARPS 3 template __global__ void gemm_device(int M, int N, int K, T* __restrict__ const A, T* B, T* out, int lda, int ldb, int ldc) { #if __CUDA_ARCH__ >= 750 using namespace nvcuda; int col_offset = blockIdx.x * 32; const int warp_id = threadIdx.x / 32; const int half_warp_id = threadIdx.x / 16; const int half_warp_lane = threadIdx.x % 16; const int batch_size_warps = (WARPS - 1) * 2; const int val_per_iter = blockDim.x - 32; T local_A[4]; T local_B[128]; const int a_tile_offset = 16; const int b_tile_offset = (16 * 32 + 16); __shared__ T smem_A[8 * 16 + (2 * 16 * (batch_size_warps - 1))]; __shared__ T smem_B[2 * batch_size_warps * 16 * 32 + (2 * 16 * (batch_size_warps - 1))]; //__shared__ T smem_C[8*32]; wmma::fragment a_frag; wmma::fragment b_frag; wmma::fragment c_frag; wmma::fill_fragment(c_frag, 0.0f); int ticktock = 0; int idx = 0 + threadIdx.x; int loaded_values = 0; // prefetch if (idx < K && warp_id < (WARPS - 1)) { if (loaded_values == 0) { local_A[0] = A[idx]; local_A[1] = A[idx + (1 * val_per_iter)]; local_A[2] = A[idx + (2 * val_per_iter)]; local_A[3] = A[idx + (3 * val_per_iter)]; #pragma unroll 32 for (int col = 0; col < 32; col++) { local_B[col] = B[(col_offset + col) * ldb + idx]; local_B[col + 32] = B[(col_offset + col) * ldb + idx + (1 * val_per_iter)]; local_B[col + 64] = B[(col_offset + col) * ldb + idx + (2 * val_per_iter)]; local_B[col + 96] = B[(col_offset + col) * ldb + idx + (3 * val_per_iter)]; } loaded_values = 3; } else { if (loaded_values == 3) { local_A[0] = local_A[1]; #pragma unroll 32 for (int col = 0; col < 32; col++) local_B[col] = local_B[col + (32)]; } else if (loaded_values == 2) { local_A[0] = local_A[2]; #pragma unroll 32 for (int col = 0; col < 32; col++) local_B[col] = local_B[col + (64)]; } else { local_A[0] = local_A[3]; #pragma unroll 32 for (int col = 0; col < 32; col++) local_B[col] = local_B[col + (96)]; } loaded_values--; } smem_A[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * a_tile_offset)] = local_A[0]; #pragma unroll 32 for (int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * b_tile_offset) + (col * 16)] = local_B[col]; } else if (warp_id < (WARPS - 1)) { local_A[0] = T(0.0); smem_A[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * a_tile_offset)] = 0.0f; #pragma unroll 32 for (int col = 0; col < 32; col++) local_B[col] = 0.0f; #pragma unroll 32 for (int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * b_tile_offset) + (col * 16)] = 0.0f; } ticktock = ticktock == 0 ? 1 : 0; // for(int base_idx = blockDim.x-32; base_idx < K; base_idx+=blockDim.x-32) for (int base_idx = blockDim.x - 32; base_idx < K; base_idx += blockDim.x - 32) { idx = base_idx + threadIdx.x; __syncthreads(); if (idx < K && warp_id < (WARPS - 1)) { // local_A[0] = A[idx]; // #pragma unroll 32 // for(int col = 0; col < 32; col++) // local_B[col] = B[(col_offset+col)*ldb+idx]; if (loaded_values == 0) { local_A[0] = A[idx]; local_A[1] = A[idx + (1 * val_per_iter)]; local_A[2] = A[idx + (2 * val_per_iter)]; local_A[3] = A[idx + (3 * val_per_iter)]; #pragma unroll 32 for (int col = 0; col < 32; col++) { local_B[col] = B[(col_offset + col) * ldb + idx]; local_B[col + 32] = B[(col_offset + col) * ldb + idx + (1 * val_per_iter)]; local_B[col + 64] = B[(col_offset + col) * ldb + idx + (2 * val_per_iter)]; local_B[col + 96] = B[(col_offset + col) * ldb + idx + (3 * val_per_iter)]; } loaded_values = 3; } else { if (loaded_values == 3) { local_A[0] = local_A[1]; #pragma unroll 32 for (int col = 0; col < 32; col++) local_B[col] = local_B[col + (32)]; } else if (loaded_values == 2) { local_A[0] = local_A[2]; #pragma unroll 32 for (int col = 0; col < 32; col++) local_B[col] = local_B[col + (64)]; } else { local_A[0] = local_A[3]; #pragma unroll 32 for (int col = 0; col < 32; col++) local_B[col] = local_B[col + (96)]; } loaded_values--; } smem_A[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * a_tile_offset)] = local_A[0]; #pragma unroll 32 for (int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * b_tile_offset) + (col * 16)] = local_B[col]; } else if (warp_id < (WARPS - 1)) { local_A[0] = T(0.0); smem_A[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * a_tile_offset)] = 0.0f; #pragma unroll 32 for (int col = 0; col < 32; col++) local_B[col] = 0.0f; #pragma unroll 32 for (int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * b_tile_offset) + (col * 16)] = 0.0f; } ticktock = ticktock == 0 ? 1 : 0; if (warp_id == (WARPS - 1)) for (int k = 0; k < batch_size_warps; k++) { wmma::load_matrix_sync( a_frag, &(smem_A[(ticktock * batch_size_warps + k) * a_tile_offset]), 16 ); // 111 mu wmma::load_matrix_sync( b_frag, &(smem_B[(ticktock * batch_size_warps + k) * b_tile_offset]), 16 ); // 35 mu wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } } __syncthreads(); if (warp_id != (WARPS - 1)) { return; } // only warp_id == (WARPS-1) from here int warp_lane = threadIdx.x % 32; ticktock = ticktock == 0 ? 1 : 0; for (int k = 0; k < batch_size_warps; k++) { wmma::load_matrix_sync(a_frag, &(smem_A[(ticktock * batch_size_warps + k) * a_tile_offset]), 16); // 111 mu wmma::load_matrix_sync(b_frag, &(smem_B[(ticktock * batch_size_warps + k) * b_tile_offset]), 16); // 35 mu wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } // 129 mu if (warp_id == (WARPS - 1)) wmma::store_matrix_sync(&(smem_A[0]), c_frag, 32, wmma::mem_row_major); if (col_offset + warp_lane < M) out[col_offset + warp_lane] = smem_A[warp_lane]; #endif } template __device__ void printnonzero(T* A, int num_values, const char* strval) { for (int i = 0; i < num_values; i++) if ((float)A[i] != 0.0) printf("%s %i %f\n", strval, i, (float)A[i]); } template __global__ void kgemm_4bit_inference( int M, int N, int K, T* __restrict__ const A, unsigned char* B, float* absmax, T* out, int lda, int ldb, int ldc, int blocksize ) { //// element-wise kernel //// 1. Load batch x k into registers //// 2. Load k x k into registers //// 3. dequantize and store in second pair of k x k //// 4. matmul //// 5. sum with cub //// 6. store outputs //// TC kernel //// use k warps per thread block //// 1. threadblock use read-only cache to read in register tile for A into shared memory //// 2. each warp loops over shared memory tiles of A of size 8x16 and loads them into fragments //// 3. each warp reads a segment of values 16x32 from B //// 4. do dequantization from register of B into second pair of registers //// 5. store (4) into fragment //// 6. matmul aggregate into fragment C //// 7. aggregate files of C into shared memory block C //// 8. sum (7) //// 9. write outputs to matmul output matrix #if __CUDA_ARCH__ >= 750 using namespace nvcuda; int col_offset = blockIdx.x * 32; const int warp_id = threadIdx.x / 32; const int warp_idx = threadIdx.x % 32; const int half_warp_id = threadIdx.x / 16; const int half_warp_lane = threadIdx.x % 16; const int batch_size_warps = (WARPS - 1) * 2; T quant_map[16]; #pragma unroll 16 for (int i = 0; i < 16; i++) quant_map[i] = nf4_dequantization_lut[i]; //__shared__ T quant_map[16*160]; T local_A[2]; T local_B[64]; unsigned char local_B_4bit[32]; const int a_tile_offset = 16; const int b_tile_offset = (16 * 32 + 16); __shared__ T smem_A[8 * 16 + (16 * (batch_size_warps - 1))]; __shared__ T smem_B[2 * batch_size_warps * 16 * 32 + (2 * 16 * (batch_size_warps - 1))]; __shared__ T smem_C[8 * 32]; wmma::fragment a_frag; wmma::fragment b_frag; wmma::fragment c_frag; wmma::fill_fragment(c_frag, 0.0f); for (int i = threadIdx.x; i < (8 * 32); i += blockDim.x) smem_C[i] = 0.0f; __syncthreads(); int ticktock = 0; int idx = 0 + threadIdx.x; int loaded_values = 0; // prefetch if (idx < K && warp_id < (WARPS - 1)) { if (loaded_values == 0) { local_A[0] = A[idx]; local_A[1] = A[idx + blockDim.x - 32]; #pragma unroll 32 for (int col = 0; col < 32; col++) local_B_4bit[col] = B[(col_offset + col) * ldb + idx]; loaded_values = 1; } else { local_A[0] = local_A[1]; loaded_values--; #pragma unroll 64 for (int col = 0; col < 64; col += 2) { // local_B[col] = dhDequantizeNF4(local_B_4bit[col/2] >> 4)*T(1.0f); // local_B[col+1] = dhDequantizeNF4(local_B_4bit[col/2] & 0x0F)*T(1.0f); // local_B[col] = d2DequantizeFP4(local_B_4bit[col/2] >> 4)*(float)(17.0); // local_B[col+1] = d2DequantizeFP4(local_B_4bit[col/2] & 0x0F)*(float)(17.0); // local_B[col] = 127*(local_B_4bit[col/2] >> 4)*(float)(17.0); // local_B[col+1] = 127*(local_B_4bit[col/2] & 0x0F)*(float)(17.0); // local_B[col] = quant_map[(local_B_4bit[col/2] >> 4)]*T(17.0); // local_B[col+1] = quant_map[(local_B_4bit[col/2] & 0x0F)]*T(17.0); local_B[col] = quant_map[160 * (local_B_4bit[col / 2] >> 4) + warp_idx] * T(17.0); local_B[col + 1] = quant_map[160 * (local_B_4bit[col / 2] & 0x0F) + warp_idx] * T(17.0); } } smem_A[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * a_tile_offset)] = local_A[0]; #pragma unroll 32 for (int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * b_tile_offset) + (col * 16)] = local_B[col]; } else if (warp_id < (WARPS - 1)) { local_A[0] = T(0.0); smem_A[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * a_tile_offset)] = 0.0f; #pragma unroll 32 for (int col = 0; col < 32; col++) local_B[col] = 0.0f; #pragma unroll 32 for (int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * b_tile_offset) + (col * 16)] = 0.0f; } ticktock = ticktock == 0 ? 1 : 0; // if(threadIdx.x == 0) // printf("aa %i %i\n", idx, loaded_values); // for(int base_idx = blockDim.x-32; base_idx < K; base_idx+=blockDim.x-32) for (int base_idx = blockDim.x - 32; base_idx < K; base_idx += blockDim.x - 32) { idx = base_idx + threadIdx.x; // if(threadIdx.x == 0) // printf("%i %i\n", idx, loaded_values); //__syncthreads(); if (idx < K && warp_id < (WARPS - 1)) { if (loaded_values == 0) { local_A[0] = A[idx]; local_A[1] = A[idx + blockDim.x - 32]; #pragma unroll 32 for (int col = 0; col < 32; col++) { local_B_4bit[col] = B[(col_offset + col) * ldb + idx]; local_B_4bit[col + 16] = B[(col_offset + col) * ldb + idx]; } loaded_values = 1; } else { local_A[0] = local_A[1]; loaded_values--; int absidx = (idx + col_offset) / blocksize; half local_absmax = __ldg(&(absmax[absidx])); #pragma unroll 64 for (int col = 0; col < 64; col += 2) { // local_B[col] = dhDequantizeNF4(local_B_4bit[col/2] >> 4)*T(absidx); // local_B[col+1] = dhDequantizeNF4(local_B_4bit[col/2] & 0x0F)*T(absidx); // local_B[col] = T(127)*T(local_B_4bit[col/2] >> 4)*T(absidx); // local_B[col+1] = T(127)*T(local_B_4bit[col/2] & 0x0F)*T(absidx); // local_B[col] = quant_map[160*(local_B_4bit[col/2] >> 4)+warp_idx]*T(local_absmax); // local_B[col+1] = quant_map[160*(local_B_4bit[col/2] & 0x0F)+warp_idx]*T(local_absmax); local_B[col] = quant_map[(local_B_4bit[col / 2] >> 4)] * T(absidx); local_B[col + 1] = quant_map[(local_B_4bit[col / 2] & 0x0F)] * T(absidx); } // printnonzero(local_B, 128, ""); } smem_A[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * a_tile_offset)] = local_A[0]; #pragma unroll 32 for (int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * b_tile_offset) + (col * 16)] = local_B[col]; } else if (warp_id < (WARPS - 1)) { local_A[0] = T(0.0); smem_A[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * a_tile_offset)] = 0.0f; #pragma unroll 32 for (int col = 0; col < 32; col++) local_B[col] = 0.0f; #pragma unroll 32 for (int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps * ticktock) + half_warp_id) * b_tile_offset) + (col * 16)] = 0.0f; } ticktock = ticktock == 0 ? 1 : 0; if (warp_id == (WARPS - 1)) for (int k = 0; k < batch_size_warps; k++) { wmma::load_matrix_sync( a_frag, &(smem_A[(ticktock * batch_size_warps + k) * a_tile_offset]), 16 ); // 111 mu wmma::load_matrix_sync( b_frag, &(smem_B[(ticktock * batch_size_warps + k) * b_tile_offset]), 16 ); // 35 mu wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } } __syncthreads(); // if(threadIdx.x == 0) //{ // printnonzero(smem_A, 8*16 + (2*16*(batch_size_warps-1)), "A: "); // printnonzero(smem_B, 2*batch_size_warps*16*32 + (2*16*(batch_size_warps-1)), "B: "); // } if (warp_id != (WARPS - 1)) { return; } // only warp_id == (WARPS-1) from here int warp_lane = threadIdx.x % 32; ticktock = ticktock == 0 ? 1 : 0; for (int k = 0; k < batch_size_warps; k++) { // if(warp_lane == 0) // printf("%i %i %i %i\n", (ticktock*batch_size_warps + k)*a_tile_offset, k, ticktock, threadIdx.x); wmma::load_matrix_sync(a_frag, &(smem_A[(ticktock * batch_size_warps + k) * a_tile_offset]), 16); // 111 mu wmma::load_matrix_sync(b_frag, &(smem_B[(ticktock * batch_size_warps + k) * b_tile_offset]), 16); // 35 mu wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } // 129 mu if (warp_id == (WARPS - 1)) wmma::store_matrix_sync(&(smem_C[0]), c_frag, 32, wmma::mem_row_major); // printnonzero(smem_C, 32, ""); if (col_offset + warp_lane < M) out[col_offset + warp_lane] = smem_C[warp_lane]; #endif } #define num_values_4bit 32 template __global__ void kgemm_4bit_inference_naive( int M, int N, int K, T* __restrict__ const A, unsigned char* B, float* absmax, const float* datatype, T* out, int lda, int ldb, int ldc, int blocksize ) { // per threadblock: // load step-by-step in chunks of [32,warps]: 1x32 * [32,warps] -> [1,warps] // 4 warps -> 4 loads per iter // 1x32 * 32x4 -> 1x4 outputs per thread block typedef cub::WarpReduce WarpReduce; __shared__ typename WarpReduce::TempStorage temp_storage[THREADS / 32]; const int warp_idx = threadIdx.x / 32; const int warp_lane = threadIdx.x % 32; const int row_B = (THREADS / 32) * blockIdx.x + warp_idx; const int offset_B = ldb * row_B; const int num_values_8bit = num_values_4bit / 2; float local_C = 0.0f; unsigned char local_B_4bit[num_values_8bit]; T local_B[num_values_4bit / 4]; T local_A[num_values_4bit / 4]; __shared__ T quant_map[16]; T local_absmax = T(0.0f); if (threadIdx.x < 16) quant_map[threadIdx.x] = T(__ldg(&datatype[threadIdx.x])); // for(int i = threadIdx.x; i < 16; i++) // quant_map[i] = T(__ldg(&datatype[i])); __syncthreads(); // A: [1, K] // B: [N, K] for (int inner_idx = warp_lane * num_values_4bit; inner_idx < K; inner_idx += 32 * num_values_4bit) { const int inner_idx_halved = inner_idx / 2; // Since blocksize will always be a power-of-2, we avoid more expensive // division by the blocksize and instead use a shift operation. // This is equivalent to (i+threadId.x*NUM_PER_TH)/blocksize. const int absidx = ((2 * offset_B) + inner_idx) >> (31 - __clz(blocksize)); local_absmax = __ldg(&(absmax[absidx])); if (row_B < M) { if ((inner_idx_halved + num_values_8bit) < (K / 2)) { // this is the most important for performance considerations reinterpret_cast(local_B_4bit)[0] = reinterpret_cast(B)[(offset_B + (inner_idx_halved)) / (num_values_8bit)]; } else { #pragma unroll for (int j = 0; j < (num_values_8bit); j++) if ((inner_idx_halved) + j < (K / 2)) local_B_4bit[j] = B[offset_B + inner_idx_halved + j]; else local_B_4bit[j] = 0b01110111; } } else { #pragma unroll for (int j = 0; j < (num_values_8bit); j++) local_B_4bit[j] = 0b01110111; } for (int i = 0; i < 4; i++) { #pragma unroll for (int k = 0; k < num_values_8bit / 4; k++) { #if BNB_BF16_AVAILABLE local_B[k * 2] = quant_map[local_B_4bit[(i * num_values_8bit / 4) + k] >> 4] * local_absmax; local_B[k * 2 + 1] = quant_map[local_B_4bit[(i * num_values_8bit / 4) + k] & 0x0F] * local_absmax; #else // bf16 multipliation not supported local_B[k * 2] = T((float)quant_map[local_B_4bit[(i * num_values_8bit / 4) + k] >> 4] * (float)local_absmax); local_B[k * 2 + 1] = T((float)quant_map[local_B_4bit[(i * num_values_8bit / 4) + k] & 0x0F] * (float)local_absmax); #endif } if (inner_idx + (num_values_4bit / 4) + (i * num_values_4bit / 4) < K) { // this is also relatively important for performance if (BITS == 16) { reinterpret_cast(local_A)[0] = reinterpret_cast(A)[inner_idx / (num_values_4bit / 4) + i]; } else { reinterpret_cast(local_A)[0] = reinterpret_cast(A)[inner_idx / (num_values_4bit / 8) + (2 * i) + 0]; reinterpret_cast(local_A)[1] = reinterpret_cast(A)[inner_idx / (num_values_4bit / 8) + (2 * i) + 1]; } } else #pragma unroll for (int k = 0; k < num_values_4bit / 4; k++) if (inner_idx + (i * num_values_4bit / 4) + k < K) local_A[k] = A[inner_idx + k + (i * num_values_4bit / 4)]; else local_A[k] = T(0.0f); // accumulate in float; small performance hit for Ampere, but lower error for outputs #pragma unroll for (int k = 0; k < num_values_4bit / 4; k++) { #if BNB_BF16_AVAILABLE local_C += (float)(local_A[k] * local_B[k]); #else // bf16 multipliation not supported local_C += ((float)local_A[k] * (float)local_B[k]); #endif } } } local_C = WarpReduce(temp_storage[warp_idx]).Sum(local_C); if (row_B < M && warp_lane == 0) out[row_B] = T(local_C); } template __global__ void kfunc(T* A, T* B, T value, long n) { for (long i = (blockDim.x * blockIdx.x) + threadIdx.x; i < n; i += (blockDim.x * gridDim.x)) { switch (FUNC) { case FILL: A[i] = (T)value; break; case ARANGE: A[i] = (T)i; break; case _MUL: A[i] = A[i] * B[i]; break; } } } //============================================================== // TEMPLATE DEFINITIONS //============================================================== template __global__ void kfunc(float* A, float* B, float value, long n); template __global__ void kfunc(unsigned char* A, unsigned char* B, unsigned char value, long n); template __global__ void kfunc(float* A, float* B, float value, long n); template __global__ void kfunc(float* A, float* B, float value, long n); // these are not used and make no sense, but the compiler needs them // template __global__ void gemm_device(int M, int N, int K, float * __restrict__ const A, float* B, // float * out, int lda, int ldb, int ldc); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); // template __global__ void gemm_device(int M, int N, int K, float * __restrict__ const A, float* B, // float * out, int lda, int ldb, int ldc); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); // these are not used and make no sense, but the compiler needs them // template __global__ void gemm_device(int M, int N, int K, float * __restrict__ const A, float* B, // float * out, int lda, int ldb, int ldc); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); // template __global__ void gemm_device(int M, int N, int K, float * __restrict__ const A, float* B, // float * out, int lda, int ldb, int ldc); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); template __global__ void gemm_device( int M, int N, int K, half* __restrict__ const A, half* B, half* out, int lda, int ldb, int ldc ); template __global__ void kgemm_4bit_inference( int M, int N, int K, half* __restrict__ const A, unsigned char* B, float* absmax, half* out, int lda, int ldb, int ldc, int blocksize ); template __global__ void kgemm_4bit_inference( int M, int N, int K, half* __restrict__ const A, unsigned char* B, float* absmax, half* out, int lda, int ldb, int ldc, int blocksize ); template __global__ void kgemm_4bit_inference( int M, int N, int K, half* __restrict__ const A, unsigned char* B, float* absmax, half* out, int lda, int ldb, int ldc, int blocksize ); template __global__ void kgemm_4bit_inference( int M, int N, int K, half* __restrict__ const A, unsigned char* B, float* absmax, half* out, int lda, int ldb, int ldc, int blocksize ); template __global__ void kgemm_4bit_inference_naive( int M, int N, int K, half* __restrict__ const A, unsigned char* B, float* absmax, const float* datatype, half* out, int lda, int ldb, int ldc, int blocksize ); template __global__ void kgemm_4bit_inference_naive<__nv_bfloat16, 128, 16>( int M, int N, int K, __nv_bfloat16* __restrict__ const A, unsigned char* B, float* absmax, const float* datatype, __nv_bfloat16* out, int lda, int ldb, int ldc, int blocksize ); template __global__ void kgemm_4bit_inference_naive( int M, int N, int K, float* __restrict__ const A, unsigned char* B, float* absmax, const float* datatype, float* out, int lda, int ldb, int ldc, int blocksize ); template __global__ void kspmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, half* B, half* out, float* __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB ); template __global__ void kspmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, half* B, half* out, float* __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB ); template __global__ void kspmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, half* B, half* out, float* __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB ); template __global__ void kspmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, signed char* B, half* out, float* __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB ); template __global__ void kspmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, signed char* B, half* out, float* __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB ); template __global__ void kspmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, signed char* B, half* out, float* __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB ); template __global__ void kdequant_mm_int32_fp16<4, 512>( int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, half* out, half* __restrict__ const bias, const int numRows, const int numCols, const int n ); template __device__ unsigned char dQuantize<0>(float* smem_code, const float rand, float x); template __device__ unsigned char dQuantize<1>(float* smem_code, const float rand, float x); #define MAKE_PreconditionOptimizer32bit1State(oname, gtype) \ template __global__ void kPreconditionOptimizer32bit1State( \ gtype * g, gtype * p, float* state1, float* unorm, const float beta1, const float beta2, const float eps, \ const float weight_decay, const int step, const float lr, const float gnorm_scale, const int n \ ); MAKE_PreconditionOptimizer32bit1State(MOMENTUM, half) MAKE_PreconditionOptimizer32bit1State(MOMENTUM, float) MAKE_PreconditionOptimizer32bit1State(MOMENTUM, __nv_bfloat16) MAKE_PreconditionOptimizer32bit1State(RMSPROP, half) MAKE_PreconditionOptimizer32bit1State(RMSPROP, float) MAKE_PreconditionOptimizer32bit1State(RMSPROP, __nv_bfloat16) MAKE_PreconditionOptimizer32bit1State(LION, half) MAKE_PreconditionOptimizer32bit1State(LION, float) MAKE_PreconditionOptimizer32bit1State(LION, __nv_bfloat16) MAKE_PreconditionOptimizer32bit1State(ADAGRAD, half) MAKE_PreconditionOptimizer32bit1State(ADAGRAD, float) MAKE_PreconditionOptimizer32bit1State(ADAGRAD, __nv_bfloat16) #define MAKE_Optimizer32bit1State(oname, gtype) \ template __global__ void kOptimizer32bit1State( \ gtype * g, gtype * p, float* state1, float* unorm, const float max_unorm, const float param_norm, \ const float beta1, const float beta2, const float eps, const float weight_decay, const int step, \ const float lr, const float gnorm_scale, const bool skip_zeros, const int n \ ); MAKE_Optimizer32bit1State(MOMENTUM, half) MAKE_Optimizer32bit1State(MOMENTUM, float) MAKE_Optimizer32bit1State(MOMENTUM, __nv_bfloat16) MAKE_Optimizer32bit1State(RMSPROP, half) MAKE_Optimizer32bit1State(RMSPROP, float) MAKE_Optimizer32bit1State(RMSPROP, __nv_bfloat16) MAKE_Optimizer32bit1State(LION, half) MAKE_Optimizer32bit1State(LION, float) MAKE_Optimizer32bit1State(LION, __nv_bfloat16) MAKE_Optimizer32bit1State(ADAGRAD, half) MAKE_Optimizer32bit1State(ADAGRAD, float) MAKE_Optimizer32bit1State(ADAGRAD, __nv_bfloat16) #define MAKE_PreconditionOptimizer32bit2State(oname, gtype) \ template __global__ void kPreconditionOptimizer32bit2State( \ gtype * g, gtype * p, float* state1, float* state2, float* unorm, const float beta1, const float beta2, \ const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, \ const int n \ ); MAKE_PreconditionOptimizer32bit2State(ADAM, float) MAKE_PreconditionOptimizer32bit2State(ADAM, half) MAKE_PreconditionOptimizer32bit2State(ADAM, __nv_bfloat16) MAKE_PreconditionOptimizer32bit2State(ADEMAMIX, float) MAKE_PreconditionOptimizer32bit2State(ADEMAMIX, half) MAKE_PreconditionOptimizer32bit2State(ADEMAMIX, __nv_bfloat16) template __global__ void kOptimizer32bit2State( float* g, float* p, float* state1, float* state2, float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n ); template __global__ void kOptimizer32bit2State( half* g, half* p, float* state1, float* state2, float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n ); template __global__ void kOptimizer32bit2State<__nv_bfloat16, ADAM>( __nv_bfloat16* g, __nv_bfloat16* p, float* state1, float* state2, float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n ); template __global__ void kOptimizer32bit2State( float* g, float* p, float* state1, float* state2, float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n ); template __global__ void kOptimizer32bit2State( half* g, half* p, float* state1, float* state2, float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n ); template __global__ void kOptimizer32bit2State<__nv_bfloat16, ADEMAMIX>( __nv_bfloat16* g, __nv_bfloat16* p, float* state1, float* state2, float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n ); #define MAKE_PreconditionStatic8bit1State(oname, gtype) \ template __global__ void kPreconditionOptimizerStatic8bit1State( \ gtype * p, gtype* __restrict__ const g, unsigned char* __restrict__ const state1, float* unorm, \ const float beta1, const float beta2, const float eps, const int step, float* __restrict__ const quantiles1, \ float* max1, float* new_max1, const float weight_decay, const float gnorm_scale, const int n \ ); MAKE_PreconditionStatic8bit1State(MOMENTUM, half) MAKE_PreconditionStatic8bit1State(MOMENTUM, float) MAKE_PreconditionStatic8bit1State(RMSPROP, half) MAKE_PreconditionStatic8bit1State(RMSPROP, float) MAKE_PreconditionStatic8bit1State(LION, half) MAKE_PreconditionStatic8bit1State(LION, float) MAKE_PreconditionStatic8bit1State(ADAGRAD, half) MAKE_PreconditionStatic8bit1State(ADAGRAD, float) #define MAKE_optimizerStatic8bit1State(oname, gtype) \ template __global__ void kOptimizerStatic8bit1State( \ gtype * p, gtype* const g, unsigned char* state1, const float* unorm, const float max_unorm, \ const float param_norm, const float beta1, const float beta2, const float eps, const int step, const float lr, \ float* __restrict__ const quantiles1, float* max1, float* new_max1, float weight_decay, \ const float gnorm_scale, const int n \ ); MAKE_optimizerStatic8bit1State(MOMENTUM, half) MAKE_optimizerStatic8bit1State(MOMENTUM, float) MAKE_optimizerStatic8bit1State(RMSPROP, half) MAKE_optimizerStatic8bit1State(RMSPROP, float) MAKE_optimizerStatic8bit1State(LION, half) MAKE_optimizerStatic8bit1State(LION, float) MAKE_optimizerStatic8bit1State(ADAGRAD, half) MAKE_optimizerStatic8bit1State(ADAGRAD, float) #define MAKE_PreconditionStatic8bit2State(oname, gtype) \ template __global__ void kPreconditionOptimizerStatic8bit2State( \ gtype * p, gtype* __restrict__ const g, unsigned char* __restrict__ const state1, \ unsigned char* __restrict__ const state2, float* unorm, const float beta1, const float beta2, const float eps, \ const int step, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* max1, \ float* max2, float* new_max1, float* new_max2, const float gnorm_scale, const int n \ ); MAKE_PreconditionStatic8bit2State(ADAM, half) MAKE_PreconditionStatic8bit2State(ADAM, float) #define MAKE_optimizerStatic8bit2State(oname, gtype) \ template __global__ void kOptimizerStatic8bit2State( \ gtype * p, gtype* const g, unsigned char* state1, unsigned char* state2, const float* unorm, \ const float max_unorm, const float param_norm, const float beta1, const float beta2, const float eps, \ const int step, const float lr, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, \ float* max1, float* max2, float* new_max1, float* new_max2, float weight_decay, const float gnorm_scale, \ const int n \ ); MAKE_optimizerStatic8bit2State(ADAM, half) MAKE_optimizerStatic8bit2State(ADAM, float) template __global__ void kPercentileClipping(float* __restrict__ g, float* gnorm_vec, int step, const int n); template __global__ void kPercentileClipping(half* __restrict__ g, float* gnorm_vec, int step, const int n); #define MAKE_kQuantizeBlockwise(dtype, blocksize, num_per_thread, stochastic, data_type_name) \ template __global__ void kQuantizeBlockwise( \ float* code, dtype* __restrict__ const A, float* absmax, unsigned char* out, float* __restrict__ const rand, \ const int rand_offset, const int n \ ); MAKE_kQuantizeBlockwise(half, 4096, 4, 0, General8bit) MAKE_kQuantizeBlockwise(half, 4096, 4, 1, General8bit) MAKE_kQuantizeBlockwise(half, 2048, 4, 0, General8bit) MAKE_kQuantizeBlockwise(half, 1024, 4, 0, General8bit) MAKE_kQuantizeBlockwise(half, 512, 2, 0, General8bit) MAKE_kQuantizeBlockwise(half, 256, 2, 0, General8bit) MAKE_kQuantizeBlockwise(half, 128, 2, 0, General8bit) MAKE_kQuantizeBlockwise(half, 64, 2, 0, General8bit) MAKE_kQuantizeBlockwise(half, 4096, 4, 0, FP4) MAKE_kQuantizeBlockwise(half, 2048, 4, 0, FP4) MAKE_kQuantizeBlockwise(half, 1024, 4, 0, FP4) MAKE_kQuantizeBlockwise(half, 512, 2, 0, FP4) MAKE_kQuantizeBlockwise(half, 256, 2, 0, FP4) MAKE_kQuantizeBlockwise(half, 128, 2, 0, FP4) MAKE_kQuantizeBlockwise(half, 64, 2, 0, FP4) MAKE_kQuantizeBlockwise(half, 4096, 4, 0, NF4) MAKE_kQuantizeBlockwise(half, 2048, 4, 0, NF4) MAKE_kQuantizeBlockwise(half, 1024, 4, 0, NF4) MAKE_kQuantizeBlockwise(half, 512, 2, 0, NF4) MAKE_kQuantizeBlockwise(half, 256, 2, 0, NF4) MAKE_kQuantizeBlockwise(half, 128, 2, 0, NF4) MAKE_kQuantizeBlockwise(half, 64, 2, 0, NF4) MAKE_kQuantizeBlockwise(float, 4096, 4, 0, General8bit) MAKE_kQuantizeBlockwise(float, 4096, 4, 1, General8bit) MAKE_kQuantizeBlockwise(float, 2048, 4, 0, General8bit) MAKE_kQuantizeBlockwise(float, 1024, 4, 0, General8bit) MAKE_kQuantizeBlockwise(float, 512, 2, 0, General8bit) MAKE_kQuantizeBlockwise(float, 256, 2, 0, General8bit) MAKE_kQuantizeBlockwise(float, 128, 2, 0, General8bit) MAKE_kQuantizeBlockwise(float, 64, 2, 0, General8bit) MAKE_kQuantizeBlockwise(float, 4096, 4, 0, FP4) MAKE_kQuantizeBlockwise(float, 2048, 4, 0, FP4) MAKE_kQuantizeBlockwise(float, 1024, 4, 0, FP4) MAKE_kQuantizeBlockwise(float, 512, 2, 0, FP4) MAKE_kQuantizeBlockwise(float, 256, 2, 0, FP4) MAKE_kQuantizeBlockwise(float, 128, 2, 0, FP4) MAKE_kQuantizeBlockwise(float, 64, 2, 0, FP4) MAKE_kQuantizeBlockwise(float, 4096, 4, 0, NF4) MAKE_kQuantizeBlockwise(float, 2048, 4, 0, NF4) MAKE_kQuantizeBlockwise(float, 1024, 4, 0, NF4) MAKE_kQuantizeBlockwise(float, 512, 2, 0, NF4) MAKE_kQuantizeBlockwise(float, 256, 2, 0, NF4) MAKE_kQuantizeBlockwise(float, 128, 2, 0, NF4) MAKE_kQuantizeBlockwise(float, 64, 2, 0, NF4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 4096, 4, 0, General8bit) MAKE_kQuantizeBlockwise(__nv_bfloat16, 4096, 4, 1, General8bit) MAKE_kQuantizeBlockwise(__nv_bfloat16, 2048, 4, 0, General8bit) MAKE_kQuantizeBlockwise(__nv_bfloat16, 1024, 4, 0, General8bit) MAKE_kQuantizeBlockwise(__nv_bfloat16, 512, 2, 0, General8bit) MAKE_kQuantizeBlockwise(__nv_bfloat16, 256, 2, 0, General8bit) MAKE_kQuantizeBlockwise(__nv_bfloat16, 128, 2, 0, General8bit) MAKE_kQuantizeBlockwise(__nv_bfloat16, 64, 2, 0, General8bit) MAKE_kQuantizeBlockwise(__nv_bfloat16, 4096, 4, 0, FP4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 2048, 4, 0, FP4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 1024, 4, 0, FP4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 512, 2, 0, FP4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 256, 2, 0, FP4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 128, 2, 0, FP4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 64, 2, 0, FP4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 4096, 4, 0, NF4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 2048, 4, 0, NF4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 1024, 4, 0, NF4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 512, 2, 0, NF4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 256, 2, 0, NF4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 128, 2, 0, NF4) MAKE_kQuantizeBlockwise(__nv_bfloat16, 64, 2, 0, NF4) template __global__ void kDequantizeBlockwise( float* code, unsigned char* A, float* absmax, half* out, const int blocksize, const int n ); template __global__ void kDequantizeBlockwise( float* code, unsigned char* A, float* absmax, half* out, const int blocksize, const int n ); template __global__ void kDequantizeBlockwise( float* code, unsigned char* A, float* absmax, half* out, const int blocksize, const int n ); template __global__ void kDequantizeBlockwise( float* code, unsigned char* A, float* absmax, float* out, const int blocksize, const int n ); template __global__ void kDequantizeBlockwise( float* code, unsigned char* A, float* absmax, float* out, const int blocksize, const int n ); template __global__ void kDequantizeBlockwise( float* code, unsigned char* A, float* absmax, float* out, const int blocksize, const int n ); template __global__ void kDequantizeBlockwise<__nv_bfloat16, 512, 64, 8, FP4>( float* code, unsigned char* A, float* absmax, __nv_bfloat16* out, const int blocksize, const int n ); template __global__ void kDequantizeBlockwise<__nv_bfloat16, 512, 64, 8, General8bit>( float* code, unsigned char* A, float* absmax, __nv_bfloat16* out, const int blocksize, const int n ); template __global__ void kDequantizeBlockwise<__nv_bfloat16, 512, 64, 8, NF4>( float* code, unsigned char* A, float* absmax, __nv_bfloat16* out, const int blocksize, const int n ); #define MAKE_OptimizerStatic8bit2StateBlockwise(oname, gtype, block_size, num_per_thread) \ template __global__ void kOptimizerStatic8bit2StateBlockwise( \ gtype * p, gtype* __restrict__ const g, unsigned char* state1, unsigned char* state2, const float beta1, \ const float beta2, const float beta3, const float alpha, const float eps, const int step, const float lr, \ float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* absmax1, float* absmax2, \ float weight_decay, const float gnorm_scale, const bool skip_zeros, const int n \ ); MAKE_OptimizerStatic8bit2StateBlockwise(ADAM, float, 256, 1) MAKE_OptimizerStatic8bit2StateBlockwise(ADAM, half, 256, 1) MAKE_OptimizerStatic8bit2StateBlockwise(ADAM, __nv_bfloat16, 256, 1) MAKE_OptimizerStatic8bit2StateBlockwise(ADEMAMIX, float, 256, 1) MAKE_OptimizerStatic8bit2StateBlockwise(ADEMAMIX, half, 256, 1) MAKE_OptimizerStatic8bit2StateBlockwise(ADEMAMIX, __nv_bfloat16, 256, 1) #define MAKE_OptimizerStatic8bit1StateBlockwise(oname, gtype, block_size, num_per_thread) \ template __global__ void kOptimizerStatic8bit1StateBlockwise( \ gtype * p, gtype* __restrict__ const g, unsigned char* state1, const float beta1, const float beta2, \ const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* absmax1, \ float weight_decay, const float gnorm_scale, const bool skip_zeros, const int n \ ); MAKE_OptimizerStatic8bit1StateBlockwise(MOMENTUM, float, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(MOMENTUM, half, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(MOMENTUM, __nv_bfloat16, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(RMSPROP, float, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(RMSPROP, half, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(RMSPROP, __nv_bfloat16, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(LION, float, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(LION, half, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(LION, __nv_bfloat16, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(ADAGRAD, float, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(ADAGRAD, half, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(ADAGRAD, __nv_bfloat16, 256, 1) template __device__ void printnonzero(float* A, int num_values, const char* strval); template __device__ void printnonzero(half* A, int num_values, const char* strval); bitsandbytes/csrc/kernels.cuh000066400000000000000000000157361506752665600167050ustar00rootroot00000000000000// Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include #include #ifndef kernels #define kernels __global__ void kQuantize(float* code, float* __restrict__ const A, unsigned char* out, const int n); __global__ void kDequantize(float* code, unsigned char* A, float* out, const int n); template __global__ void kQuantizeBlockwise( float* code, T* __restrict__ const A, float* absmax, unsigned char* out, float* __restrict__ const rand, const int rand_offset, const int n ); template __global__ void kDequantizeBlockwise(float* code, unsigned char* A, float* absmax, T* out, const int blocksize, const int n); template __global__ void kPreconditionOptimizer32bit2State( T* g, T* p, float* state1, float* state2, float* unorm, const float beta1, const float beta2, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const int n ); template __global__ void kOptimizer32bit2State( T* g, T* p, float* state1, float* state2, float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n ); template __global__ void kPreconditionOptimizer32bit1State( T* g, T* p, float* state1, float* unorm, const float beta1, const float beta2, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const int n ); template __global__ void kOptimizer32bit1State( T* g, T* p, float* state1, float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n ); template __global__ void kPreconditionOptimizerStatic8bit1State( T* p, T* __restrict__ const g, unsigned char* __restrict__ const state1, float* unorm, const float beta1, const float beta2, const float eps, const int step, float* __restrict__ const quantiles1, float* max1, float* new_max1, const float weight_decay, const float gnorm_scale, const int n ); template __global__ void kOptimizerStatic8bit1State( T* p, T* const g, unsigned char* state1, const float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* max1, float* new_max1, float weight_decay, const float gnorm_scale, const int n ); template __global__ void kPreconditionOptimizerStatic8bit2State( T* p, T* __restrict__ const g, unsigned char* __restrict__ const state1, unsigned char* __restrict__ const state2, float* unorm, const float beta1, const float beta2, const float eps, const int step, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, const float gnorm_scale, const int n ); template __global__ void kOptimizerStatic8bit2State( T* p, T* const g, unsigned char* state1, unsigned char* state2, const float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, float weight_decay, const float gnorm_scale, const int n ); template __global__ void kOptimizerStatic8bit2StateBlockwise( T* p, T* __restrict__ const g, unsigned char* state1, unsigned char* state2, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* absmax1, float* absmax2, float weight_decay, const float gnorm_scale, const bool skip_zeros, const int n ); template __global__ void kOptimizerStatic8bit1StateBlockwise( T* p, T* __restrict__ const g, unsigned char* state1, const float beta1, const float beta2, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* absmax1, float weight_decay, const float gnorm_scale, const bool skip_zeros, const int n ); template __global__ void kPercentileClipping(T* __restrict__ g, float* gnorm_vec, int step, const int n); template __global__ void kspmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, T* B, half* out, float* __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB ); template __global__ void kdequant_mm_int32_fp16( int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, half* out, half* __restrict__ const bias, const int numRows, const int numCols, const int n ); template __global__ void kgetRowStats(T* __restrict__ A, float* rowStats, float threshold, int rows, int cols); template __global__ void kInt8VectorQuant(T* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols); template __global__ void kTransformRowToFormat( char* __restrict__ const A, char* out, int rows, int cols, int tiledCols, int outRows, int outCols ); template __global__ void gemm_device(int M, int N, int K, T* __restrict__ const A, T* B, T* out, int lda, int ldb, int ldc); template __global__ void kgemm_4bit_inference( int M, int N, int K, T* __restrict__ const A, unsigned char* B, float* absmax, T* out, int lda, int ldb, int ldc, int blocksize ); template __global__ void kgemm_4bit_inference_naive( int M, int N, int K, T* __restrict__ const A, unsigned char* B, float* absmax, const float* datatype, T* out, int lda, int ldb, int ldc, int blocksize ); template __global__ void kfunc(T* A, T* B, T value, long n); #endif bitsandbytes/csrc/kernels.hip000066400000000000000000003562631506752665600167110ustar00rootroot00000000000000// !!! This is a file automatically generated by hipify!!! #include "hip/hip_runtime.h" // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include "kernels_hip.cuh" #include "common_hip.cuh" #include #include #include //#include #define HLF_MAX 65504 #define TH 1024 #define NUM 4 #define NUM_BLOCK 4096 __device__ static float fp4_dequantization_lut[8] = { 0.0f, // 0b000 0.005208333333f, // 0b001 0.66666667f, // 0b010 1.0f, // 0b011 0.33333333f, // 0b100 0.5f, // 0b101 0.16666667f, // 0b110 0.25f // 0b111 }; __device__ static float nf4_dequantization_lut[16] = { -1.0f, // 0b0000 -0.6961928009986877f, // 0b0001 -0.5250730514526367f, // 0b0010 -0.39491748809814453f, // 0b0011 -0.28444138169288635f, // 0b0100 -0.18477343022823334f, // 0b0101 -0.09105003625154495f, // 0b0110 0.0f, // 0b0111 0.07958029955625534f, // 0b1000 0.16093020141124725f, // 0b1001 0.24611230194568634f, // 0b1010 0.33791524171829224f, // 0b1011 0.44070982933044434f, // 0b1100 0.5626170039176941f, // 0b1101 0.7229568362236023f, // 0b1110 1.0f // 0b1111 }; // source: https://stackoverflow.com/questions/17399119/how-do-i-use-atomicmax-on-floating-point-values-in-cuda // Luckily we have atomicmax and atomicmin in ROCm __device__ __forceinline__ float dDequantizeFP4Tree(unsigned char val) { float sign = 1.0f - 2 * ((val & 0b1000) >> 3); return fp4_dequantization_lut[val & 0b111] * sign; } __device__ unsigned char dQuantizeFP4(float x) { // FP4 with bias of 3 // first bit is a sign // subnormals // 0b000 = 0 // 0b001 = 0.0625 // 0b110 = 2 // 0b111 = 3 // 0b100 = 4 // 0b101 = 6 // 0b010 = 8 // 0b011 = 12 // we do a binary search // the pivots are divided by 12 (the FP4 absmax) // since we assume input data is in [-1.0, 1.0] // !be careful here, its easy to make a mistake // that is difficult to notice if you add an extra // zero somewhere! int sign = x < 0 ? 0b1000 : 0b0000; x = fabsf(x); if(x > 0.29166667f) if( x > 0.583333f) if( x > 0.8333333f) return 0b0011+sign; else return 0b0010+sign; else if(x > 0.4166667f) return 0b101+sign; else return 0b100+sign; else if(x > 0.0859375f) if(x > 0.20833333f) return 0b0111+sign; else return 0b0110+sign; else if(x > 0.00260417f) return 0b0001+sign; else return 0b0000+sign; } __device__ __forceinline__ float dDequantizeNF4(unsigned char val) { return nf4_dequantization_lut[val & 0x0F]; } __device__ unsigned char dQuantizeNF4(float x) { // the values for this tree was generated by test_normal_map_tree // in the file tests/test_functional.py if(x > 0.03979014977812767f) if(x > 0.3893125355243683f) // 1 if(x > 0.6427869200706482f) // 11 if(x > 0.8614784181118011f) // 111 return 0b1111; else return 0b1110; else if(x > 0.5016634166240692f) // 110 return 0b1101; else return 0b1100; else if(x > 0.2035212516784668f) // 10 if(x > 0.2920137718319893f) // 101 return 0b1011; else return 0b1010; else if(x > 0.1202552504837513f) // 100 return 0b1001; else return 0b1000; else if(x > -0.33967943489551544f) // 0 if(x > -0.13791173323988914f) // 01 if(x > -0.045525018125772476f) // 011 return 0b0111; else return 0b0110; else if(x > -0.23460740596055984f) // 010 return 0b0101; else return 0b0100; else if(x > -0.6106329262256622f) // 00 if(x > -0.4599952697753906f) // 001 return 0b0011; else return 0b0010; else if(x > -0.8480964004993439f) // 000 return 0b0001; else return 0b0000; } // sign function for lion // taken from https://stackoverflow.com/a/4609795, but not sure if there's a proper way to do this in CUDA template __device__ int sgn(T val) { return (T(0) < val) - (val < T(0)); } template __device__ unsigned char dQuantize(float* smem_code, const float rand, float x) { int pivot = 127; int upper_pivot = 255; int lower_pivot = 0; float lower = -1.0f; float upper = 1.0f; float val = smem_code[pivot]; // i>>=1 = {32, 16, 8, 4, 2, 1} for(int i = 64; i > 0; i>>=1) { if(x > val) { lower_pivot = pivot; lower = val; pivot+=i; } else { upper_pivot = pivot; upper = val; pivot-=i; } val = smem_code[pivot]; } if(upper_pivot == 255) upper = smem_code[upper_pivot]; if(lower_pivot == 0) lower = smem_code[lower_pivot]; if(!STOCHASTIC) { if(x > val) { float midpoint = (upper+val)*0.5f; if(x > midpoint) { return upper_pivot; } else return pivot; } else { float midpoint = (lower+val)*0.5f; if(x < midpoint) return lower_pivot; else return pivot; } } else { if(x > val) { float dist_to_upper = fabsf(upper-x); float dist_full = upper-val; if(rand >= dist_to_upper/dist_full) return upper_pivot; else return pivot; } else { float dist_to_lower = fabsf(lower-x); float dist_full = val-lower; if(rand >= dist_to_lower/dist_full) return lower_pivot; else return pivot; } } } template __device__ __forceinline__ unsigned char quantize_2D(float *__restrict__ quadrants, float *__restrict__ const smem_code, float x) { int pivot = 127; int upper_pivot = 255; int lower_pivot = 0; float lower = SIGNED ? -1.0f : 0.0f; float upper = 1.0f; float midpoint; float val = quadrants[1]; int local_pivot = 1; int offset = 1; // i>>=1 = {32, 16, 8, 4, 2, 1} for(int i = 64; i > 0; i>>=1) { if(x > val) { lower_pivot = pivot; lower = val; pivot+=i; //val = i == 64 ? quadrants[2] : smem_code[pivot]; local_pivot += offset; } else { upper_pivot = pivot; upper = val; pivot-=i; //val = i == 64 ? quadrants[0] : smem_code[pivot]; local_pivot -= offset; } val = i >= 64 ? quadrants[local_pivot] : smem_code[pivot]; offset -= 1; } if(x > val) { midpoint = (upper+val)*0.5f; if(x > midpoint) return upper_pivot; else return pivot; } else { midpoint = (lower+val)*0.5f; if(x < midpoint) return lower_pivot; else return pivot; } } __launch_bounds__(TH, 4) __global__ void kQuantize(float * code, float * __restrict__ const A, unsigned char *out, const int n) { const int n_full = (NUM_BLOCK*(n/NUM_BLOCK)) + (n % NUM_BLOCK == 0 ? 0 : NUM_BLOCK); int valid_items = (blockIdx.x+1 == gridDim.x) ? n - (blockIdx.x*NUM_BLOCK) : NUM_BLOCK; const int base_idx = (blockIdx.x * NUM_BLOCK); float vals[NUM]; unsigned char qvals[NUM]; //const int lane_id = threadIdx.x % 2; typedef hipcub::BlockLoad LoadFloat; typedef hipcub::BlockStore StoreChar; __shared__ typename LoadFloat::TempStorage loadf; __shared__ typename StoreChar::TempStorage storec; __shared__ float smem_code[256]; //__shared__ float smem_code[2][257]; if(threadIdx.x < 256) { smem_code[threadIdx.x] = code[threadIdx.x]; //smem_code[0][threadIdx.x] = code[threadIdx.x]; //smem_code[1][threadIdx.x] = smem_code[0][threadIdx.x]; } for (unsigned int i = base_idx; i < n_full; i += gridDim.x*NUM_BLOCK) { // number of values already processed in blocks + // number of values already processed in this block + // rand_offset % mod value valid_items = n - i > NUM_BLOCK ? NUM_BLOCK : n - i; __syncthreads(); LoadFloat(loadf).Load(&(A[i]), vals, valid_items); #pragma unroll 4 for(int j = 0; j < NUM; j++) qvals[j] = dQuantize<0>(smem_code, 0.0f, vals[j]); __syncthreads(); StoreChar(storec).Store(&(out[i]), qvals, valid_items); } } template //__launch_bounds__(TH, 4) __global__ void kQuantizeBlockwise(float * code, T * __restrict__ const A, float *absmax, unsigned char *out, float * __restrict__ const rand, const int rand_offset, const int n) { const int n_full = gridDim.x * BLOCK_SIZE; int valid_items = 0; const int base_idx = (blockIdx.x * BLOCK_SIZE); T vals[NUM_PER_TH]; float rand_vals[NUM_PER_TH]; unsigned char qvals[(DATA_TYPE > 0) ? NUM_PER_TH/2 : NUM_PER_TH]; //float local_abs_max = -FLT_MAX; float local_abs_max = 0.0f; int local_rand_idx = 0; typedef hipcub::BlockLoad LoadT; typedef hipcub::BlockStore 0) ? NUM_PER_TH/2 : NUM_PER_TH, hipcub::BLOCK_STORE_WARP_TRANSPOSE> StoreChar; typedef hipcub::BlockReduce BlockReduce; typedef hipcub::BlockLoad LoadFloat; __shared__ typename LoadT::TempStorage loadt; __shared__ typename LoadFloat::TempStorage loadf; __shared__ typename StoreChar::TempStorage storec; __shared__ typename BlockReduce::TempStorage reduce; __shared__ float smem_code[256]; __shared__ float smem_absmax_value[1]; if(DATA_TYPE == General8bit) for(int i = threadIdx.x; i < 256; i+=blockDim.x) smem_code[i] = code[i]; for (int i = base_idx; i < n_full; i += gridDim.x*BLOCK_SIZE) { valid_items = n - i > BLOCK_SIZE ? BLOCK_SIZE : n - i; local_abs_max = -FLT_MAX; __syncthreads(); LoadT(loadt).Load(&(A[i]), vals, valid_items, (T)0.0f); // 1. compute local max // 2. broadcast local max // 3. normalize inputs and quantize #pragma unroll NUM_PER_TH for(int j = 0; j < NUM_PER_TH; j++) local_abs_max = fmaxf(local_abs_max, fabsf((float)vals[j])); local_abs_max = BlockReduce(reduce).Reduce(local_abs_max, hipcub::Max(), valid_items); if(threadIdx.x == 0) { smem_absmax_value[0] = 1.0f / local_abs_max; absmax[i / BLOCK_SIZE] = local_abs_max; } __syncthreads(); local_abs_max = smem_absmax_value[0]; if(STOCHASTIC) { local_rand_idx = ((blockIdx.x*NUM_BLOCK) + (threadIdx.x*NUM) + rand_offset) % (1024-4); LoadFloat(loadf).Load(&rand[local_rand_idx], rand_vals, BLOCK_SIZE, 0); } switch(DATA_TYPE) { case General8bit: #pragma unroll NUM_PER_TH for(int j = 0; j < NUM_PER_TH; j++) { if(!STOCHASTIC) qvals[j] = dQuantize<0>(smem_code, 0.0f, ((float)vals[j])*local_abs_max); else qvals[j] = dQuantize<1>(smem_code, rand_vals[j], ((float)vals[j])*local_abs_max); } break; case FP4: #pragma unroll NUM_PER_TH for(int j = 0; j < NUM_PER_TH/2; j++) { qvals[j] = dQuantizeFP4(((float)vals[2 * j]) * local_abs_max) << 4; qvals[j] |= dQuantizeFP4(((float)vals[2 * j + 1]) * local_abs_max); } break; case NF4: #pragma unroll NUM_PER_TH for(int j = 0; j < NUM_PER_TH/2; j++) { qvals[j] = dQuantizeNF4(((float)vals[2 * j]) * local_abs_max) << 4; qvals[j] |= dQuantizeNF4(((float)vals[2 * j + 1]) * local_abs_max); } break; } __syncthreads(); StoreChar(storec).Store(&(out[(DATA_TYPE > 0) ? i/2 : i]), qvals, (DATA_TYPE > 0) ? (valid_items+1)/2 : valid_items); } } template __global__ void kDequantizeBlockwise(float *code, unsigned char * A, float * absmax, T *out, const int blocksize, const int n) { const int n_load = (gridDim.x * TILE_SIZE); int valid_items_load = 0; int valid_items_store = 0; const int base_idx = (blockIdx.x * TILE_SIZE); T vals[NUM_PER_TH*((DATA_TYPE > 0) ? 2 : 1)]; unsigned char qvals[NUM_PER_TH]; float local_abs_max = -FLT_MAX; typedef hipcub::BlockLoad LoadChar; typedef hipcub::BlockStore 0) ? 2 : 1), hipcub::BLOCK_STORE_WARP_TRANSPOSE> StoreT; __shared__ typename LoadChar::TempStorage loadchar; __shared__ typename StoreT::TempStorage storet; for (int i = base_idx; i < n_load; i += gridDim.x*TILE_SIZE) { if (DATA_TYPE > 0) { valid_items_load = min(TILE_SIZE, (n + 1) / 2 - i); valid_items_store = min(TILE_SIZE * 2, n - i * 2); } else { valid_items_load = min(TILE_SIZE, n - i); valid_items_store = valid_items_load; } // Since blocksize will always be a power-of-2, we avoid more expensive // division by the blocksize and instead use a shift operation. // This is equivalent to (i+threadId.x*NUM_PER_TH)/blocksize. local_abs_max = __ldg(&absmax[(i+threadIdx.x*NUM_PER_TH) >> (31 - __clz(blocksize))]); __syncthreads(); LoadChar(loadchar).Load(&(A[i]), qvals, valid_items_load, 128); switch (DATA_TYPE) { case General8bit: // load code through read-only cache via __ldg #pragma unroll NUM_PER_TH for(int j = 0; j < NUM_PER_TH; j++) vals[j] = __ldg(&code[qvals[j]])*local_abs_max; break; case FP4: #pragma unroll NUM_PER_TH for(int j = 0; j < NUM_PER_TH; j++) { vals[j * 2] = dDequantizeFP4Tree(qvals[j] >> 4) * local_abs_max; vals[j * 2 + 1] = dDequantizeFP4Tree(qvals[j] & 0x0F) * local_abs_max; } break; case NF4: #pragma unroll NUM_PER_TH for(int j = 0; j < NUM_PER_TH; j++) { vals[j*2] = dDequantizeNF4(qvals[j] >> 4)* local_abs_max; vals[j*2 + 1] = dDequantizeNF4(qvals[j] & 0x0F)* local_abs_max; } break; } __syncthreads(); StoreT(storet).Store(&(out[(DATA_TYPE > 0) ? i*2 : i]), vals, valid_items_store); } } __global__ void kDequantize(float *code, unsigned char *A, float *out, const int n) { const unsigned int numThreads = blockDim.x * gridDim.x; const int idx = (blockIdx.x * blockDim.x) + threadIdx.x; __shared__ float smem_code[256]; if(threadIdx.x < 256) { smem_code[threadIdx.x] = code[threadIdx.x]; } __syncthreads(); for (int i = idx;i < n; i += numThreads) { out[i] = smem_code[A[i]]; } } template __launch_bounds__(BLOCK_SIZE/NUM_VALS, 1) __global__ void kPreconditionOptimizer32bit2State(T* g, T* p, float* state1, float* state2, float *unorm, const float beta1, const float beta2, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const int n) { const int n_full = (BLOCK_SIZE*(n/BLOCK_SIZE)) + (n % BLOCK_SIZE == 0 ? 0 : BLOCK_SIZE); const int base_idx = (blockIdx.x * blockDim.x * NUM_VALS); int valid_items = 0; T g_vals[NUM_VALS]; float s1_vals[NUM_VALS]; float s2_vals[NUM_VALS]; const float correction1 = 1.0f/(1.0f - powf(beta1, step)); const float correction2 = 1.0f/(1.0f - powf(beta2, step)); typedef hipcub::BlockLoad Load; typedef hipcub::BlockLoad LoadFloat; typedef hipcub::BlockReduce BlockReduce; __shared__ union { typename Load::TempStorage load; typename LoadFloat::TempStorage loadf; typename BlockReduce::TempStorage reduce; } temp_storage; for (unsigned int i = base_idx; i < n_full; i += gridDim.x*BLOCK_SIZE) { valid_items = n - i >= (BLOCK_SIZE) ? (BLOCK_SIZE) : n - i; __syncthreads(); Load(temp_storage.load).Load(&(g[i]), g_vals, valid_items, 0.0f); __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state1[i]), s1_vals, valid_items, 0.0f); __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state2[i]), s2_vals, valid_items, 0.0f); # pragma unroll NUM_VALS for(unsigned int j = 0; j < NUM_VALS; j++) g_vals[j] = gnorm_scale*((float)g_vals[j]); # pragma unroll NUM_VALS for(unsigned int j = 0; j < NUM_VALS; j++) { switch(OPTIMIZER) { case ADAM: s1_vals[j] = s1_vals[j]*beta1 + ((1.0f -beta1)*((float)g_vals[j])); s2_vals[j] = s2_vals[j]*beta2 + ((1.0f -beta2)*(((float)g_vals[j])*((float)g_vals[j]))); s1_vals[j] *= correction1; s2_vals[j] *= correction2; s1_vals[j] = s1_vals[j]/(sqrtf(s2_vals[j])+eps); // update s1_vals[j] *= s1_vals[j]; // update l2 norm (update*update) break; } } # pragma unroll NUM_VALS-1 for(unsigned int j = 1; j < NUM_VALS; j++) s1_vals[0] += s1_vals[j]; __syncthreads(); s1_vals[0] = BlockReduce(temp_storage.reduce).Sum(s1_vals[0]); if(threadIdx.x == 0) atomicAdd(&unorm[0], s1_vals[0]); //__syncwarp(); } } #define NUM_PER_THREAD 4 template __launch_bounds__(TH, 1) __global__ void kOptimizer32bit2State(T* g, T* p, float* state1, float* state2, float *unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n) { const int n_full = ((TH*NUM_PER_THREAD)*(n/(TH*NUM_PER_THREAD))) + (n % (TH*NUM_PER_THREAD) == 0 ? 0 : (TH*NUM_PER_THREAD)); const int base_idx = (blockIdx.x * blockDim.x * NUM_PER_THREAD); int valid_items = 0; float update_scale = 0.0f; T g_vals[NUM_PER_THREAD]; T p_vals[NUM_PER_THREAD]; float s1_vals[NUM_PER_THREAD]; float s2_vals[NUM_PER_THREAD]; // AdEMAMix has an additional state buffer, which we packed // into state1. We need thread-local storage here for these. // TODO: Mark with [[maybe_unused]] after upgrade to min compiler. float s3_vals[NUM_PER_THREAD]; const float correction1 = 1.0f - powf(beta1, step); const float correction2 = sqrtf(1.0f - powf(beta2, step)); const float step_size = -lr*correction2/correction1; if(max_unorm > 0.0f) { update_scale = max_unorm > 0.0f ? sqrtf(unorm[0]) : 1.0f; if(update_scale > max_unorm*param_norm){ update_scale = (max_unorm*param_norm)/update_scale; } else{ update_scale = 1.0f; } } else{ update_scale = 1.0f; } typedef hipcub::BlockLoad Load; typedef hipcub::BlockStore Store; typedef hipcub::BlockLoad LoadFloat; typedef hipcub::BlockStore StoreFloat; __shared__ union { typename Load::TempStorage load; typename Store::TempStorage store; typename LoadFloat::TempStorage loadf; typename StoreFloat::TempStorage storef; } temp_storage; for (unsigned int i = base_idx; i < n_full; i += gridDim.x*TH*NUM_PER_THREAD) { valid_items = n - i >= (TH*NUM_PER_THREAD) ? (TH*NUM_PER_THREAD) : n - i; __syncthreads(); Load(temp_storage.load).Load(&(g[i]), g_vals, valid_items); __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state1[i]), s1_vals, valid_items); __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state2[i]), s2_vals, valid_items); __syncthreads(); Load(temp_storage.load).Load(&(p[i]), p_vals, valid_items); // Load additional state1 data for AdEMAMix // TODO: Make constexpr after updating min compiler if (OPTIMIZER == ADEMAMIX) { __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state1[n + i]), s3_vals, valid_items); } # pragma unroll 4 for(unsigned int j = 0; j < NUM_PER_THREAD; j++) g_vals[j] = gnorm_scale*((float)g_vals[j]); # pragma unroll 4 for(unsigned int j = 0; j < NUM_PER_THREAD; j++) { switch(OPTIMIZER) { case ADEMAMIX: // m1 update: m1 = beta1 * m1 + (1-beta1) * g s1_vals[j] = (s1_vals[j] * beta1) + ((1.0f - beta1) * (float)g_vals[j]); // m2 update: m2 = m2 * beta3 + (1-beta3) * g s3_vals[j] = (s3_vals[j] * beta3) + ((1.0f - beta3) * (float)g_vals[j]); // nu update: nu = beta2 * nu + (1-beta2) * g^2 s2_vals[j] = (s2_vals[j] * beta2) + ((1.0f - beta2) * (float)g_vals[j] * (float)g_vals[j]); p_vals[j] = (float)p_vals[j] - lr * ( ((s1_vals[j] / correction1) + (alpha * s3_vals[j])) / ( (sqrtf(s2_vals[j]) / correction2) + eps ) ); if (weight_decay > 0.0f) p_vals[j] = ((float)p_vals[j]) * (1.0f - (lr * weight_decay)); break; case ADAM: if(!skip_zeros || (skip_zeros && ((float)g_vals[j] != 0.0f))) { s1_vals[j] = s1_vals[j]*beta1 + ((1.0f -beta1)*((float)g_vals[j])); s2_vals[j] = s2_vals[j]*beta2 + ((1.0f -beta2)*(((float)g_vals[j])*((float)g_vals[j]))); p_vals[j] = ((float)p_vals[j]) + (update_scale*step_size*(s1_vals[j]/(sqrtf(s2_vals[j])+(eps*correction2)))); if(weight_decay > 0.0f) p_vals[j] = ((float)p_vals[j])*(1.0f-(lr*weight_decay)); } break; } } __syncthreads(); Store(temp_storage.store).Store(&(p[i]), p_vals, valid_items); __syncthreads(); StoreFloat(temp_storage.storef).Store(&(state1[i]), s1_vals, valid_items); __syncthreads(); StoreFloat(temp_storage.storef).Store(&(state2[i]), s2_vals, valid_items); if (OPTIMIZER == ADEMAMIX) { __syncthreads(); StoreFloat(temp_storage.storef).Store(&(state1[n + i]), s3_vals, valid_items); } } } template __launch_bounds__(BLOCK_SIZE/NUM_VALS, 1) __global__ void kPreconditionOptimizer32bit1State(T* g, T* p, float* state1, float *unorm, const float beta1, const float beta2, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const int n) { const int n_full = (BLOCK_SIZE*(n/BLOCK_SIZE)) + (n % BLOCK_SIZE == 0 ? 0 : BLOCK_SIZE); const int base_idx = (blockIdx.x * blockDim.x * NUM_VALS); int valid_items = 0; T g_vals[NUM_VALS]; float s1_vals[NUM_VALS]; typedef hipcub::BlockLoad Load; typedef hipcub::BlockLoad LoadFloat; typedef hipcub::BlockReduce BlockReduce; __shared__ union { typename Load::TempStorage load; typename LoadFloat::TempStorage loadf; typename BlockReduce::TempStorage reduce; } temp_storage; for (unsigned int i = base_idx; i < n_full; i += gridDim.x*BLOCK_SIZE) { valid_items = n - i >= (BLOCK_SIZE) ? (BLOCK_SIZE) : n - i; __syncthreads(); Load(temp_storage.load).Load(&(g[i]), g_vals, valid_items, 0.0f); __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state1[i]), s1_vals, valid_items, 0.0f); # pragma unroll NUM_VALS for(unsigned int j = 0; j < NUM_VALS; j++) g_vals[j] = gnorm_scale*((float)g_vals[j]); # pragma unroll NUM_VALS for(unsigned int j = 0; j < NUM_VALS; j++) { switch(OPTIMIZER) { case MOMENTUM: if(step == 1) s1_vals[j] = (float)g_vals[j]; // state update else s1_vals[j] = s1_vals[j]*beta1 + ((float)g_vals[j]); // state update s1_vals[j] = s1_vals[j]*s1_vals[j]; // update norm break; case LION: s1_vals[j] = s1_vals[j]*beta2 + ((1.0f-beta2)*(float)g_vals[j]); // state update break; case RMSPROP: s1_vals[j] = s1_vals[j]*beta1 + ((1.0f-beta1)*((float)g_vals[j])*((float)g_vals[j])); // state update s1_vals[j] = __fdividef((float)g_vals[j],sqrtf(s1_vals[j])+eps); // update value s1_vals[j] = s1_vals[j]*s1_vals[j]; // update norm break; case ADAGRAD: s1_vals[j] = s1_vals[j] + ((float)g_vals[j])*((float)g_vals[j]); // state update s1_vals[j] = __fdividef((float)g_vals[j],sqrtf(s1_vals[j])+eps); // update value s1_vals[j] = s1_vals[j]*s1_vals[j]; // update norm break; } } # pragma unroll for(unsigned int j = 1; j < NUM_VALS; j++) s1_vals[0] += s1_vals[j]; __syncthreads(); s1_vals[0] = BlockReduce(temp_storage.reduce).Sum(s1_vals[0], valid_items); if(threadIdx.x == 0) atomicAdd(&unorm[0], s1_vals[0]); //__syncwarp(); } } template __launch_bounds__(TH, 1) __global__ void kOptimizer32bit1State(T *g, T *p, float *state1, float *unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n) { const int n_full = ((TH*NUM_PER_THREAD)*(n/(TH*NUM_PER_THREAD))) + (n % (TH*NUM_PER_THREAD) == 0 ? 0 : (TH*NUM_PER_THREAD)); const int base_idx = (blockIdx.x * blockDim.x * NUM_PER_THREAD); int valid_items = 0; float update_scale = 0.0f; if(max_unorm > 0.0f) { update_scale = max_unorm > 0.0f ? sqrtf(unorm[0]) : 1.0f; if(update_scale > max_unorm*param_norm+eps){ update_scale = (max_unorm*param_norm+eps)/update_scale; } else{ update_scale = 1.0f; } } else{ update_scale = 1.0f; } T g_vals[NUM_PER_THREAD]; T p_vals[NUM_PER_THREAD]; float s1_vals[NUM_PER_THREAD]; typedef hipcub::BlockLoad Load; typedef hipcub::BlockStore Store; typedef hipcub::BlockLoad LoadFloat; typedef hipcub::BlockStore StoreFloat; __shared__ union { typename Load::TempStorage load; typename Store::TempStorage store; typename LoadFloat::TempStorage loadf; typename StoreFloat::TempStorage storef; } temp_storage; for (unsigned int i = base_idx; i < n_full; i += gridDim.x*TH*NUM_PER_THREAD) { valid_items = n - i >= (TH*NUM_PER_THREAD) ? (TH*NUM_PER_THREAD) : n - i; __syncthreads(); Load(temp_storage.load).Load(&(g[i]), g_vals, valid_items); __syncthreads(); LoadFloat(temp_storage.loadf).Load(&(state1[i]), s1_vals, valid_items); __syncthreads(); Load(temp_storage.load).Load(&(p[i]), p_vals, valid_items); # pragma unroll 4 for(unsigned int j = 0; j < NUM_PER_THREAD; j++) { g_vals[j] = gnorm_scale*((float)g_vals[j]); if(weight_decay > 0.0f) g_vals[j] = (float)g_vals[j] + (((float)p_vals[j])*weight_decay); } # pragma unroll 4 for(unsigned int j = 0; j < NUM_PER_THREAD; j++) { if(!skip_zeros || (skip_zeros && ((float)g_vals[j] != 0.0f))) { switch(OPTIMIZER) { case MOMENTUM: if(step == 1) s1_vals[j] = (float)g_vals[j]; else s1_vals[j] = s1_vals[j]*beta1 + ((float)g_vals[j]); p_vals[j] = ((float)p_vals[j]) + update_scale*(-lr*(s1_vals[j])); break; case LION: p_vals[j] = ((float)p_vals[j]) - update_scale*(lr*sgn(((float)s1_vals[j])*beta1 + ((1.0f-beta1)*((float)g_vals[j])))); s1_vals[j] = s1_vals[j]*beta2 + ((1.0f-beta2)*((float)g_vals[j])); break; case RMSPROP: s1_vals[j] = s1_vals[j]*beta1 + ((1.0f-beta1)*((float)g_vals[j])*((float)g_vals[j])); p_vals[j] = ((float)p_vals[j]) - update_scale*(lr*__fdividef((float)g_vals[j],sqrtf((float)s1_vals[j])+eps)); break; case ADAGRAD: s1_vals[j] = s1_vals[j] + ((float)g_vals[j])*((float)g_vals[j]); p_vals[j] = ((float)p_vals[j]) - lr*__fdividef((float)g_vals[j],sqrtf((float)s1_vals[j])+eps); break; } } } __syncthreads(); Store(temp_storage.store).Store(&(p[i]), p_vals, valid_items); __syncthreads(); StoreFloat(temp_storage.storef).Store(&(state1[i]), s1_vals, valid_items); } } #define NUM8BIT 16 #define NUM_THREADS 256 #define NUM_PER_BLOCK 4096 template __global__ void __launch_bounds__(NUM_THREADS, 2) kPreconditionOptimizerStatic8bit2State(T* p, T* __restrict__ const g, unsigned char*__restrict__ const state1, unsigned char* __restrict__ const state2, float *unorm, const float beta1, const float beta2, const float eps, const int step, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, const float gnorm_scale, const int n) { const int n_full = gridDim.x * NUM_PER_BLOCK; const int base_idx = (blockIdx.x * blockDim.x * NUM_PER_THREAD); int valid_items = n - (blockIdx.x*NUM_PER_BLOCK) > NUM_PER_BLOCK ? NUM_PER_BLOCK : n - (blockIdx.x*NUM_PER_BLOCK); float g_val = 0.0f; float local_max_s1 = -FLT_MAX; float local_max_s2 = -FLT_MAX; float local_unorm = 0.0f; float s2_vals[NUM8BIT]; float s1_vals[NUM8BIT]; T g_vals[NUM8BIT]; unsigned char m_c1[NUM8BIT]; unsigned char r_c2[NUM8BIT]; typedef hipcub::BlockLoad LoadT; typedef hipcub::BlockLoad LoadUInt8; typedef hipcub::BlockReduce BlockReduce; __shared__ union { typename LoadT::TempStorage loadh; typename LoadUInt8::TempStorage loadc; typename BlockReduce::TempStorage reduce; } temp_storage; __shared__ float smem_quantiles1[256]; __shared__ float smem_quantiles2[256]; if(threadIdx.x < 256) { smem_quantiles1[threadIdx.x] = quantiles1[threadIdx.x]; smem_quantiles2[threadIdx.x] = quantiles2[threadIdx.x]; } __syncthreads(); for (unsigned int i = base_idx; i < n_full; i += NUM_THREADS*gridDim.x*NUM8BIT) { valid_items = n - i >= (TH*NUM_PER_THREAD) ? (TH*NUM_PER_THREAD) : n - i; LoadT(temp_storage.loadh).Load(&(g[i]), g_vals, valid_items, (T)0.0f); __syncthreads(); LoadUInt8(temp_storage.loadc).Load(&(state1[i]), m_c1, valid_items, 128); __syncthreads(); LoadUInt8(temp_storage.loadc).Load(&(state2[i]), r_c2, valid_items, 128); __syncthreads(); #pragma unroll 16 for(int j = 0; j < NUM8BIT; j++) { g_val = g_vals[j]; g_val *= gnorm_scale; s1_vals[j] = smem_quantiles1[m_c1[j]]*max1[0]*beta1; s1_vals[j] += (1.0f-beta1)*g_val; local_max_s1 = fmaxf(local_max_s1, fabsf(s1_vals[j])); } #pragma unroll 16 for(int j = 0; j < NUM8BIT; j++) { g_val = g_vals[j]; g_val *= gnorm_scale; s2_vals[j] = smem_quantiles2[r_c2[j]]*max2[0]*beta2; s2_vals[j] += (1.0f-beta2)*g_val*g_val; local_max_s2 = fmaxf(local_max_s2, fabsf(s2_vals[j])); } if(unorm != NULL) { #pragma unroll 16 for(int j = 0; j < NUM8BIT; j++) { float correction1 = __fdividef(1.0f, 1.0f - powf(beta1, step)); float correction2 = __fdividef(1.0f, 1.0f - powf(beta2, step)); s1_vals[j] *= correction1; s2_vals[j] *= correction2; float update_val = s1_vals[j]/(sqrtf(s2_vals[j])+eps); // update local_unorm += update_val*update_val; } } } __syncthreads(); local_max_s1 = BlockReduce(temp_storage.reduce).Reduce(local_max_s1, hipcub::Max(), valid_items); __syncthreads(); local_max_s2 = BlockReduce(temp_storage.reduce).Reduce(local_max_s2, hipcub::Max(), valid_items); if(unorm != NULL) { __syncthreads(); local_unorm = BlockReduce(temp_storage.reduce).Reduce(local_unorm, hipcub::Sum(), valid_items); } if(threadIdx.x == 0) { atomicMax(&new_max1[0], local_max_s1); atomicMax(&new_max2[0], local_max_s2); if(unorm != NULL){ atomicAdd(&unorm[0], local_unorm); } } } #define NUM_PER_THREAD2 4 #define NUM_THREADS2 1024 #define NUM_PER_BLOCK2 4096 template __global__ void __launch_bounds__(NUM_THREADS2, 1) kOptimizerStatic8bit2State(T* p, T* const g, unsigned char* state1, unsigned char* state2, const float *unorm, const float max_unorm, const float param_norm, \ const float beta1, const float beta2, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, float weight_decay, const float gnorm_scale, const int n) { const int n_full = (blockDim.x * gridDim.x)*NUM_PER_THREAD2; const int base_idx = (blockIdx.x * blockDim.x * NUM_PER_THREAD2); int valid_items = 0; float g_val = 0.0f; float s1_vals[NUM_PER_THREAD2]; float s2_vals[NUM_PER_THREAD2]; const float correction1 = 1.0f - powf(beta1, step); const float correction2 = sqrtf(1.0f - powf(beta2, step)); const float step_size = -lr*correction2/correction1; //const float step_size = -lr*correction2/correction1; float new_max_val1 = 1.0f/new_max1[0]; float new_max_val2 = 1.0f/new_max2[0]; float update_scale = 1.0f; if(max_unorm > 0.0f) { update_scale = max_unorm > 0.0f ? sqrtf(unorm[0]) : 1.0f; if(update_scale > max_unorm*param_norm){ update_scale = (max_unorm*param_norm)/update_scale; } else{ update_scale = 1.0f; } } else{ update_scale = 1.0f; } unsigned char c1s[NUM_PER_THREAD2]; unsigned char c2s[NUM_PER_THREAD2]; T p_vals[NUM_PER_THREAD2]; T g_vals[NUM_PER_THREAD2]; typedef hipcub::BlockLoad LoadT; typedef hipcub::BlockLoad LoadChar; typedef hipcub::BlockStore StoreChar; typedef hipcub::BlockStore StoreT; __shared__ float smem_quantiles1[256]; __shared__ float smem_quantiles2[256]; __shared__ union { typename LoadT::TempStorage loadh; typename LoadChar::TempStorage loadc; typename StoreChar::TempStorage storec; typename StoreT::TempStorage storeh; } temp_storage; if(threadIdx.x < 512) { if(threadIdx.x < 256) smem_quantiles1[threadIdx.x] = quantiles1[threadIdx.x]; else smem_quantiles2[threadIdx.x-256] = quantiles2[threadIdx.x-256]; } __syncthreads(); for (unsigned int i = base_idx; i < n_full; i += gridDim.x*NUM_THREADS2*NUM_PER_THREAD2) { valid_items = n - i >= (TH*NUM_PER_THREAD) ? (TH*NUM_PER_THREAD) : n - i; LoadT(temp_storage.loadh).Load(&(g[i]), g_vals, valid_items, (T)0.0f); __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state1[i]), c1s, valid_items, 128); __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state2[i]), c2s, valid_items, 0); __syncthreads(); LoadT(temp_storage.loadh).Load(&(p[i]), p_vals, valid_items); if((i + (threadIdx.x*NUM_PER_THREAD2) + NUM_PER_THREAD2) > n){ continue; } # pragma unroll 4 for(unsigned int j = 0; j < NUM_PER_THREAD2; j++) { g_val = float(g_vals[j]); g_val *= gnorm_scale; s1_vals[j] = smem_quantiles1[c1s[j]]; s1_vals[j] = s1_vals[j]*max1[0]; s1_vals[j] = (s1_vals[j]*beta1) + (((1.0f-beta1)*g_val)); c1s[j] = dQuantize<0>(smem_quantiles1, 0.0f, s1_vals[j]*new_max_val1); // make sure state1 term has still the same sign after quantization // (not needed for state2 term which has only positive values) if(signbit(smem_quantiles1[c1s[j]]) != signbit(s1_vals[j])) { if(s1_vals[j] > 0.0f) c1s[j] += 1; else c1s[j] -= 1; } s2_vals[j] = smem_quantiles2[c2s[j]]; s2_vals[j] = s2_vals[j]*max2[0]; s2_vals[j] = (s2_vals[j]*beta2) + (((1.0f-beta2)*g_val*g_val)); c2s[j] = dQuantize<0>(smem_quantiles2, 0.0f, s2_vals[j]*new_max_val2); } # pragma unroll 4 for(unsigned int j = 0; j < NUM_PER_THREAD2; j++) { p_vals[j] = (T)(((float)p_vals[j]) + ((update_scale*step_size*(s1_vals[j]/(sqrtf(s2_vals[j])+(correction2*eps)))))); if(weight_decay > 0.0f) p_vals[j] = update_scale*((float)p_vals[j])*(1.0f-(lr*weight_decay)); } StoreT(temp_storage.storeh).Store(&(p[i]), p_vals, valid_items); __syncthreads(); StoreChar(temp_storage.storec).Store(&(state1[i]), c1s, valid_items); __syncthreads(); StoreChar(temp_storage.storec).Store(&(state2[i]), c2s, valid_items); __syncthreads(); } } template __global__ void __launch_bounds__(NUM_THREADS, 2) kPreconditionOptimizerStatic8bit1State(T* p, T* __restrict__ const g, unsigned char*__restrict__ const state1, float *unorm, const float beta1, const float beta2, const float eps, const int step, float* __restrict__ const quantiles1, float* max1, float* new_max1, const float weight_decay, const float gnorm_scale, const int n) { const int n_full = gridDim.x * NUM_PER_BLOCK; const int base_idx = (blockIdx.x * blockDim.x * NUM_PER_THREAD); int valid_items = n - (blockIdx.x*NUM_PER_BLOCK) > NUM_PER_BLOCK ? NUM_PER_BLOCK : n - (blockIdx.x*NUM_PER_BLOCK); float g_val = 0.0f; float local_max_s1 = -FLT_MAX; float local_unorm = 0.0f; float s1_vals[NUM8BIT]; T g_vals[NUM8BIT]; unsigned char m_c1[NUM8BIT]; typedef hipcub::BlockLoad LoadT; typedef hipcub::BlockLoad LoadUInt8; typedef hipcub::BlockReduce BlockReduce; __shared__ union { typename LoadT::TempStorage loadh; typename LoadUInt8::TempStorage loadc; typename BlockReduce::TempStorage reduce; } temp_storage; __shared__ float smem_quantiles1[256]; if(threadIdx.x < 256) smem_quantiles1[threadIdx.x] = quantiles1[threadIdx.x]; __syncthreads(); for (unsigned int i = base_idx; i < n_full; i += gridDim.x*NUM_THREADS*NUM8BIT) { valid_items = n - i >= (TH*NUM_PER_THREAD) ? (TH*NUM_PER_THREAD) : n - i; __syncthreads(); LoadT(temp_storage.loadh).Load(&(g[i]), g_vals, valid_items, (T)0.0f); __syncthreads(); LoadUInt8(temp_storage.loadc).Load(&(state1[i]), m_c1, valid_items, 128); #pragma unroll 16 for(int j = 0; j < NUM8BIT; j++) { g_val = g_vals[j]; g_val *= gnorm_scale; s1_vals[j] = smem_quantiles1[m_c1[j]]*max1[0]; switch(OPTIMIZER) { case ADAGRAD: case MOMENTUM: if(step == 1) s1_vals[j] = (float)g_vals[j]; else s1_vals[j] = s1_vals[j]*beta1 + ((float)g_vals[j]); if(unorm != NULL) local_unorm += s1_vals[j]*s1_vals[j]; break; case LION: s1_vals[j] = s1_vals[j]*beta2 + ((1.0f-beta2)*g_val); break; case RMSPROP: s1_vals[j] = s1_vals[j]*beta1 + ((1.0f-beta1)*(g_val*g_val)); break; } local_max_s1 = fmaxf(local_max_s1, fabsf(s1_vals[j])); } } __syncthreads(); local_max_s1 = BlockReduce(temp_storage.reduce).Reduce(local_max_s1, hipcub::Max(), valid_items); if(threadIdx.x == 0){ atomicMax(&new_max1[0], local_max_s1); } if(unorm != NULL) { __syncthreads(); local_unorm = BlockReduce(temp_storage.reduce).Reduce(local_unorm, hipcub::Sum(), valid_items); if(threadIdx.x == 0){ atomicAdd(&unorm[0], local_unorm); } } } template __global__ void __launch_bounds__(1024, 1) kOptimizerStatic8bit1State(T* p, T* const g, unsigned char* state1, const float *unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* max1, float* new_max1, float weight_decay, const float gnorm_scale, const int n) { const int n_full = (blockDim.x * gridDim.x)*NUM_PER_THREAD2; const int base_idx = (blockIdx.x * blockDim.x * NUM_PER_THREAD2); int valid_items = 0; float g_val = 0.0f; float s1_vals[NUM_PER_THREAD2]; float new_max_val1 = 1.0f/new_max1[0]; float update_scale = 1.0f; if(max_unorm > 0.0f) { update_scale = max_unorm > 0.0f ? sqrtf(unorm[0]) : 1.0f; if(update_scale > max_unorm*param_norm){ update_scale = (max_unorm*param_norm)/update_scale; } else{ update_scale = 1.0f; } } else{ update_scale = 1.0f; } unsigned char c1s[NUM_PER_THREAD2]; T p_vals[NUM_PER_THREAD2]; T g_vals[NUM_PER_THREAD2]; typedef hipcub::BlockLoad LoadT; typedef hipcub::BlockLoad LoadChar; typedef hipcub::BlockStore StoreChar; typedef hipcub::BlockStore StoreT; __shared__ float smem_quantiles1[256]; __shared__ union { typename LoadT::TempStorage loadh; typename LoadChar::TempStorage loadc; typename StoreChar::TempStorage storec; typename StoreT::TempStorage storeh; } temp_storage; if(threadIdx.x < 256) smem_quantiles1[threadIdx.x] = quantiles1[threadIdx.x]; __syncthreads(); for (unsigned int i = base_idx; i < n_full; i += gridDim.x*NUM_THREADS2*NUM_PER_THREAD2) { valid_items = n - i >= (TH*NUM_PER_THREAD) ? (TH*NUM_PER_THREAD) : n - i; LoadT(temp_storage.loadh).Load(&(g[i]), g_vals, valid_items, (T)0.0f); __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state1[i]), c1s, valid_items, 128); __syncthreads(); LoadT(temp_storage.loadh).Load(&(p[i]), p_vals, valid_items); if((i + (threadIdx.x*NUM_PER_THREAD2) + NUM_PER_THREAD2) > n){ continue; } # pragma unroll 4 for(unsigned int j = 0; j < NUM_PER_THREAD2; j++) { g_val = float(g_vals[j]); g_val *= gnorm_scale; if(weight_decay > 0.0f) { switch(OPTIMIZER) { case ADAGRAD: case MOMENTUM: case RMSPROP: g_val += ((float)p_vals[j])*weight_decay; break; case LION: p_vals[j] = ((float)p_vals[j])*(1.0f-lr*weight_decay); break; } } s1_vals[j] = smem_quantiles1[c1s[j]]*max1[0]; switch(OPTIMIZER){ case ADAGRAD: case MOMENTUM: if(step == 1) s1_vals[j] = g_vals[j]; else s1_vals[j] = s1_vals[j]*beta1 + ((float)g_vals[j]); p_vals[j] = ((float)p_vals[j]) + (-lr*update_scale*(s1_vals[j])); break; case LION: p_vals[j] = ((float)p_vals[j]) - (lr*sgn(((float)s1_vals[j])*beta1 + ((1.0f-beta1)*((float)g_val)))); s1_vals[j] = s1_vals[j]*beta2 + ((1.0f-beta2)*g_val); break; case RMSPROP: s1_vals[j] = s1_vals[j]*beta1 + ((1.0f-beta1)*(g_val*g_val)); p_vals[j] = ((float)p_vals[j]) - (lr*__fdividef(g_val,sqrtf(s1_vals[j])+eps)); break; } c1s[j] = dQuantize<0>(smem_quantiles1, 0.0f, s1_vals[j]*new_max_val1); // make sure state1 term has still the same sign after quantization if(signbit(smem_quantiles1[c1s[j]]) != signbit(s1_vals[j])) { if(s1_vals[j] > 0.0f) c1s[j] += 1; else c1s[j] -= 1; } } StoreT(temp_storage.storeh).Store(&(p[i]), p_vals, valid_items); __syncthreads(); StoreChar(temp_storage.storec).Store(&(state1[i]), c1s, valid_items); __syncthreads(); } } template __global__ void kPercentileClipping(T * __restrict__ g, float *gnorm_vec, int step, const int n) { const int n_full = (BLOCK_SIZE*(n/BLOCK_SIZE)) + (n % BLOCK_SIZE == 0 ? 0 : BLOCK_SIZE); int valid_items = 0; typedef hipcub::BlockReduce BlockReduce; typedef hipcub::BlockLoad LoadT; __shared__ typename BlockReduce::TempStorage reduce; __shared__ typename LoadT::TempStorage loadT; T vals[NUM_VALS]; float local_sum = 0.0f; for (unsigned int i = (blockIdx.x * BLOCK_SIZE); i < n_full; i += gridDim.x*BLOCK_SIZE) { valid_items = n - i > BLOCK_SIZE ? BLOCK_SIZE : n - i; local_sum = 0.0f; __syncthreads(); LoadT(loadT).Load(&(g[i]), vals, valid_items, (T)0.0f); #pragma unroll NUM_VALS for(int j = 0; j < NUM_VALS; j++) local_sum += ((float)vals[j])*((float)vals[j]); local_sum = BlockReduce(reduce).Sum(local_sum, valid_items); if(threadIdx.x == 0) { if(step == 1) { // initialize with the same norm for all positions //#pragma unroll 10 for(int j = 0; j < 100; j++) atomicAdd(&gnorm_vec[j], local_sum); } else atomicAdd(&gnorm_vec[step % 100], local_sum); } } } #define LANES 2 #define QUAD 3 template __launch_bounds__(256, 3) __global__ void kOptimizerStatic8bit2StateBlockwise( T* p, T* __restrict__ const g, unsigned char* state1, unsigned char* state2, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* absmax1, float* absmax2, float weight_decay, const float gnorm_scale, const bool skip_zeros, const int n ) { //const int n_full = n + (n%BLOCK_SIZE); const int n_full = gridDim.x * BLOCK_SIZE; const int base_idx = (blockIdx.x * BLOCK_SIZE); int valid_items = 0; float g_val = 0.0f; float s1_vals[N_PER_TH]; float s2_vals[N_PER_TH]; float s3_vals[N_PER_TH]; // 2-5% const float correction1 = 1.0f - __powf(beta1, step); const float correction2 = sqrtf(1.0f -__powf(beta2, step)); const float step_size = __fdividef(-lr*correction2,correction1); const int lane_id = threadIdx.x % LANES; float new_local_abs_max1 = -FLT_MAX; float new_local_abs_max2 = -FLT_MAX; float new_local_abs_max3 = -FLT_MAX; float quadrants1[QUAD]; float quadrants2[QUAD]; unsigned char c1s[N_PER_TH]; unsigned char c2s[N_PER_TH]; unsigned char c3s[N_PER_TH]; T g_vals[N_PER_TH]; T p_vals[N_PER_TH]; typedef hipcub::BlockLoad LoadT; typedef hipcub::BlockLoad LoadChar; typedef hipcub::BlockStore StoreChar; typedef hipcub::BlockStore StoreT; __shared__ float smem_quantiles1[LANES][257]; __shared__ float smem_quantiles2[LANES][257]; typedef hipcub::BlockReduce BlockReduce1; typedef hipcub::BlockReduce BlockReduce2; typedef hipcub::BlockReduce BlockReduce3; __shared__ typename BlockReduce1::TempStorage reduce1; __shared__ typename BlockReduce2::TempStorage reduce2; __shared__ typename BlockReduce2::TempStorage reduce3; __shared__ float smem_exchange1[1]; __shared__ float smem_exchange2[1]; __shared__ float smem_exchange3[1]; // [[maybe_unused]] __shared__ union { typename LoadT::TempStorage loadh; typename LoadChar::TempStorage loadc; typename StoreChar::TempStorage storec; typename StoreT::TempStorage storeh; } temp_storage; // init: 0.2 -> 0.23 // 0.23 -> 0.23 smem_quantiles1[0][threadIdx.x] = quantiles1[threadIdx.x]; smem_quantiles2[0][threadIdx.x] = quantiles2[threadIdx.x]; # pragma unroll for(unsigned int j = 1; j < LANES; j++) { smem_quantiles1[j][threadIdx.x] = smem_quantiles1[0][threadIdx.x]; smem_quantiles2[j][threadIdx.x] = smem_quantiles2[0][threadIdx.x]; } __syncthreads(); #pragma unroll for(int k = 0; k < QUAD; k++) { quadrants1[k] = smem_quantiles1[lane_id][(k*256/(QUAD+1)) + (256/(QUAD+1)-1)]; quadrants2[k] = smem_quantiles2[lane_id][(k*256/(QUAD+1)) + (256/(QUAD+1)-1)]; } for (unsigned int i = base_idx; i < n_full; i += gridDim.x*BLOCK_SIZE) { // loads: 0.23 -> 0.85/1.44 valid_items = n - i >= BLOCK_SIZE ? BLOCK_SIZE : n - i; __syncthreads(); LoadT(temp_storage.loadh).Load(&(g[i]), g_vals, valid_items, (T)0.0f); __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state1[i]), c1s, valid_items, 128); __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state2[i]), c2s, valid_items, 0); // AdEMAMix has an additional state packed into state1. if (OPTIMIZER == ADEMAMIX) { __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state1[n + i]), c3s, valid_items, 128); } new_local_abs_max1 = -FLT_MAX; new_local_abs_max2 = -FLT_MAX; new_local_abs_max3 = -FLT_MAX; // update: 2.48/1.57 -> 2.51/1.60 # pragma unroll N_PER_TH for(unsigned int j = 0; j < N_PER_TH; j++) { if(!isnan((float)g_vals[j]) && !isinf((float)g_vals[j])) { s2_vals[j] = smem_quantiles2[lane_id][c2s[j]]*absmax2[i/BLOCK_SIZE]; g_val = g_vals[j]; //float ratio = (g_val*g_val)/fmaxf(s2_vals[j], eps*eps); //g_val = ratio > 2.0f ? 2.0f*g_val/ratio : g_val; g_val *= gnorm_scale; s2_vals[j] = (s2_vals[j]*beta2) + (((1.0f-beta2)*g_val*g_val)); s1_vals[j] = smem_quantiles1[lane_id][c1s[j]]*absmax1[i/BLOCK_SIZE]; s1_vals[j] = (s1_vals[j]*beta1) + (((1.0f-beta1)*g_val)); if (OPTIMIZER == ADEMAMIX) { // The absmax for the third state is appended to absmax1 s3_vals[j] = smem_quantiles1[lane_id][c3s[j]] * absmax1[(n + i)/BLOCK_SIZE]; s3_vals[j] = (s3_vals[j] * beta3) + (((1.0f - beta3) * g_val)); } } else { s1_vals[j] = 0.0f; s2_vals[j] = 0.0f; if (OPTIMIZER == ADEMAMIX) { s3_vals[j] = 0.0f; } } new_local_abs_max1 = fmaxf(new_local_abs_max1, fabsf(s1_vals[j])); new_local_abs_max2 = fmaxf(new_local_abs_max2, fabsf(s2_vals[j])); if (OPTIMIZER == ADEMAMIX) { new_local_abs_max3 = fmaxf(new_local_abs_max3, fabsf(s3_vals[j])); } } // reduce: 2.51/1.60 -> 2.67/1.69 new_local_abs_max1 = BlockReduce1(reduce1).Reduce(new_local_abs_max1, hipcub::Max()); new_local_abs_max2 = BlockReduce2(reduce2).Reduce(new_local_abs_max2, hipcub::Max()); if (OPTIMIZER == ADEMAMIX) { new_local_abs_max3 = BlockReduce3(reduce3).Reduce(new_local_abs_max3, hipcub::Max()); } if(threadIdx.x == 0) { smem_exchange1[0] = new_local_abs_max1; smem_exchange2[0] = new_local_abs_max2; if (OPTIMIZER == ADEMAMIX) { smem_exchange3[0] = new_local_abs_max3; } } __syncthreads(); if(threadIdx.x == 0) { absmax1[i/BLOCK_SIZE] = new_local_abs_max1; absmax2[i/BLOCK_SIZE] = new_local_abs_max2; if (OPTIMIZER == ADEMAMIX) { absmax1[(n + i)/BLOCK_SIZE] = new_local_abs_max3; } } else { new_local_abs_max1 = smem_exchange1[0]; new_local_abs_max2 = smem_exchange2[0]; if (OPTIMIZER == ADEMAMIX) { new_local_abs_max3 = smem_exchange3[0]; } } __syncthreads(); LoadT(temp_storage.loadh).Load(&(p[i]), p_vals, valid_items, (T)0.0f); // reduce: 2.67/1.69 -> 2.67/1.70 # pragma unroll N_PER_TH for(unsigned int j = 0; j < N_PER_TH; j++) { //if(!skip_zeros || (skip_zeros && ((float)g_vals[j] != 0.0f))) if(!isnan((float)g_vals[j]) && !isinf((float)g_vals[j])) { if (OPTIMIZER == ADEMAMIX) { p_vals[j] = T((float)p_vals[j] - lr * ( ((s1_vals[j] / correction1) + (alpha * s3_vals[j])) / ( (sqrtf(s2_vals[j]) / correction2) + eps ) )); } else { p_vals[j] = (T)(((float)p_vals[j]) + ((step_size*(__fdividef(s1_vals[j],(sqrtf(s2_vals[j])+(correction2*eps))))))); } if(weight_decay > 0.0f) p_vals[j] = ((float)p_vals[j])*(1.0f-(lr*weight_decay)); } } // store: 0.85/1.44 -> 2.48/1.57 __syncthreads(); StoreT(temp_storage.storeh).Store(&(p[i]), p_vals, valid_items); // quantizaztion: 2.67/1.70 -> 3.4/3.3 # pragma unroll N_PER_TH for(unsigned int j = 0; j < N_PER_TH; j++) { c1s[j] = quantize_2D<1>(quadrants1, smem_quantiles1[lane_id], __fdividef(s1_vals[j],new_local_abs_max1)); c2s[j] = quantize_2D<0>(quadrants2, smem_quantiles2[lane_id], __fdividef(s2_vals[j],new_local_abs_max2)); // make sure state1 term has still the same sign after quantization // (not needed for state2 term which has only positive values) if(signbit(smem_quantiles1[lane_id][c1s[j]]) != signbit(s1_vals[j])) { if(s1_vals[j] > 0.0f) c1s[j] += 1; else c1s[j] -= 1; } if (OPTIMIZER == ADEMAMIX) { c3s[j] = quantize_2D<1>(quadrants1, smem_quantiles1[lane_id], __fdividef(s3_vals[j],new_local_abs_max3)); if (signbit(smem_quantiles1[lane_id][c3s[j]]) != signbit(s3_vals[j])) { c3s[j] += (s3_vals[j] > 0.0f) ? 1 : -1; } } } __syncthreads(); StoreChar(temp_storage.storec).Store(&(state1[i]), c1s, valid_items); __syncthreads(); StoreChar(temp_storage.storec).Store(&(state2[i]), c2s, valid_items); if (OPTIMIZER == ADEMAMIX) { __syncthreads(); StoreChar(temp_storage.storec).Store(&(state1[n + i]), c3s, valid_items); } } } #define LANES 2 #define QUAD 3 template __launch_bounds__(256, 3) __global__ void kOptimizerStatic8bit1StateBlockwise(T* p, T* __restrict__ const g, unsigned char* state1, const float beta1, const float beta2, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* absmax1, float weight_decay, const float gnorm_scale, const bool skip_zeros, const int n) { //const int n_full = n + (n%BLOCK_SIZE); const int n_full = gridDim.x * BLOCK_SIZE; const int base_idx = (blockIdx.x * BLOCK_SIZE); int valid_items = 0; float g_val = 0.0f; float s1_vals[N_PER_TH]; // 2-5% const int lane_id = threadIdx.x % LANES; float new_local_abs_max1 = -FLT_MAX; float quadrants1[QUAD]; unsigned char c1s[N_PER_TH]; T g_vals[N_PER_TH]; T p_vals[N_PER_TH]; typedef hipcub::BlockLoad LoadT; typedef hipcub::BlockLoad LoadChar; typedef hipcub::BlockStore StoreChar; typedef hipcub::BlockStore StoreT; __shared__ float smem_quantiles1[LANES][257]; typedef hipcub::BlockReduce BlockReduce1; __shared__ typename BlockReduce1::TempStorage reduce1; __shared__ float smem_exchange1[1]; __shared__ union { typename LoadT::TempStorage loadh; typename LoadChar::TempStorage loadc; typename StoreChar::TempStorage storec; typename StoreT::TempStorage storeh; } temp_storage; // init: 0.2 -> 0.23 // 0.23 -> 0.23 smem_quantiles1[0][threadIdx.x] = quantiles1[threadIdx.x]; # pragma unroll for(unsigned int j = 1; j < LANES; j++) smem_quantiles1[j][threadIdx.x] = smem_quantiles1[0][threadIdx.x]; __syncthreads(); #pragma unroll for(int k = 0; k < QUAD; k++) quadrants1[k] = smem_quantiles1[lane_id][(k*256/(QUAD+1)) + (256/(QUAD+1)-1)]; for (unsigned int i = base_idx; i < n_full; i += gridDim.x*BLOCK_SIZE) { // loads: 0.23 -> 0.85/1.44 valid_items = n - i >= BLOCK_SIZE ? BLOCK_SIZE : n - i; __syncthreads(); LoadT(temp_storage.loadh).Load(&(g[i]), g_vals, valid_items, (T)0.0f); __syncthreads(); LoadChar(temp_storage.loadc).Load(&(state1[i]), c1s, valid_items, 128); __syncthreads(); LoadT(temp_storage.loadh).Load(&(p[i]), p_vals, valid_items, (T)0.0f); new_local_abs_max1 = -FLT_MAX; // update: 2.48/1.57 -> 2.51/1.60 # pragma unroll N_PER_TH for(unsigned int j = 0; j < N_PER_TH; j++) { g_val = float(g_vals[j]); g_val *= gnorm_scale; if(!skip_zeros || (skip_zeros && ((float)g_vals[j] != 0.0f))) { if(weight_decay > 0.0f) { switch(OPTIMIZER) { case MOMENTUM: case ADAGRAD: case RMSPROP: g_val += ((float)p_vals[j])*weight_decay; break; case LION: p_vals[j] = ((float)p_vals[j])*(1.0f-lr*weight_decay); break; } } s1_vals[j] = smem_quantiles1[lane_id][c1s[j]]*absmax1[i/BLOCK_SIZE]; switch(OPTIMIZER) { case MOMENTUM: if(step == 1) s1_vals[j] = g_val; else s1_vals[j] = (s1_vals[j]*beta1) + g_val; break; case LION: // here, using gvals[j] to store the gradient smoothed by beta1 for the following parameter update, before the momentum is updated by beta2 g_vals[j] = lr*sgn(((float)s1_vals[j])*beta1 + ((1.0f-beta1)*g_val)); s1_vals[j] = s1_vals[j]*beta2 + ((1.0f-beta2)*g_val); break; case RMSPROP: s1_vals[j] = s1_vals[j]*beta1 + ((1.0f-beta1)*(g_val*g_val)); break; case ADAGRAD: s1_vals[j] = s1_vals[j] + (g_val*g_val); break; } } new_local_abs_max1 = fmaxf(new_local_abs_max1, fabsf(s1_vals[j])); } // reduce: 2.51/1.60 -> 2.67/1.69 new_local_abs_max1 = BlockReduce1(reduce1).Reduce(new_local_abs_max1, hipcub::Max()); if(threadIdx.x == 0) smem_exchange1[0] = new_local_abs_max1; __syncthreads(); if(threadIdx.x == 0) absmax1[i/BLOCK_SIZE] = new_local_abs_max1; else new_local_abs_max1 = smem_exchange1[0]; // reduce: 2.67/1.69 -> 2.67/1.70 # pragma unroll N_PER_TH for(unsigned int j = 0; j < N_PER_TH; j++) { if(!skip_zeros || (skip_zeros && ((float)g_vals[j] != 0.0f))) { switch(OPTIMIZER) { case MOMENTUM: p_vals[j] = ((float)p_vals[j]) - lr*(s1_vals[j]); break; case LION: p_vals[j] = ((float)p_vals[j]) - ((float)g_vals[j]); break; case RMSPROP: g_val = g_vals[j]; p_vals[j] = ((float)p_vals[j]) - lr*(__fdividef(g_val, sqrtf(s1_vals[j])+eps)); break; case ADAGRAD: g_val = g_vals[j]; p_vals[j] = ((float)p_vals[j]) - lr*(__fdividef(g_val, sqrtf(s1_vals[j])+eps)); break; } } } // store: 0.85/1.44 -> 2.48/1.57 __syncthreads(); StoreT(temp_storage.storeh).Store(&(p[i]), p_vals, valid_items); // quantizaztion: 2.67/1.70 -> 3.4/3.3 # pragma unroll N_PER_TH for(unsigned int j = 0; j < N_PER_TH; j++) { c1s[j] = quantize_2D<1>(quadrants1, smem_quantiles1[lane_id], __fdividef(s1_vals[j],new_local_abs_max1)); // make sure state1 term has still the same sign after quantization // (not needed for state2 term which has only positive values) if(signbit(smem_quantiles1[lane_id][c1s[j]]) != signbit(s1_vals[j])) { if(s1_vals[j] > 0.0f) c1s[j] += 1; else c1s[j] -= 1; } } __syncthreads(); StoreChar(temp_storage.storec).Store(&(state1[i]), c1s, valid_items); } } // Inputs: // A [rows, cols] // Outputs: // rowStats [rows] // out [rows, cols] template __launch_bounds__(1024, BNB_MAX_THREADS_PER_SM / 1024) __global__ void kInt8VectorQuant(T * __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols) { // For sm50/sm52 and CUDA < 12.2 we need to do the reduction in fp32. // Otherwise `T` is `fp16`. This can be removed when Maxwell is dropped. #if (__CUDACC_VER_MAJOR__ >= 12 && __CUDACC_VER_MINOR >= 2) || BNB_FP16_AVAILABLE using TReduction = T; #else using TReduction = float; #endif using BlockReduceT = hipcub::BlockReduce; // One block per row. // Threads load column values in a striped arrangement. // e.g. t0 reads row[0], row[0+nthreads], .. // and t1 reads row[1], row[1+nthreads], .. // Each thread will determine its local absmax. // We then do a blockwise reduction to determine the row's absmax. __shared__ typename BlockReduceT::TempStorage temp_storage; __shared__ TReduction smem_row_absmax; const int row_id = blockIdx.x; const T* row_data = A + (row_id * cols); // Threads will read the row values in a striped access pattern and find a local absmax. TReduction row_local_absmax = -FLT_MIN; for (int i = threadIdx.x; i < cols; i += THREADS) { const TReduction absval = fabsf(__ldcs(&(row_data[i]))); // For sparse decomposition, values outside of the threshold are not to be // included when calculating the row's absmax. if constexpr (SPARSE_DECOMP) { row_local_absmax = fmaxf(row_local_absmax, absval < TReduction(threshold) ? absval : row_local_absmax); } else { row_local_absmax = fmaxf(row_local_absmax, absval); } } // Reduce thread-local absmax across the block. const TReduction row_absmax = BlockReduceT(temp_storage).Reduce(row_local_absmax, hipcub::Max(), cols); if (threadIdx.x == 0) { // Save our block's absmax to shared memory for the quantization step. rowStats[row_id] = smem_row_absmax = row_absmax; } __syncthreads(); // Quantize row-wise. const float scale = __fdividef(127.0f, smem_row_absmax); for (int i = threadIdx.x; i < cols; i += THREADS) { float val = row_data[i]; if constexpr (SPARSE_DECOMP) { // For sparse decomposition, we do not want to quantize the outliers. // Instead they're zeroed out. out[row_id * cols + i] = fabs(val) < threshold ? __float2int_rn(val * scale) : 0; } else { out[row_id * cols + i] = __float2int_rn(val * scale); } } } template __launch_bounds__(1024, BNB_MAX_THREADS_PER_SM / 1024) __global__ void kgetRowStats(T * __restrict__ A, float *rowStats, float threshold, int rows, int cols) { using BlockReduceT = hipcub::BlockReduce; // One block per row. // Threads load column values in a striped arrangement. // e.g. t0 reads row[0], row[0+nthreads], .. // and t1 reads row[1], row[1+nthreads], .. // Each thread will determine its local absmax. // We then do a blockwise reduction to determine the row's absmax. __shared__ typename BlockReduceT::TempStorage temp_storage; const int row_id = blockIdx.x; const T* __restrict__ row_data = A + (row_id * cols); // Threads will read the row values in a striped access pattern and find a local absmax. float row_local_absmax = -FLT_MIN; for (int i = threadIdx.x; i < cols; i += THREADS) { const float absval = fabsf(row_data[i]); // For sparse decomposition, values outside of the threshold are not to be // included when calculating the row's absmax. if constexpr (SPARSE_DECOMP) { row_local_absmax = fmaxf(row_local_absmax, absval < threshold ? absval : row_local_absmax); } else { row_local_absmax = fmaxf(row_local_absmax, absval); } } // Reduce thread-local absmax across the block. // TODO: Consider algorithm BLOCK_REDUCE_RAKING_COMMUTATIVE_ONLY const float row_absmax = BlockReduceT(temp_storage).Reduce(row_local_absmax, hipcub::Max(), cols); if (threadIdx.x == 0) { // Save our block's absmax to shared memory for the quantization step. rowStats[row_id] = row_absmax; } } template __global__ void kgetRowStats(half * __restrict__ A, float *rowStats, float threshold, int rows, int cols); template __global__ void kgetRowStats(half * __restrict__ A, float *rowStats, float threshold, int rows, int cols); template __global__ void kInt8VectorQuant(half * __restrict__ A, int8_t *out, float *rowStats, float threshold, int rows, int cols); template __global__ void kInt8VectorQuant(half * __restrict__ A, int8_t *out, float *rowStats, float threshold, int rows, int cols); #define MM_DEQUANT_CONST 6.200012e-05f //1.0f/(127.0f*127.0f) template __global__ void kdequant_mm_int32_fp16( int* __restrict__ const A, float *__restrict__ const rowStats, float *__restrict__ const colStats, half *out, half *__restrict__ const bias, const int numRows, const int numCols, const int n ) { const int n_out = numRows * numCols; int block_offset = blockIdx.x * THREADS * ITEMS_PER_THREAD; int thread_offset = threadIdx.x * ITEMS_PER_THREAD; int local_values[ITEMS_PER_THREAD]; half local_output[ITEMS_PER_THREAD]; float local_rowStats[ITEMS_PER_THREAD]; float local_colStats[ITEMS_PER_THREAD]; float local_biasValue[ITEMS_PER_THREAD]; typedef hipcub::BlockLoad LoadInt32; __shared__ typename LoadInt32::TempStorage loadint32; int row_idx, col_idx; #pragma unroll ITEMS_PER_THREAD for(int j = 0; j < ITEMS_PER_THREAD; j++) { row_idx = (block_offset + thread_offset + j) / numCols; col_idx = (block_offset + thread_offset + j) % numCols; local_colStats[j] = col_idx >= numCols ? 0.0f : colStats[col_idx]; local_rowStats[j] = row_idx >= numRows ? 0.0f : rowStats[row_idx]; local_biasValue[j] = ((bias == nullptr) || (col_idx >= numCols)) ? 0.0f : __half2float(bias[col_idx]); } // Each block loads THREADS * ITEMS_PER_THREAD values from A int valid_items = block_offset + THREADS * ITEMS_PER_THREAD < n_out ? THREADS * ITEMS_PER_THREAD : n_out - block_offset; LoadInt32(loadint32).Load(&(A[block_offset]), local_values, valid_items, 0); #pragma unroll ITEMS_PER_THREAD for (int j = 0; j < ITEMS_PER_THREAD; ++j) { local_output[j] = __float2half( fmaf(local_values[j] * local_rowStats[j] * local_colStats[j], MM_DEQUANT_CONST, local_biasValue[j]) ); } #pragma unroll ITEMS_PER_THREAD for (int j = 0; j < ITEMS_PER_THREAD; j++) { int outIdx = block_offset + thread_offset + j; if (outIdx < n_out) { out[outIdx] = local_output[j]; } } } #define DENORM 1.0f/127.0f #define MAX_SPARSE_COUNT 32 #define SMEM_SIZE 8*256 #if defined(__GFX9__) #define WARP_SIZE 64 #else #define WARP_SIZE 32 #endif template __global__ void kspmm_coo_very_sparse_naive(int *max_count, int *max_idx, int *offset_rowidx, int *rowidx, int *colidx, half *values, T *B, half *out, float * __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB) { // 0. load balancing: We process rows with most columns first (count_vec)and we process one row per block // If a block finishes, the next one is scheduled. Since the last blocks like have fewer // elements they finish faster "fillin up" the gaps left by larger blocks // without tensor cores // 1. use rowidx_length to find what to load (as many blocks as there are rows) // 2. Load A into registers // 3. each warp loads all required rows of B but each warp is offset by k // 4. Do mma operations that accumulate into registers // 5. Each warp stores its output row into matrix C const int count = max_count[blockIdx.x]; const int local_max_idx = max_idx[blockIdx.x]; const int offset = local_max_idx == 0 ? 0 : offset_rowidx[local_max_idx-1]; const int local_row_idx = rowidx[offset]; const int warp_id = threadIdx.x / WARP_SIZE; const int warp_idx = threadIdx.x % WARP_SIZE; const int warp_offset = (warp_id*WARP_SIZE)*SPMM_ITEMS; const int num_items = BITS == 8 ? 8 : 8; int idx_col_B = warp_offset; int local_idx_col_B_offset = 0; half local_valA[MAX_SPARSE_COUNT]; int local_colidxA[MAX_SPARSE_COUNT]; half local_valC[SPMM_ITEMS]; T local_valsB[num_items]; half local_valOut[num_items]; // 128 byte loads per warp == 4 bytes per thread // 2. Load A into registers for(int j = 0; j < MAX_SPARSE_COUNT; j++) { local_valA[j] = j < count ? values[offset+j] : __float2half(0.0f); local_colidxA[j] = j < count ? colidx[offset+j] : 0; } // each thread processes SPMM_ITEMS=32 per iteration. We have 256 threads. 32*256=x192 // we expect each warp to be SPMM_ITEMS*WARP_SIZE apart // we have a total of 128 bytes for the bank with a bank size of 4 bytes // added 3 bytes = 6 values between warps should reduce bank conflicts __shared__ half smem_dequant_stats[SMEM_SIZE]; while(idx_col_B < colsB) { if(dequant_stats != NULL) { for(int i = threadIdx.x; i < SMEM_SIZE; i+=blockDim.x) if((idx_col_B+i-local_idx_col_B_offset) < colsB) smem_dequant_stats[i] = dequant_stats[idx_col_B+i-local_idx_col_B_offset]; __syncthreads(); } #pragma unroll SPMM_ITEMS for(int j = 0; j < SPMM_ITEMS; j++) local_valC[j] = 0.0f; #pragma unroll for(int i = 0; i < count; i++) { // 3. each warp loads all required rows of B but each warp is offset by k int row_offset = colsB*local_colidxA[i]; #pragma unroll SPMM_ITEMS for(int j = 0; j < SPMM_ITEMS; j+=num_items) { // 4. Multiply the tile -> accumulate outputs in shared memory until 128 bytes it reached int idx = idx_col_B + (warp_idx*SPMM_ITEMS) + j; if(idx >= colsB){ break; } if((idx+num_items < colsB)) { if(BITS == 8) reinterpret_cast(local_valsB)[0] = reinterpret_cast(B)[(row_offset+ idx)/num_items]; else reinterpret_cast(local_valsB)[0] = reinterpret_cast(B)[(row_offset+ idx)/num_items]; } else { #pragma unroll num_items for(int k = 0; k < num_items; k++) if(idx+k < colsB) local_valsB[k] = B[row_offset+idx+k]; else local_valsB[k] = 0.0f; } #pragma unroll num_items for(int k = 0; k < num_items; k++) { if(BITS == 8 && dequant_stats != NULL) // we do texture cache reads (__ldg) on dequant_stats which should be super fast { float valB = local_valsB[k]; float valA = local_valA[i]; if(valB != 0.0 && valA != 0.0) local_valC[j+k] = (float)local_valC[j+k] + ((float)smem_dequant_stats[idx+k-local_idx_col_B_offset])*DENORM*valB*valA; } else local_valC[j+k] = (float)local_valC[j+k] + (float)local_valsB[k]*(float)local_valA[i]; } } } int idx_row_C = (colsB*local_row_idx); #pragma unroll SPMM_ITEMS for(int j = 0; j < SPMM_ITEMS; j+=num_items) { //int idx_col_C = idx_col_B + (32*j) + warp_idx; int idx_col_C = idx_col_B + warp_idx*SPMM_ITEMS + j; int idx_val = idx_col_C + idx_row_C; if(idx_col_C +num_items < colsB) { // load outputs to do inplace addition reinterpret_cast(local_valOut)[0] = reinterpret_cast(out)[idx_val/num_items]; #pragma unroll num_items for(int k = 0; k < num_items; k++) local_valC[(j/num_items) + k] = (float)local_valC[(j/num_items) + k] + (float)local_valOut[k]; reinterpret_cast(out)[idx_val/num_items] = reinterpret_cast(local_valC)[j/num_items]; } else { #pragma unroll num_items for(int k = 0; k < num_items; k++) if(idx_col_C + k < colsB) out[idx_val+k] = (float)out[idx_val+k]+(float)local_valC[j+k]; } } idx_col_B += blockDim.x*SPMM_ITEMS; local_idx_col_B_offset += blockDim.x*SPMM_ITEMS; } } #define WARPS 3 template __global__ void gemm_device(int M, int N, int K, T * __restrict__ const A, T* B, T * out, int lda, int ldb, int ldc) { #if __CUDA_ARCH__ >= 750 using namespace nvcuda; int col_offset = blockIdx.x *32; const int warp_id = threadIdx.x / 32; const int half_warp_id = threadIdx.x / 16; const int half_warp_lane = threadIdx.x % 16; const int batch_size_warps = (WARPS-1)*2; const int val_per_iter = blockDim.x-32; T local_A[4]; T local_B[128]; const int a_tile_offset = 16; const int b_tile_offset = (16*32 + 16); __shared__ T smem_A[8*16 + (2*16*(batch_size_warps-1))]; __shared__ T smem_B[2*batch_size_warps*16*32 + (2*16*(batch_size_warps-1))]; //__shared__ T smem_C[8*32]; rocwmma::fragment a_frag; rocwmma::fragment b_frag; rocwmma::fragment c_frag; rocwmma::fill_fragment(c_frag, 0.0f); int ticktock = 0; int idx = 0 + threadIdx.x; int loaded_values = 0; // prefetch if(idx < K && warp_id < (WARPS-1)) { if(loaded_values == 0) { local_A[0] = A[idx]; local_A[1] = A[idx+(1*val_per_iter)]; local_A[2] = A[idx+(2*val_per_iter)]; local_A[3] = A[idx+(3*val_per_iter)]; #pragma unroll 32 for(int col = 0; col < 32; col++) { local_B[col] = B[(col_offset+col)*ldb+idx]; local_B[col+32] = B[(col_offset+col)*ldb+idx+(1*val_per_iter)]; local_B[col+64] = B[(col_offset+col)*ldb+idx+(2*val_per_iter)]; local_B[col+96] = B[(col_offset+col)*ldb+idx+(3*val_per_iter)]; } loaded_values = 3; } else { if(loaded_values == 3) { local_A[0] = local_A[1]; #pragma unroll 32 for(int col = 0; col < 32; col++) local_B[col] = local_B[col+(32)]; } else if(loaded_values == 2) { local_A[0] = local_A[2]; #pragma unroll 32 for(int col = 0; col < 32; col++) local_B[col] = local_B[col+(64)]; } else { local_A[0] = local_A[3]; #pragma unroll 32 for(int col = 0; col < 32; col++) local_B[col] = local_B[col+(96)]; } loaded_values--; } smem_A[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*a_tile_offset)] = local_A[0]; #pragma unroll 32 for(int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*b_tile_offset) + (col*16)] = local_B[col]; } else if(warp_id < (WARPS-1)) { local_A[0] = T(0.0); smem_A[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*a_tile_offset)] = 0.0f; #pragma unroll 32 for(int col = 0; col < 32; col++) local_B[col] = 0.0f; #pragma unroll 32 for(int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*b_tile_offset) + (col*16)] = 0.0f; } ticktock = ticktock == 0 ? 1 : 0; //for(int base_idx = blockDim.x-32; base_idx < K; base_idx+=blockDim.x-32) for(int base_idx = blockDim.x-32; base_idx < K; base_idx+=blockDim.x-32) { idx = base_idx + threadIdx.x; __syncthreads(); if(idx < K && warp_id < (WARPS-1)) { //local_A[0] = A[idx]; //#pragma unroll 32 //for(int col = 0; col < 32; col++) // local_B[col] = B[(col_offset+col)*ldb+idx]; if(loaded_values == 0) { local_A[0] = A[idx]; local_A[1] = A[idx+(1*val_per_iter)]; local_A[2] = A[idx+(2*val_per_iter)]; local_A[3] = A[idx+(3*val_per_iter)]; #pragma unroll 32 for(int col = 0; col < 32; col++) { local_B[col] = B[(col_offset+col)*ldb+idx]; local_B[col+32] = B[(col_offset+col)*ldb+idx+(1*val_per_iter)]; local_B[col+64] = B[(col_offset+col)*ldb+idx+(2*val_per_iter)]; local_B[col+96] = B[(col_offset+col)*ldb+idx+(3*val_per_iter)]; } loaded_values = 3; } else { if(loaded_values == 3) { local_A[0] = local_A[1]; #pragma unroll 32 for(int col = 0; col < 32; col++) local_B[col] = local_B[col+(32)]; } else if(loaded_values == 2) { local_A[0] = local_A[2]; #pragma unroll 32 for(int col = 0; col < 32; col++) local_B[col] = local_B[col+(64)]; } else { local_A[0] = local_A[3]; #pragma unroll 32 for(int col = 0; col < 32; col++) local_B[col] = local_B[col+(96)]; } loaded_values--; } smem_A[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*a_tile_offset)] = local_A[0]; #pragma unroll 32 for(int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*b_tile_offset) + (col*16)] = local_B[col]; } else if(warp_id < (WARPS-1)) { local_A[0] = T(0.0); smem_A[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*a_tile_offset)] = 0.0f; #pragma unroll 32 for(int col = 0; col < 32; col++) local_B[col] = 0.0f; #pragma unroll 32 for(int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*b_tile_offset) + (col*16)] = 0.0f; } ticktock = ticktock == 0 ? 1 : 0; if(warp_id == (WARPS-1)) for(int k = 0; k < batch_size_warps; k++) { rocwmma::load_matrix_sync(a_frag, &(smem_A[(ticktock*batch_size_warps + k)*a_tile_offset]), 16); // 111 mu rocwmma::load_matrix_sync(b_frag, &(smem_B[(ticktock*batch_size_warps + k)*b_tile_offset]), 16); // 35 mu rocwmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } } __syncthreads(); if(warp_id != (WARPS-1)){ return; } // only warp_id == (WARPS-1) from here int warp_lane = threadIdx.x % 32; ticktock = ticktock == 0 ? 1 : 0; for(int k = 0; k < batch_size_warps; k++) { rocwmma::load_matrix_sync(a_frag, &(smem_A[(ticktock*batch_size_warps + k)*a_tile_offset]), 16); // 111 mu rocwmma::load_matrix_sync(b_frag, &(smem_B[(ticktock*batch_size_warps + k)*b_tile_offset]), 16); // 35 mu rocwmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } // 129 mu if(warp_id == (WARPS-1)) rocwmma::store_matrix_sync(&(smem_A[0]), c_frag, 32, rocwmma::mem_row_major); if(col_offset + warp_lane < M) out[col_offset + warp_lane] = smem_A[warp_lane]; #endif } template __device__ void printnonzero(T *A, int num_values, const char * strval) { for(int i = 0; i < num_values; i++) if((float)A[i] != 0.0) printf("%s %i %f\n", strval, i, (float)A[i]); } template __global__ void kgemm_4bit_inference(int M, int N, int K, T * __restrict__ const A, unsigned char *B, float *absmax, T * out, int lda, int ldb, int ldc, int blocksize) { //// element-wise kernel //// 1. Load batch x k into registers + //// 2. Load k x k into registers //// 3. dequantize and store in second pair of k x k + //// 4. matmul //// 5. sum with cub //// 6. store outputs //// TC kernel //// use k warps per thread block //// 1. threadblock use read-only cache to read in register tile for A into shared memory //// 2. each warp loops over shared memory tiles of A of size 8x16 and loads them into fragments //// 3. each warp reads a segment of values 16x32 from B //// 4. do dequantization from register of B into second pair of registers //// 5. store (4) into fragment //// 6. matmul aggregate into fragment C //// 7. aggregate files of C into shared memory block C //// 8. sum (7) //// 9. write outputs to matmul output matrix #if __CUDA_ARCH__ >= 750 using namespace nvcuda; int col_offset = blockIdx.x *32; const int warp_id = threadIdx.x / 32; const int warp_idx = threadIdx.x % 32; const int half_warp_id = threadIdx.x / 16; const int half_warp_lane = threadIdx.x % 16; const int batch_size_warps = (WARPS-1)*2; T quant_map[16]; #pragma unroll 16 for(int i = 0; i < 16; i++) quant_map[i] = nf4_dequantization_lut[i]; //__shared__ T quant_map[16*160]; T local_A[2]; T local_B[64]; unsigned char local_B_4bit[32]; const int a_tile_offset = 16; const int b_tile_offset = (16*32 + 16); __shared__ T smem_A[8*16 + (16*(batch_size_warps-1))]; __shared__ T smem_B[2*batch_size_warps*16*32 + (2*16*(batch_size_warps-1))]; __shared__ T smem_C[8*32]; rocwmma::fragment a_frag; rocwmma::fragment b_frag; rocwmma::fragment c_frag; rocwmma::fill_fragment(c_frag, 0.0f); for(int i = threadIdx.x; i < (8*32); i+=blockDim.x) smem_C[i] = 0.0f; __syncthreads(); int ticktock = 0; int idx = 0 + threadIdx.x; int loaded_values = 0; // prefetch if(idx < K && warp_id < (WARPS-1)) { if(loaded_values == 0) { local_A[0] = A[idx]; local_A[1] = A[idx+blockDim.x-32]; #pragma unroll 32 for(int col = 0; col < 32; col++) local_B_4bit[col] = B[(col_offset+col)*ldb+idx]; loaded_values = 1; } else { local_A[0] = local_A[1]; loaded_values--; #pragma unroll 64 for(int col = 0; col < 64; col+=2) { //local_B[col] = dhDequantizeNF4(local_B_4bit[col/2] >> 4)*T(1.0f); //local_B[col+1] = dhDequantizeNF4(local_B_4bit[col/2] & 0x0F)*T(1.0f); //local_B[col] = d2DequantizeFP4(local_B_4bit[col/2] >> 4)*(float)(17.0); //local_B[col+1] = d2DequantizeFP4(local_B_4bit[col/2] & 0x0F)*(float)(17.0); //local_B[col] = 127*(local_B_4bit[col/2] >> 4)*(float)(17.0); //local_B[col+1] = 127*(local_B_4bit[col/2] & 0x0F)*(float)(17.0); //local_B[col] = quant_map[(local_B_4bit[col/2] >> 4)]*T(17.0); //local_B[col+1] = quant_map[(local_B_4bit[col/2] & 0x0F)]*T(17.0); local_B[col] = quant_map[160*(local_B_4bit[col/2] >> 4)+warp_idx]*T(17.0); local_B[col+1] = quant_map[160*(local_B_4bit[col/2] & 0x0F)+warp_idx]*T(17.0); } } smem_A[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*a_tile_offset)] = local_A[0]; #pragma unroll 32 for(int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*b_tile_offset) + (col*16)] = local_B[col]; } else if(warp_id < (WARPS-1)) { local_A[0] = T(0.0); smem_A[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*a_tile_offset)] = 0.0f; #pragma unroll 32 for(int col = 0; col < 32; col++) local_B[col] = 0.0f; #pragma unroll 32 for(int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*b_tile_offset) + (col*16)] = 0.0f; } ticktock = ticktock == 0 ? 1 : 0; //if(threadIdx.x == 0) //printf("aa %i %i\n", idx, loaded_values); //for(int base_idx = blockDim.x-32; base_idx < K; base_idx+=blockDim.x-32) for(int base_idx = blockDim.x-32; base_idx < K; base_idx+=blockDim.x-32) { idx = base_idx + threadIdx.x; //if(threadIdx.x == 0) //printf("%i %i\n", idx, loaded_values); //__syncthreads(); if(idx < K && warp_id < (WARPS-1)) { if(loaded_values == 0) { local_A[0] = A[idx]; local_A[1] = A[idx+blockDim.x-32]; #pragma unroll 32 for(int col = 0; col < 32; col++) { local_B_4bit[col] = B[(col_offset+col)*ldb+idx]; local_B_4bit[col+16] = B[(col_offset+col)*ldb+idx]; } loaded_values = 1; } else { local_A[0] = local_A[1]; loaded_values--; int absidx = (idx + col_offset)/blocksize; half local_absmax = __ldg(&(absmax[absidx])); #pragma unroll 64 for(int col = 0; col < 64; col+=2) { //local_B[col] = dhDequantizeNF4(local_B_4bit[col/2] >> 4)*T(absidx); //local_B[col+1] = dhDequantizeNF4(local_B_4bit[col/2] & 0x0F)*T(absidx); //local_B[col] = T(127)*T(local_B_4bit[col/2] >> 4)*T(absidx); //local_B[col+1] = T(127)*T(local_B_4bit[col/2] & 0x0F)*T(absidx); //local_B[col] = quant_map[160*(local_B_4bit[col/2] >> 4)+warp_idx]*T(local_absmax); //local_B[col+1] = quant_map[160*(local_B_4bit[col/2] & 0x0F)+warp_idx]*T(local_absmax); local_B[col] = quant_map[(local_B_4bit[col/2] >> 4)]*T(absidx); local_B[col+1] = quant_map[(local_B_4bit[col/2] & 0x0F)]*T(absidx); } //printnonzero(local_B, 128, ""); } smem_A[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*a_tile_offset)] = local_A[0]; #pragma unroll 32 for(int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*b_tile_offset) + (col*16)] = local_B[col]; } else if(warp_id < (WARPS-1)) { local_A[0] = T(0.0); smem_A[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*a_tile_offset)] = 0.0f; #pragma unroll 32 for(int col = 0; col < 32; col++) local_B[col] = 0.0f; #pragma unroll 32 for(int col = 0; col < 32; col++) smem_B[half_warp_lane + (((batch_size_warps*ticktock)+half_warp_id)*b_tile_offset) + (col*16)] = 0.0f; } ticktock = ticktock == 0 ? 1 : 0; if(warp_id == (WARPS-1)) for(int k = 0; k < batch_size_warps; k++) { rocwmma::load_matrix_sync(a_frag, &(smem_A[(ticktock*batch_size_warps + k)*a_tile_offset]), 16); // 111 mu rocwmma::load_matrix_sync(b_frag, &(smem_B[(ticktock*batch_size_warps + k)*b_tile_offset]), 16); // 35 mu rocwmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } } __syncthreads(); //if(threadIdx.x == 0) //{ // printnonzero(smem_A, 8*16 + (2*16*(batch_size_warps-1)), "A: "); // printnonzero(smem_B, 2*batch_size_warps*16*32 + (2*16*(batch_size_warps-1)), "B: "); //} if(warp_id != (WARPS-1)){ return; } // only warp_id == (WARPS-1) from here int warp_lane = threadIdx.x % 32; ticktock = ticktock == 0 ? 1 : 0; for(int k = 0; k < batch_size_warps; k++) { //if(warp_lane == 0) //printf("%i %i %i %i\n", (ticktock*batch_size_warps + k)*a_tile_offset, k, ticktock, threadIdx.x); rocwmma::load_matrix_sync(a_frag, &(smem_A[(ticktock*batch_size_warps + k)*a_tile_offset]), 16); // 111 mu rocwmma::load_matrix_sync(b_frag, &(smem_B[(ticktock*batch_size_warps + k)*b_tile_offset]), 16); // 35 mu rocwmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } // 129 mu if(warp_id == (WARPS-1)) rocwmma::store_matrix_sync(&(smem_C[0]), c_frag, 32, rocwmma::mem_row_major); //printnonzero(smem_C, 32, ""); if(col_offset + warp_lane < M) out[col_offset + warp_lane] = smem_C[warp_lane]; #endif } // No of 4bit values processed by each thread #define num_values_4bit 32 template __global__ void kgemm_4bit_inference_naive(int M, int N, int K, T * __restrict__ const A, unsigned char *B, float *absmax, const float *datatype, T * out, int lda, int ldb, int ldc, int blocksize) { // per threadblock: // load step-by-step in chunks of [warp_size,warps]: 1xwarp_size * [warp_size,warps] -> [1,warps] // 4 warps -> 4 loads per iter // 1xwarp_size * warp_sizex4 -> 1x4 outputs per thread block typedef hipcub::WarpReduce WarpReduce; __shared__ typename WarpReduce::TempStorage temp_storage[THREADS/WARP_SIZE]; const int warp_idx = threadIdx.x / WARP_SIZE; const int warp_lane = threadIdx.x % WARP_SIZE; const int row_B = (THREADS/WARP_SIZE)*blockIdx.x + warp_idx; const int offset_B = ldb * row_B; const int num_values_8bit = num_values_4bit/2; float local_C = 0.0f; unsigned char local_B_4bit[num_values_8bit]; T local_B[num_values_4bit/4]; T local_A[num_values_4bit/4]; __shared__ T quant_map[16]; T local_absmax = T(0.0f); if (threadIdx.x < 16) quant_map[threadIdx.x] = T(__ldg(&datatype[threadIdx.x])); //for(int i = threadIdx.x; i < 16; i++) //quant_map[i] = T(datatype[i]); __syncthreads(); // A: [1, K] // B: [M, K] for(int inner_idx = warp_lane*num_values_4bit; inner_idx < K; inner_idx += WARP_SIZE*num_values_4bit) { const int inner_idx_halved = inner_idx/2; // Since blocksize will always be a power-of-2, we avoid more expensive // division by the blocksize and instead use a shift operation. // This is equivalent to (i+threadId.x*NUM_PER_TH)/blocksize. const int absidx = ((2*offset_B)+inner_idx) >> (31 - __clz(blocksize)); local_absmax = __ldg(&(absmax[absidx])); if(row_B < M) { if((inner_idx_halved + num_values_8bit) < (K/2)) { // this is the most important for performance considerations reinterpret_cast(local_B_4bit)[0] = reinterpret_cast(B)[(offset_B+(inner_idx_halved))/(num_values_8bit)]; } else { #pragma unroll for(int j = 0; j < (num_values_8bit); j++) if((inner_idx_halved) + j < (K/2)) local_B_4bit[j] = B[offset_B+inner_idx_halved + j]; else local_B_4bit[j] = 0b01110111; } } else { #pragma unroll for(int j = 0; j < (num_values_8bit); j++) local_B_4bit[j] = 0b01110111; } for(int i = 0; i < 4; i++) { #pragma unroll for(int k = 0; k < num_values_8bit/4; k++) { #if BNB_BF16_AVAILABLE local_B[k*2] = quant_map[local_B_4bit[(i*num_values_8bit/4) + k] >> 4]*local_absmax; local_B[k*2 + 1] = quant_map[local_B_4bit[(i*num_values_8bit/4) + k] & 0x0F]*local_absmax; #else // bf16 multipliation not supported local_B[k*2] = T((float)quant_map[local_B_4bit[(i*num_values_8bit/4) + k] >> 4]*(float)local_absmax); local_B[k*2 + 1] = T((float)quant_map[local_B_4bit[(i*num_values_8bit/4) + k] & 0x0F]*(float)local_absmax); #endif } if(inner_idx+(num_values_4bit/4) + (i*num_values_4bit/4) < K) { // this is also relatively important for performance if(BITS==16) { reinterpret_cast(local_A)[0] = reinterpret_cast(A)[inner_idx/(num_values_4bit/4) + i]; } else { reinterpret_cast(local_A)[0] = reinterpret_cast(A)[inner_idx/(num_values_4bit/8) + (2*i) + 0]; reinterpret_cast(local_A)[1] = reinterpret_cast(A)[inner_idx/(num_values_4bit/8) + (2*i) + 1]; } } else #pragma unroll for(int k = 0; k < num_values_4bit/4; k++) if(inner_idx + (i*num_values_4bit/4) + k < K) local_A[k] = A[inner_idx + k + (i*num_values_4bit/4)]; else local_A[k] = T(0.0f); // accumulate in float; small performance hit for Ampere, but lower error for outputs #pragma unroll for(int k = 0; k < num_values_4bit/4; k++) { #if BNB_BF16_AVAILABLE local_C += (float)(local_A[k]*local_B[k]); #else // bf16 multipliation not supported local_C += ((float)local_A[k]*(float)local_B[k]); #endif } } } local_C = WarpReduce(temp_storage[warp_idx]).Sum(local_C); if(row_B < M && warp_lane == 0) out[row_B] = T(local_C); } template __global__ void kfunc(T *A, T *B, T value, long n) { for(long i = (blockDim.x*blockIdx.x) + threadIdx.x; i < n; i+=(blockDim.x*gridDim.x)) { switch(FUNC) { case FILL: A[i] = (T)value; break; case ARANGE: A[i] = (T)i; break; case _MUL: A[i] = A[i]*B[i]; break; } } } //============================================================== // TEMPLATE DEFINITIONS //============================================================== template __global__ void kfunc(float *A, float *B, float value, long n); template __global__ void kfunc(unsigned char *A, unsigned char *B, unsigned char value, long n); template __global__ void kfunc(float *A, float *B, float value, long n); template __global__ void kfunc(float *A, float *B, float value, long n); // these are not used and make no sense, but the compiler needs them //template __global__ void gemm_device(int M, int N, int K, float * __restrict__ const A, float* B, float * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); //template __global__ void gemm_device(int M, int N, int K, float * __restrict__ const A, float* B, float * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); // these are not used and make no sense, but the compiler needs them //template __global__ void gemm_device(int M, int N, int K, float * __restrict__ const A, float* B, float * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); //template __global__ void gemm_device(int M, int N, int K, float * __restrict__ const A, float* B, float * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); template __global__ void gemm_device(int M, int N, int K, half * __restrict__ const A, half* B, half * out, int lda, int ldb, int ldc); template __global__ void kgemm_4bit_inference(int M, int N, int K, half * __restrict__ const A, unsigned char *B, float *absmax, half * out, int lda, int ldb, int ldc, int blocksize); template __global__ void kgemm_4bit_inference(int M, int N, int K, half * __restrict__ const A, unsigned char *B, float *absmax, half * out, int lda, int ldb, int ldc, int blocksize); template __global__ void kgemm_4bit_inference(int M, int N, int K, half * __restrict__ const A, unsigned char *B, float *absmax, half * out, int lda, int ldb, int ldc, int blocksize); template __global__ void kgemm_4bit_inference(int M, int N, int K, half * __restrict__ const A, unsigned char *B, float *absmax, half * out, int lda, int ldb, int ldc, int blocksize); template __global__ void kgemm_4bit_inference_naive(int M, int N, int K, half * __restrict__ const A, unsigned char *B, float *absmax, const float *datatype, half * out, int lda, int ldb, int ldc, int blocksize); template __global__ void kgemm_4bit_inference_naive(int M, int N, int K, hip_bfloat16 * __restrict__ const A, unsigned char *B, float *absmax, const float *datatype, hip_bfloat16 * out, int lda, int ldb, int ldc, int blocksize); template __global__ void kgemm_4bit_inference_naive(int M, int N, int K, float * __restrict__ const A, unsigned char *B, float *absmax, const float *datatype, float * out, int lda, int ldb, int ldc, int blocksize); template __global__ void kspmm_coo_very_sparse_naive(int *max_count, int *max_idx, int *offset_rowidx, int *rowidx, int *colidx, half *values, half *B, half *out, float * __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB); template __global__ void kspmm_coo_very_sparse_naive(int *max_count, int *max_idx, int *offset_rowidx, int *rowidx, int *colidx, half *values, half *B, half *out, float * __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB); template __global__ void kspmm_coo_very_sparse_naive(int *max_count, int *max_idx, int *offset_rowidx, int *rowidx, int *colidx, half *values, half *B, half *out, float * __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB); template __global__ void kspmm_coo_very_sparse_naive(int *max_count, int *max_idx, int *offset_rowidx, int *rowidx, int *colidx, half *values, signed char *B, half *out, float * __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB); template __global__ void kspmm_coo_very_sparse_naive(int *max_count, int *max_idx, int *offset_rowidx, int *rowidx, int *colidx, half *values, signed char *B, half *out, float * __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB); template __global__ void kspmm_coo_very_sparse_naive(int *max_count, int *max_idx, int *offset_rowidx, int *rowidx, int *colidx, half *values, signed char *B, half *out, float * __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB); template __global__ void kdequant_mm_int32_fp16<4, 512>(int *__restrict__ const A, float *__restrict__ const rowStats, float *__restrict__ const colStats, half *out, half * __restrict__ const bias, const int numRows, const int numCols, const int n); template __device__ unsigned char dQuantize<0>(float* smem_code, const float rand, float x); template __device__ unsigned char dQuantize<1>(float* smem_code, const float rand, float x); #define MAKE_PreconditionOptimizer32bit1State(oname, gtype) \ template __global__ void kPreconditionOptimizer32bit1State(gtype* g, gtype* p, \ float* state1, float *unorm, \ const float beta1, const float beta2, const float eps, const float weight_decay, \ const int step, const float lr, const float gnorm_scale, const int n); \ MAKE_PreconditionOptimizer32bit1State(MOMENTUM, half) MAKE_PreconditionOptimizer32bit1State(MOMENTUM, float) MAKE_PreconditionOptimizer32bit1State(MOMENTUM, hip_bfloat16) MAKE_PreconditionOptimizer32bit1State(RMSPROP, half) MAKE_PreconditionOptimizer32bit1State(RMSPROP, float) MAKE_PreconditionOptimizer32bit1State(RMSPROP, hip_bfloat16) MAKE_PreconditionOptimizer32bit1State(LION, half) MAKE_PreconditionOptimizer32bit1State(LION, float) MAKE_PreconditionOptimizer32bit1State(LION, hip_bfloat16) MAKE_PreconditionOptimizer32bit1State(ADAGRAD, half) MAKE_PreconditionOptimizer32bit1State(ADAGRAD, float) MAKE_PreconditionOptimizer32bit1State(ADAGRAD, hip_bfloat16) #define MAKE_Optimizer32bit1State(oname, gtype) \ template __global__ void kOptimizer32bit1State(gtype* g, gtype* p, float* state1, float *unorm, const float max_unorm, const float param_norm, \ const float beta1, const float beta2, const float eps, const float weight_decay,const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n); \ MAKE_Optimizer32bit1State(MOMENTUM, half) MAKE_Optimizer32bit1State(MOMENTUM, float) MAKE_Optimizer32bit1State(MOMENTUM, hip_bfloat16) MAKE_Optimizer32bit1State(RMSPROP, half) MAKE_Optimizer32bit1State(RMSPROP, float) MAKE_Optimizer32bit1State(RMSPROP, hip_bfloat16) MAKE_Optimizer32bit1State(LION, half) MAKE_Optimizer32bit1State(LION, float) MAKE_Optimizer32bit1State(LION, hip_bfloat16) MAKE_Optimizer32bit1State(ADAGRAD, half) MAKE_Optimizer32bit1State(ADAGRAD, float) MAKE_Optimizer32bit1State(ADAGRAD, hip_bfloat16) #define MAKE_PreconditionOptimizer32bit2State(oname, gtype) \ template __global__ void kPreconditionOptimizer32bit2State(gtype* g, gtype* p, \ float* state1, float* state2, float *unorm, \ const float beta1, const float beta2, const float eps, const float weight_decay, \ const int step, const float lr, const float gnorm_scale, const int n); \ MAKE_PreconditionOptimizer32bit2State(ADAM, float) MAKE_PreconditionOptimizer32bit2State(ADAM, half) MAKE_PreconditionOptimizer32bit2State(ADAM, hip_bfloat16) MAKE_PreconditionOptimizer32bit2State(ADEMAMIX, float) MAKE_PreconditionOptimizer32bit2State(ADEMAMIX, half) MAKE_PreconditionOptimizer32bit2State(ADEMAMIX, hip_bfloat16) template __global__ void kOptimizer32bit2State(float* g, float* p, float* state1, float* state2, float *unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay,const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n); template __global__ void kOptimizer32bit2State(half* g, half* p, float* state1, float* state2, float *unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay,const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n); template __global__ void kOptimizer32bit2State(hip_bfloat16* g, hip_bfloat16* p, float* state1, float* state2, float *unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay,const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n); template __global__ void kOptimizer32bit2State(float* g, float* p, float* state1, float* state2, float *unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay,const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n); template __global__ void kOptimizer32bit2State(half* g, half* p, float* state1, float* state2, float *unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay,const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n); template __global__ void kOptimizer32bit2State(hip_bfloat16* g, hip_bfloat16* p, float* state1, float* state2, float *unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay,const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n); #define MAKE_PreconditionStatic8bit1State(oname, gtype) \ template __global__ void kPreconditionOptimizerStatic8bit1State(gtype* p, gtype* __restrict__ const g, unsigned char*__restrict__ const state1, \ float *unorm, \ const float beta1, \ const float beta2, \ const float eps, const int step, \ float* __restrict__ const quantiles1, \ float* max1, float* new_max1, \ const float weight_decay, \ const float gnorm_scale, \ const int n); \ MAKE_PreconditionStatic8bit1State(MOMENTUM, half) MAKE_PreconditionStatic8bit1State(MOMENTUM, float) MAKE_PreconditionStatic8bit1State(RMSPROP, half) MAKE_PreconditionStatic8bit1State(RMSPROP, float) MAKE_PreconditionStatic8bit1State(LION, half) MAKE_PreconditionStatic8bit1State(LION, float) MAKE_PreconditionStatic8bit1State(ADAGRAD, half) MAKE_PreconditionStatic8bit1State(ADAGRAD, float) #define MAKE_optimizerStatic8bit1State(oname, gtype) \ template __global__ void kOptimizerStatic8bit1State(gtype* p, gtype* const g, unsigned char* state1, \ const float *unorm, const float max_unorm, const float param_norm, \ const float beta1, \ const float beta2, \ const float eps, const int step, const float lr, \ float* __restrict__ const quantiles1, \ float* max1, float* new_max1, \ float weight_decay, \ const float gnorm_scale, \ const int n); \ MAKE_optimizerStatic8bit1State(MOMENTUM, half) MAKE_optimizerStatic8bit1State(MOMENTUM, float) MAKE_optimizerStatic8bit1State(RMSPROP, half) MAKE_optimizerStatic8bit1State(RMSPROP, float) MAKE_optimizerStatic8bit1State(LION, half) MAKE_optimizerStatic8bit1State(LION, float) MAKE_optimizerStatic8bit1State(ADAGRAD, half) MAKE_optimizerStatic8bit1State(ADAGRAD, float) #define MAKE_PreconditionStatic8bit2State(oname, gtype) \ template __global__ void kPreconditionOptimizerStatic8bit2State(gtype* p, gtype* __restrict__ const g, unsigned char*__restrict__ const state1, unsigned char* __restrict__ const state2, \ float *unorm, \ const float beta1, const float beta2, \ const float eps, const int step, \ float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, \ float* max1, float* max2, float* new_max1, float* new_max2, \ const float gnorm_scale, \ const int n); \ MAKE_PreconditionStatic8bit2State(ADAM, half) MAKE_PreconditionStatic8bit2State(ADAM, float) #define MAKE_optimizerStatic8bit2State(oname, gtype) \ template __global__ void kOptimizerStatic8bit2State(gtype* p, gtype* const g, unsigned char* state1, unsigned char* state2, \ const float *unorm, const float max_unorm, const float param_norm, \ const float beta1, const float beta2, \ const float eps, const int step, const float lr, \ float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, \ float* max1, float* max2, float* new_max1, float* new_max2, \ float weight_decay, \ const float gnorm_scale, \ const int n); \ MAKE_optimizerStatic8bit2State(ADAM, half) MAKE_optimizerStatic8bit2State(ADAM, float) template __global__ void kPercentileClipping(float * __restrict__ g, float *gnorm_vec, int step, const int n); template __global__ void kPercentileClipping(half * __restrict__ g, float *gnorm_vec, int step, const int n); #define MAKE_kQuantizeBlockwise(dtype, blocksize, num_per_thread, stochastic, data_type_name) \ template __global__ void kQuantizeBlockwise(float * code, dtype * __restrict__ const A, float *absmax, unsigned char *out, float * __restrict__ const rand, const int rand_offset, const int n); \ MAKE_kQuantizeBlockwise(half, 4096, 4, 0, General8bit) MAKE_kQuantizeBlockwise(half, 4096, 4, 1, General8bit) MAKE_kQuantizeBlockwise(half, 2048, 4, 0, General8bit) MAKE_kQuantizeBlockwise(half, 1024, 4, 0, General8bit) MAKE_kQuantizeBlockwise(half, 512, 2, 0, General8bit) MAKE_kQuantizeBlockwise(half, 256, 2, 0, General8bit) MAKE_kQuantizeBlockwise(half, 128, 2, 0, General8bit) //MAKE_kQuantizeBlockwise(half, 64, 2, 0, General8bit) MAKE_kQuantizeBlockwise(half, 4096, 4, 0, FP4) MAKE_kQuantizeBlockwise(half, 2048, 4, 0, FP4) MAKE_kQuantizeBlockwise(half, 1024, 4, 0, FP4) MAKE_kQuantizeBlockwise(half, 512, 2, 0, FP4) MAKE_kQuantizeBlockwise(half, 256, 2, 0, FP4) MAKE_kQuantizeBlockwise(half, 128, 2, 0, FP4) //MAKE_kQuantizeBlockwise(half, 64, 2, 0, FP4) MAKE_kQuantizeBlockwise(half, 4096, 4, 0, NF4) MAKE_kQuantizeBlockwise(half, 2048, 4, 0, NF4) MAKE_kQuantizeBlockwise(half, 1024, 4, 0, NF4) MAKE_kQuantizeBlockwise(half, 512, 2, 0, NF4) MAKE_kQuantizeBlockwise(half, 256, 2, 0, NF4) MAKE_kQuantizeBlockwise(half, 128, 2, 0, NF4) //MAKE_kQuantizeBlockwise(half, 64, 2, 0, NF4) MAKE_kQuantizeBlockwise(float, 4096, 4, 0, General8bit) MAKE_kQuantizeBlockwise(float, 4096, 4, 1, General8bit) MAKE_kQuantizeBlockwise(float, 2048, 4, 0, General8bit) MAKE_kQuantizeBlockwise(float, 1024, 4, 0, General8bit) MAKE_kQuantizeBlockwise(float, 512, 2, 0, General8bit) MAKE_kQuantizeBlockwise(float, 256, 2, 0, General8bit) MAKE_kQuantizeBlockwise(float, 128, 2, 0, General8bit) //MAKE_kQuantizeBlockwise(float, 64, 2, 0, General8bit) MAKE_kQuantizeBlockwise(float, 4096, 4, 0, FP4) MAKE_kQuantizeBlockwise(float, 2048, 4, 0, FP4) MAKE_kQuantizeBlockwise(float, 1024, 4, 0, FP4) MAKE_kQuantizeBlockwise(float, 512, 2, 0, FP4) MAKE_kQuantizeBlockwise(float, 256, 2, 0, FP4) MAKE_kQuantizeBlockwise(float, 128, 2, 0, FP4) //MAKE_kQuantizeBlockwise(float, 64, 2, 0, FP4) MAKE_kQuantizeBlockwise(float, 4096, 4, 0, NF4) MAKE_kQuantizeBlockwise(float, 2048, 4, 0, NF4) MAKE_kQuantizeBlockwise(float, 1024, 4, 0, NF4) MAKE_kQuantizeBlockwise(float, 512, 2, 0, NF4) MAKE_kQuantizeBlockwise(float, 256, 2, 0, NF4) MAKE_kQuantizeBlockwise(float, 128, 2, 0, NF4) //MAKE_kQuantizeBlockwise(float, 64, 2, 0, NF4) MAKE_kQuantizeBlockwise(hip_bfloat16, 4096, 4, 0, General8bit) MAKE_kQuantizeBlockwise(hip_bfloat16, 4096, 4, 1, General8bit) MAKE_kQuantizeBlockwise(hip_bfloat16, 2048, 4, 0, General8bit) MAKE_kQuantizeBlockwise(hip_bfloat16, 1024, 4, 0, General8bit) MAKE_kQuantizeBlockwise(hip_bfloat16, 512, 2, 0, General8bit) MAKE_kQuantizeBlockwise(hip_bfloat16, 256, 2, 0, General8bit) MAKE_kQuantizeBlockwise(hip_bfloat16, 128, 2, 0, General8bit) //MAKE_kQuantizeBlockwise(hip_bfloat16, 64, 2, 0, General8bit) MAKE_kQuantizeBlockwise(hip_bfloat16, 4096, 4, 0, FP4) MAKE_kQuantizeBlockwise(hip_bfloat16, 2048, 4, 0, FP4) MAKE_kQuantizeBlockwise(hip_bfloat16, 1024, 4, 0, FP4) MAKE_kQuantizeBlockwise(hip_bfloat16, 512, 2, 0, FP4) MAKE_kQuantizeBlockwise(hip_bfloat16, 256, 2, 0, FP4) MAKE_kQuantizeBlockwise(hip_bfloat16, 128, 2, 0, FP4) //MAKE_kQuantizeBlockwise(hip_bfloat16, 64, 2, 0, FP4) MAKE_kQuantizeBlockwise(hip_bfloat16, 4096, 4, 0, NF4) MAKE_kQuantizeBlockwise(hip_bfloat16, 2048, 4, 0, NF4) MAKE_kQuantizeBlockwise(hip_bfloat16, 1024, 4, 0, NF4) MAKE_kQuantizeBlockwise(hip_bfloat16, 512, 2, 0, NF4) MAKE_kQuantizeBlockwise(hip_bfloat16, 256, 2, 0, NF4) MAKE_kQuantizeBlockwise(hip_bfloat16, 128, 2, 0, NF4) //MAKE_kQuantizeBlockwise(hip_bfloat16, 64, 2, 0, NF4) template __global__ void kDequantizeBlockwise(float *code, unsigned char * A, float * absmax, half *out, const int blocksize, const int n); template __global__ void kDequantizeBlockwise(float *code, unsigned char * A, float * absmax, half *out, const int blocksize, const int n); template __global__ void kDequantizeBlockwise(float *code, unsigned char * A, float * absmax, half *out, const int blocksize, const int n); template __global__ void kDequantizeBlockwise(float *code, unsigned char * A, float * absmax, float *out, const int blocksize, const int n); template __global__ void kDequantizeBlockwise(float *code, unsigned char * A, float * absmax, float *out, const int blocksize, const int n); template __global__ void kDequantizeBlockwise(float *code, unsigned char * A, float * absmax, float *out, const int blocksize, const int n); template __global__ void kDequantizeBlockwise(float *code, unsigned char * A, float * absmax, hip_bfloat16 *out, const int blocksize, const int n); template __global__ void kDequantizeBlockwise(float *code, unsigned char * A, float * absmax, hip_bfloat16 *out, const int blocksize, const int n); template __global__ void kDequantizeBlockwise(float *code, unsigned char * A, float * absmax, hip_bfloat16 *out, const int blocksize, const int n); #define MAKE_OptimizerStatic8bit2StateBlockwise(oname, gtype, block_size, num_per_thread) \ template __global__ void kOptimizerStatic8bit2StateBlockwise(gtype* p, gtype* __restrict__ const g, unsigned char* state1, unsigned char* state2, \ const float beta1, const float beta2, const float beta3, const float alpha, \ const float eps, const int step, const float lr, \ float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, \ float* absmax1, float* absmax2, \ float weight_decay, \ const float gnorm_scale, const bool skip_zeros, const int n); \ MAKE_OptimizerStatic8bit2StateBlockwise(ADAM, float, 256, 1) MAKE_OptimizerStatic8bit2StateBlockwise(ADAM, half, 256, 1) MAKE_OptimizerStatic8bit2StateBlockwise(ADAM, hip_bfloat16, 256, 1) MAKE_OptimizerStatic8bit2StateBlockwise(ADEMAMIX, float, 256, 1) MAKE_OptimizerStatic8bit2StateBlockwise(ADEMAMIX, half, 256, 1) MAKE_OptimizerStatic8bit2StateBlockwise(ADEMAMIX, hip_bfloat16, 256, 1) #define MAKE_OptimizerStatic8bit1StateBlockwise(oname, gtype, block_size, num_per_thread) \ template __global__ void kOptimizerStatic8bit1StateBlockwise( \ gtype* p, gtype* __restrict__ const g, unsigned char* state1, \ const float beta1, const float beta2, \ const float eps, const int step, const float lr, \ float* __restrict__ const quantiles1, \ float* absmax1, \ float weight_decay, \ const float gnorm_scale, const bool skip_zeros, const int n); \ MAKE_OptimizerStatic8bit1StateBlockwise(MOMENTUM, float, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(MOMENTUM, half, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(MOMENTUM, hip_bfloat16, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(RMSPROP, float, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(RMSPROP, half, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(RMSPROP, hip_bfloat16, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(LION, float, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(LION, half, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(LION, hip_bfloat16, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(ADAGRAD, float, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(ADAGRAD, half, 256, 1) MAKE_OptimizerStatic8bit1StateBlockwise(ADAGRAD, hip_bfloat16, 256, 1) template __device__ void printnonzero(float *A, int num_values, const char*strval); template __device__ void printnonzero(half *A, int num_values, const char*strval); bitsandbytes/csrc/kernels_hip.cuh000066400000000000000000000160721506752665600175370ustar00rootroot00000000000000// !!! This is a file automatically generated by hipify!!! #include "hip/hip_runtime.h" // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include #include #ifndef kernels #define kernels __global__ void kQuantize(float* code, float* __restrict__ const A, unsigned char* out, const int n); __global__ void kDequantize(float* code, unsigned char* A, float* out, const int n); template __global__ void kQuantizeBlockwise( float* code, T* __restrict__ const A, float* absmax, unsigned char* out, float* __restrict__ const rand, const int rand_offset, const int n ); template __global__ void kDequantizeBlockwise(float* code, unsigned char* A, float* absmax, T* out, const int blocksize, const int n); template __global__ void kPreconditionOptimizer32bit2State( T* g, T* p, float* state1, float* state2, float* unorm, const float beta1, const float beta2, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const int n ); template __global__ void kOptimizer32bit2State( T* g, T* p, float* state1, float* state2, float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n ); template __global__ void kPreconditionOptimizer32bit1State( T* g, T* p, float* state1, float* unorm, const float beta1, const float beta2, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const int n ); template __global__ void kOptimizer32bit1State( T* g, T* p, float* state1, float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n ); template __global__ void kPreconditionOptimizerStatic8bit1State( T* p, T* __restrict__ const g, unsigned char* __restrict__ const state1, float* unorm, const float beta1, const float beta2, const float eps, const int step, float* __restrict__ const quantiles1, float* max1, float* new_max1, const float weight_decay, const float gnorm_scale, const int n ); template __global__ void kOptimizerStatic8bit1State( T* p, T* const g, unsigned char* state1, const float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* max1, float* new_max1, float weight_decay, const float gnorm_scale, const int n ); template __global__ void kPreconditionOptimizerStatic8bit2State( T* p, T* __restrict__ const g, unsigned char* __restrict__ const state1, unsigned char* __restrict__ const state2, float* unorm, const float beta1, const float beta2, const float eps, const int step, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, const float gnorm_scale, const int n ); template __global__ void kOptimizerStatic8bit2State( T* p, T* const g, unsigned char* state1, unsigned char* state2, const float* unorm, const float max_unorm, const float param_norm, const float beta1, const float beta2, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, float weight_decay, const float gnorm_scale, const int n ); template __global__ void kOptimizerStatic8bit2StateBlockwise( T* p, T* __restrict__ const g, unsigned char* state1, unsigned char* state2, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* __restrict__ const quantiles2, float* absmax1, float* absmax2, float weight_decay, const float gnorm_scale, const bool skip_zeros, const int n ); template __global__ void kOptimizerStatic8bit1StateBlockwise( T* p, T* __restrict__ const g, unsigned char* state1, const float beta1, const float beta2, const float eps, const int step, const float lr, float* __restrict__ const quantiles1, float* absmax1, float weight_decay, const float gnorm_scale, const bool skip_zeros, const int n ); template __global__ void kPercentileClipping(T* __restrict__ g, float* gnorm_vec, int step, const int n); template __global__ void kspmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, T* B, half* out, float* __restrict__ const dequant_stats, int nnz, int rowsA, int rowsB, int colsB ); template __global__ void kdequant_mm_int32_fp16( int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, half* out, half* __restrict__ const bias, const int numRows, const int numCols, const int n ); template __global__ void kgetRowStats(T* __restrict__ A, float* rowStats, float threshold, int rows, int cols); template __global__ void kInt8VectorQuant(T* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols); template __global__ void kTransformRowToFormat( char* __restrict__ const A, char* out, int rows, int cols, int tiledCols, int outRows, int outCols ); template __global__ void gemm_device(int M, int N, int K, T* __restrict__ const A, T* B, T* out, int lda, int ldb, int ldc); template __global__ void kgemm_4bit_inference( int M, int N, int K, T* __restrict__ const A, unsigned char* B, float* absmax, T* out, int lda, int ldb, int ldc, int blocksize ); template __global__ void kgemm_4bit_inference_naive( int M, int N, int K, T* __restrict__ const A, unsigned char* B, float* absmax, const float* datatype, T* out, int lda, int ldb, int ldc, int blocksize ); template __global__ void kfunc(T* A, T* B, T value, long n); #endif bitsandbytes/csrc/mps_kernels.metal000066400000000000000000000053541506752665600201020ustar00rootroot00000000000000#include using namespace metal; #define HLF_MAX 65504 #define TH 1024 #define NUM 4 #define NUM_BLOCK 4096 template static unsigned char quantize_scalar( float rand, device float* code, float x) { int pivot = 127; int upper_pivot = 255; int lower_pivot = 0; float lower = -1.0f; float upper = 1.0f; float val = code[pivot]; // i>>=1 = {32, 16, 8, 4, 2, 1} for(int i = 64; i > 0; i>>=1) { if(x > val) { lower_pivot = pivot; lower = val; pivot+=i; } else { upper_pivot = pivot; upper = val; pivot-=i; } val = code[pivot]; } if(upper_pivot == 255) upper = code[upper_pivot]; if(lower_pivot == 0) lower = code[lower_pivot]; if(!STOCHASTIC) { if(x > val) { float midpoint = (upper+val)*0.5f; if(x > midpoint) { return upper_pivot; } else return pivot; } else { float midpoint = (lower+val)*0.5f; if(x < midpoint) return lower_pivot; else return pivot; } } else { if(x > val) { float dist_to_upper = fabs(upper-x); float dist_full = upper-val; if(rand >= dist_to_upper/dist_full) return upper_pivot; else return pivot; } else { float dist_to_lower = fabs(lower-x); float dist_full = val-lower; if(rand >= dist_to_lower/dist_full) return lower_pivot; else return pivot; } } } kernel void quantize(device float* code [[buffer(0)]], device float* A [[buffer(1)]], device uchar* out [[buffer(2)]], constant uint& n [[buffer(3)]], uint id [[thread_position_in_grid]]) { const uint n_full = (NUM_BLOCK * (n / NUM_BLOCK)) + (n % NUM_BLOCK == 0 ? 0 : NUM_BLOCK); uint valid_items = (id / NUM_BLOCK + 1 == (n + NUM_BLOCK - 1) / NUM_BLOCK) ? n - (id / NUM_BLOCK * NUM_BLOCK) : NUM_BLOCK; const uint base_idx = (id / NUM_BLOCK * NUM_BLOCK); float vals[NUM]; uchar qvals[NUM]; for (uint i = base_idx; i < n_full; i += ((n + NUM_BLOCK - 1) / NUM_BLOCK) * NUM_BLOCK) { valid_items = n - i > NUM_BLOCK ? NUM_BLOCK : n - i; threadgroup_barrier(mem_flags::mem_threadgroup); for (uint j = 0; j < valid_items; j++) { vals[j] = A[i + j]; } for (uint j = 0; j < valid_items; j++) { qvals[j] = quantize_scalar(0.0f, code, vals[j]); } threadgroup_barrier(mem_flags::mem_threadgroup); for (uint j = 0; j < valid_items; j++) { out[i + j] = qvals[j]; } } } bitsandbytes/csrc/mps_ops.h000066400000000000000000000000001506752665600163440ustar00rootroot00000000000000bitsandbytes/csrc/mps_ops.mm000066400000000000000000000033021506752665600165360ustar00rootroot00000000000000#import #define HLF_MAX 65504 #define TH 1024 #define NUM 4 #define NUM_BLOCK 4096 static inline MPSGraph* get_graph() { static MPSGraph* cur = nil; if (!cur) { cur = [[MPSGraph alloc] init]; } return cur; } static inline id get_device() { NSError* error = nil; static id device = nil; if (!device) { device = MTLCreateSystemDefaultDevice(); } if (!device) { NSLog(@"Failed to get MPS device"); abort(); } return device; } static inline id get_library() { NSError* error = nil; static id library = nil; if (!library) { library = [get_device() newLibraryWithURL:[NSURL fileURLWithPath:@"bitsandbytes.metallib"] error:&error]; } if (!library) { NSLog(@"Failed to load bitsandbytes.metallib"); abort(); } return library; } /*MPSGraphTensor* dequantize_mps(MPSGraphTensor* code, MPSGraphTensor* A, int n) { id out = [get_graph() dequantizeTensor:(MPSGraphTensor*)A scaleTensor:(MPSGraphTensor*)code zeroPoint:0.0 dataType:MPSDataTypeInt8 axis:0 name:@"out"]; return out; }*/ // MPSGraph function for quantize extern "C" MPSGraphTensor* quantize_mps(MPSGraph* graph, MPSGraphTensor* code, MPSGraphTensor* A, int n) { id device = get_device(); id library = get_library(); static id kernel = nil; if (!kernel) { kernel = [library newFunctionWithName:@"quantize"]; if (!kernel) { NSLog(@"Failed to load bitsandbytes.metallib"); abort(); } } NSLog(@"Not implemented"); return nil; } bitsandbytes/csrc/ops.cu000066400000000000000000001004321506752665600156570ustar00rootroot00000000000000// Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include #include #include #include #include #include #include #define ERR_NOT_IMPLEMENTED 100 using namespace BinSearch; using std::cout; using std::endl; void quantize(float* code, float* A, unsigned char* out, int n) { int num_blocks = n / 1024; num_blocks = n % 1024 == 0 ? num_blocks : num_blocks + 1; kQuantize<<>>(code, A, out, n); CUDA_CHECK_RETURN(cudaPeekAtLastError()); } void dequantize(float* code, unsigned char* A, float* out, int n, cudaStream_t stream) { int num_blocks = n / 1024; num_blocks = n % 1024 == 0 ? num_blocks : num_blocks + 1; kDequantize<<>>(code, A, out, n); CUDA_CHECK_RETURN(cudaPeekAtLastError()); } template void quantizeBlockwise( float* code, T* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ) { int num_blocks = n / blocksize; num_blocks = n % blocksize == 0 ? num_blocks : num_blocks + 1; if (blocksize == 4096) kQuantizeBlockwise <<>>(code, A, absmax, out, rand, rand_offset, n); else if (blocksize == 2048) kQuantizeBlockwise<<>>(code, A, absmax, out, rand, rand_offset, n); else if (blocksize == 1024) kQuantizeBlockwise<<>>(code, A, absmax, out, rand, rand_offset, n); else if (blocksize == 512) kQuantizeBlockwise<<>>(code, A, absmax, out, rand, rand_offset, n); else if (blocksize == 256) kQuantizeBlockwise<<>>(code, A, absmax, out, rand, rand_offset, n); else if (blocksize == 128) kQuantizeBlockwise<<>>(code, A, absmax, out, rand, rand_offset, n); else if (blocksize == 64) kQuantizeBlockwise<<>>(code, A, absmax, out, rand, rand_offset, n); CUDA_CHECK_RETURN(cudaPeekAtLastError()); } template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, T* out, int blocksize, const int n, cudaStream_t stream ) { // printf("stream==%d\n",stream); int num_blocks = n / blocksize; num_blocks = n % blocksize == 0 ? num_blocks : num_blocks + 1; int tile_size = (DATA_TYPE > 0) ? 1024 : 512; if (DATA_TYPE > 0) kDequantizeBlockwise <<<(n + tile_size - 1) / tile_size, 64, 0, stream>>>(code, A, absmax, out, blocksize / 2, n); else kDequantizeBlockwise <<<(n + tile_size - 1) / tile_size, 64, 0, stream>>>(code, A, absmax, out, blocksize, n); CUDA_CHECK_RETURN(cudaPeekAtLastError()); } template void optimizer32bit( T* g, T* p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, bool skip_zeros, const int n ) { int num_blocks = n / 4096; num_blocks = n % 4096 == 0 ? num_blocks : num_blocks + 1; switch (OPTIMIZER) { case ADAM: case ADEMAMIX: if (max_unorm > 0.0f) { CUDA_CHECK_RETURN(cudaMemset(unorm, 0, 1 * sizeof(float))); kPreconditionOptimizer32bit2State<<>>( g, p, state1, state2, unorm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, n ); CUDA_CHECK_RETURN(cudaPeekAtLastError()); } kOptimizer32bit2State<<>>( g, p, state1, state2, unorm, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, lr, gnorm_scale, skip_zeros, n ); CUDA_CHECK_RETURN(cudaPeekAtLastError()); break; case MOMENTUM: case RMSPROP: case ADAGRAD: if (max_unorm > 0.0f) { CUDA_CHECK_RETURN(cudaMemset(unorm, 0, 1 * sizeof(float))); kPreconditionOptimizer32bit1State <<>>(g, p, state1, unorm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, n); CUDA_CHECK_RETURN(cudaPeekAtLastError()); } kOptimizer32bit1State<<>>( g, p, state1, unorm, max_unorm, param_norm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, skip_zeros, n ); CUDA_CHECK_RETURN(cudaPeekAtLastError()); break; case LION: // in lion, the momentum update after the parameter update kOptimizer32bit1State<<>>( g, p, state1, unorm, max_unorm, param_norm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, skip_zeros, n ); CUDA_CHECK_RETURN(cudaPeekAtLastError()); if (max_unorm > 0.0f) { CUDA_CHECK_RETURN(cudaMemset(unorm, 0, 1 * sizeof(float))); kPreconditionOptimizer32bit1State <<>>(g, p, state1, unorm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, n); CUDA_CHECK_RETURN(cudaPeekAtLastError()); } break; } } template void optimizerStatic8bit( T* p, T* g, unsigned char* state1, unsigned char* state2, float* unorm, float max_unorm, float param_norm, float beta1, float beta2, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, float weight_decay, const float gnorm_scale, int n ) { int num_blocks = n / 4096; num_blocks = n % 4096 == 0 ? num_blocks : num_blocks + 1; if (max_unorm > 0.0f) { CUDA_CHECK_RETURN(cudaMemset(unorm, 0, 1 * sizeof(float))); } switch (OPTIMIZER) { case ADAM: CUDA_CHECK_RETURN(cudaMemset(new_max1, 0, 1 * sizeof(float))); CUDA_CHECK_RETURN(cudaMemset(new_max2, 0, 1 * sizeof(float))); kPreconditionOptimizerStatic8bit2State<<>>( p, g, state1, state2, unorm, beta1, beta2, eps, step, quantiles1, quantiles2, max1, max2, new_max1, new_max2, gnorm_scale, n ); CUDA_CHECK_RETURN(cudaPeekAtLastError()); kOptimizerStatic8bit2State<<>>( p, g, state1, state2, unorm, max_unorm, param_norm, beta1, beta2, eps, step, lr, quantiles1, quantiles2, max1, max2, new_max1, new_max2, weight_decay, gnorm_scale, n ); CUDA_CHECK_RETURN(cudaPeekAtLastError()); break; case MOMENTUM: case RMSPROP: case ADAGRAD: CUDA_CHECK_RETURN(cudaMemset(new_max1, 0, 1 * sizeof(float))); kPreconditionOptimizerStatic8bit1State<<>>( p, g, state1, unorm, beta1, beta2, eps, step, quantiles1, max1, new_max1, weight_decay, gnorm_scale, n ); CUDA_CHECK_RETURN(cudaPeekAtLastError()); kOptimizerStatic8bit1State<<>>( p, g, state1, unorm, max_unorm, param_norm, beta1, beta2, eps, step, lr, quantiles1, max1, new_max1, weight_decay, gnorm_scale, n ); CUDA_CHECK_RETURN(cudaPeekAtLastError()); break; case LION: // in lion, the momentum update happens after the parameter update kOptimizerStatic8bit1State<<>>( p, g, state1, unorm, max_unorm, param_norm, beta1, beta2, eps, step, lr, quantiles1, max1, new_max1, weight_decay, gnorm_scale, n ); CUDA_CHECK_RETURN(cudaPeekAtLastError()); CUDA_CHECK_RETURN(cudaMemset(new_max1, 0, 1 * sizeof(float))); kPreconditionOptimizerStatic8bit1State<<>>( p, g, state1, unorm, beta1, beta2, eps, step, quantiles1, max1, new_max1, weight_decay, gnorm_scale, n ); CUDA_CHECK_RETURN(cudaPeekAtLastError()); break; default: break; } } #define BLOCKSIZE_2STATE 256 #define NUM_2STATE 1 #define BLOCKSIZE_1STATE 256 #define NUM_1STATE 1 template void optimizerStatic8bitBlockwise( T* p, T* g, unsigned char* state1, unsigned char* state2, float beta1, float beta2, float beta3, float alpha, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* absmax1, float* absmax2, float weight_decay, const float gnorm_scale, bool skip_zeros, int n ) { int num_blocks = 0; switch (OPTIMIZER) { case ADAM: case ADEMAMIX: num_blocks = n / BLOCKSIZE_2STATE; num_blocks = n % BLOCKSIZE_2STATE == 0 ? num_blocks : num_blocks + 1; kOptimizerStatic8bit2StateBlockwise <<>>( p, g, state1, state2, beta1, beta2, beta3, alpha, eps, step, lr, quantiles1, quantiles2, absmax1, absmax2, weight_decay, gnorm_scale, skip_zeros, n ); CUDA_CHECK_RETURN(cudaPeekAtLastError()); break; case MOMENTUM: case RMSPROP: case ADAGRAD: case LION: num_blocks = n / BLOCKSIZE_1STATE; num_blocks = n % BLOCKSIZE_1STATE == 0 ? num_blocks : num_blocks + 1; kOptimizerStatic8bit1StateBlockwise <<>>( p, g, state1, beta1, beta2, eps, step, lr, quantiles1, absmax1, weight_decay, gnorm_scale, skip_zeros, n ); CUDA_CHECK_RETURN(cudaPeekAtLastError()); break; } } template void percentileClipping(T* g, float* gnorm_vec, int step, const int n) { int num_blocks = n / 2048; num_blocks = n % 2048 == 0 ? num_blocks : num_blocks + 1; CUDA_CHECK_RETURN(cudaMemset(&gnorm_vec[step % 100], 0, 1 * sizeof(float))); kPercentileClipping<<>>(g, gnorm_vec, step, n); CUDA_CHECK_RETURN(cudaPeekAtLastError()); } void gemmex( Context* context, bool transposeA, bool transposeB, int m, int n, int k, void* A, void* B, void* C, int lda, int ldb, int ldc ) { const int falpha = 1; const int fbeta = 0; const void* alpha = &falpha; const void* beta = &fbeta; cublasStatus_t status; status = cublasGemmEx( context->m_handle, transposeA ? CUBLAS_OP_T : CUBLAS_OP_N, transposeB ? CUBLAS_OP_T : CUBLAS_OP_N, m, n, k, alpha, A, CUDA_R_8I, lda, B, CUDA_R_8I, ldb, beta, C, CUDA_R_32I, ldc, CUDA_R_32I, CUBLAS_GEMM_DEFAULT_TENSOR_OP ); if (status != CUBLAS_STATUS_SUCCESS) { std::cout << "CUBLAS ERROR: Status " << status << std::endl; } } void strided_gemmex( Context* context, bool transposeA, bool transposeB, int m, int n, int k, void* A, void* B, void* C, int lda, int ldb, int ldc, long long int strideA, long long int strideB, long long int strideC, int batchCount ) { const int falpha = 1; const int fbeta = 0; const void* alpha = &falpha; const void* beta = &fbeta; cublasStatus_t status; // cout << transposeA << transposeB << endl; // printf("%i %i %i\n", m,n,k); // printf("%i %i %i\n", lda,ldb,ldc); // printf("%i %i %i\n", strideA, strideB, strideC); // printf("%i\n", batchCount); status = cublasGemmStridedBatchedEx( context->m_handle, transposeA ? CUBLAS_OP_T : CUBLAS_OP_N, transposeB ? CUBLAS_OP_T : CUBLAS_OP_N, m, n, k, alpha, A, CUDA_R_8I, lda, (long long int)strideA, B, CUDA_R_8I, ldb, (long long int)strideB, beta, C, CUDA_R_32I, ldc, (long long int)strideC, batchCount, CUDA_R_32I, CUBLAS_GEMM_DEFAULT ); if (status != CUBLAS_STATUS_SUCCESS) { std::cout << "CUBLAS ERROR: Status " << status << std::endl; } } int roundoff(int v, int d) { return (v + d - 1) / d * d; } template cublasLtOrder_t get_order() { switch (ORDER) { case ROW: return CUBLASLT_ORDER_ROW; break; case COL: return CUBLASLT_ORDER_COL; break; case COL32: return CUBLASLT_ORDER_COL32; break; case COL_TURING: return CUBLASLT_ORDER_COL4_4R2_8C; break; case COL_AMPERE: return CUBLASLT_ORDER_COL32_2R_4R4; break; default: break; } return CUBLASLT_ORDER_ROW; } template cublasLtOrder_t get_order(); template cublasLtOrder_t get_order(); template cublasLtOrder_t get_order(); template cublasLtOrder_t get_order(); template cublasLtOrder_t get_order(); template int get_leading_dim(int dim1, int dim2) { switch (ORDER) { case ROW: return dim2; break; case COL: return dim1; break; case COL32: // 32*row tiles return dim1 * 32; break; case COL_TURING: return 32 * roundoff(dim1, 8); break; case COL_AMPERE: // 32*32 tiles return 32 * roundoff(dim1, 32); break; default: return 0; break; } } template int igemmlt( cublasLtHandle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, cudaStream_t stream ) { // Calculate C = A^T @ B, in col-major layout. // // Use the IMMA kernels requires: // * A must be transposed and B must be non-transposed. // * Dimensions m and k must be multiples of 4. // * All pointers must be 4-byte aligned; 16-byte alignment preferred. int has_error = 0; cublasLtMatmulDesc_t matmulDesc; cublasLtMatrixLayout_t aDesc, bDesc, cDesc; cublasOperation_t opT = CUBLAS_OP_T; cudaDataType_t outType = DTYPE_OUT == 32 ? CUDA_R_32I : CUDA_R_8I; cudaDataType_t scaleType = DTYPE_OUT == 32 ? CUDA_R_32I : CUDA_R_32F; cublasLtPointerMode_t pointerMode = CUBLASLT_POINTER_MODE_ALPHA_DEVICE_VECTOR_BETA_ZERO; has_error |= checkCublasStatus(cublasLtMatrixLayoutCreate(&aDesc, CUDA_R_8I, m, k, lda)); has_error |= checkCublasStatus(cublasLtMatrixLayoutCreate(&bDesc, CUDA_R_8I, m, n, ldb)); has_error |= checkCublasStatus(cublasLtMatrixLayoutCreate(&cDesc, outType, k, n, ldc)); // Default layout order is col major has_error |= checkCublasStatus(cublasLtMatmulDescCreate(&matmulDesc, CUBLAS_COMPUTE_32I, scaleType)); has_error |= checkCublasStatus(cublasLtMatmulDescSetAttribute(matmulDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opT, sizeof(opT))); if (DTYPE_OUT == 32) { int alpha = 1, beta = 0; has_error |= checkCublasStatus(cublasLtMatmul( ltHandle, matmulDesc, &alpha, A, aDesc, B, bDesc, &beta, (int32_t*)C, cDesc, (int32_t*)C, cDesc, NULL, NULL, 0, stream )); } else { // This path is unlikely to be used, as 8-bit accumulation can lead to likely overflows. if (!SCALE_ROWS) { float alpha = 1.0f, beta = 0.0f; has_error |= checkCublasStatus(cublasLtMatmul( ltHandle, matmulDesc, &alpha, A, aDesc, B, bDesc, &beta, (int8_t*)C, cDesc, (int8_t*)C, cDesc, NULL, NULL, 0, stream )); } else { cublasLtPointerMode_t alphaVec = CUBLASLT_POINTER_MODE_ALPHA_DEVICE_VECTOR_BETA_HOST; float beta = 0.0f; has_error |= checkCublasStatus(cublasLtMatmulDescSetAttribute( matmulDesc, CUBLASLT_MATMUL_DESC_POINTER_MODE, &pointerMode, sizeof(alphaVec) )); has_error |= checkCublasStatus(cublasLtMatmul( ltHandle, matmulDesc, row_scale, A, aDesc, B, bDesc, &beta, (int8_t*)C, cDesc, (int8_t*)C, cDesc, NULL, NULL, 0, stream )); } } has_error |= checkCublasStatus(cublasLtMatrixLayoutDestroy(cDesc)); has_error |= checkCublasStatus(cublasLtMatrixLayoutDestroy(bDesc)); has_error |= checkCublasStatus(cublasLtMatrixLayoutDestroy(aDesc)); has_error |= checkCublasStatus(cublasLtMatmulDescDestroy(matmulDesc)); if (has_error == 1) printf("error detected"); return has_error; } int fill_up_to_nearest_multiple(int value, int multiple) { return value + (value % multiple == 0 ? 0 : (multiple - (value % multiple))); } void dequant_mm_int32_fp16( int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, cudaStream_t stream ) { const int threads = 512; const int num_per_thread = 4; const int num_per_block = threads * num_per_thread; const int n = numRows * numCols; const int num_blocks = (n + num_per_block - 1) / num_per_block; kdequant_mm_int32_fp16 <<>>(A, rowStats, colStats, out, bias, numRows, numCols, n); CUDA_CHECK_RETURN(cudaPeekAtLastError()); } void int8VectorQuant( half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, cudaStream_t stream ) { if (threshold == 0.0) { kInt8VectorQuant<<>>(A, out, rowStats, threshold, rows, cols); } else { kInt8VectorQuant<<>>(A, out, rowStats, threshold, rows, cols); } CUDA_CHECK_RETURN(cudaPeekAtLastError()); } void getRowStats(half* A, float* rowStats, float threshold, int rows, int cols, cudaStream_t stream) { if (threshold == 0.0) kgetRowStats<<>>(A, rowStats, threshold, rows, cols); else kgetRowStats<<>>(A, rowStats, threshold, rows, cols); CUDA_CHECK_RETURN(cudaPeekAtLastError()); } void spmm_coo( cusparseHandle_t handle, int* A_rowidx, int* A_colidx, half* A_vals, int A_nnz, int A_rows, int A_cols, int B_cols, int ldb, half* B, int ldc, half* C, bool transposed_B ) { cusparseSpMatDescr_t descA; cusparseDnMatDescr_t descB, descC; float alpha = 1.0f; float beta = 0.0f; void* dBuffer = NULL; size_t bufferSize = 0; CHECK_CUSPARSE(cusparseCreateCoo( &descA, A_rows, A_cols, A_nnz, A_rowidx, A_colidx, A_vals, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, CUDA_R_16F )); // Create dense matrix C CHECK_CUSPARSE(cusparseCreateDnMat(&descC, A_rows, B_cols, ldc, C, CUDA_R_16F, CUSPARSE_ORDER_ROW)); // Create dense matrix B if (transposed_B) { int tmp = A_cols; A_cols = B_cols; B_cols = tmp; } CHECK_CUSPARSE(cusparseCreateDnMat(&descB, A_cols, B_cols, ldb, B, CUDA_R_16F, CUSPARSE_ORDER_ROW)); // allocate an external buffer if needed CHECK_CUSPARSE(cusparseSpMM_bufferSize( handle, CUSPARSE_OPERATION_NON_TRANSPOSE, transposed_B ? CUSPARSE_OPERATION_TRANSPOSE : CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, descA, descB, &beta, descC, CUDA_R_32F, CUSPARSE_SPMM_ALG_DEFAULT, &bufferSize )); CUDA_CHECK_RETURN(cudaMalloc(&dBuffer, bufferSize)); // execute SpMM CHECK_CUSPARSE(cusparseSpMM( handle, CUSPARSE_OPERATION_NON_TRANSPOSE, transposed_B ? CUSPARSE_OPERATION_TRANSPOSE : CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, descA, descB, &beta, descC, CUDA_R_32F, CUSPARSE_SPMM_ALG_DEFAULT, dBuffer )); // destroy matrix/vector descriptors CHECK_CUSPARSE(cusparseDestroySpMat(descA)); CHECK_CUSPARSE(cusparseDestroyDnMat(descB)); CHECK_CUSPARSE(cusparseDestroyDnMat(descC)); CUDA_CHECK_RETURN(cudaFree(dBuffer)); } template void spmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, T* B, half* out, float* dequant_stats, int nnz_rows, int nnz, int rowsA, int rowsB, int colsB ) { kspmm_coo_very_sparse_naive<<>>( max_count, max_idx, offset_rowidx, rowidx, colidx, values, B, out, dequant_stats, nnz, rowsA, rowsB, colsB ); CUDA_CHECK_RETURN(cudaPeekAtLastError()); } template void gemm_host(int m, int n, int k, T* A, T* B, T* out, int lda, int ldb, int ldc, int bits) { int num_blocks = (m + 31) / 32; if (bits == 32) gemm_device<<>>(m, n, k, A, B, out, lda, ldb, ldc); if (bits == 16) gemm_device<<>>(m, n, k, A, B, out, lda, ldb, ldc); } template void gemm_4bit_inference( int m, int n, int k, T* A, unsigned char* B, float* absmax, T* out, int lda, int ldb, int ldc, int blocksize ) { int num_blocks = (m + 31) / 32; kgemm_4bit_inference<<>>(m, n, k, A, B, absmax, out, lda, ldb, ldc, blocksize); } template void gemm_4bit_inference_naive( int m, int n, int k, T* A, unsigned char* B, float* absmax, float* datatype, T* out, int lda, int ldb, int ldc, int blocksize, cudaStream_t stream ) { int num_blocks = (m + 3) / 4; kgemm_4bit_inference_naive <<>>(m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize); CUDA_CHECK_RETURN(cudaPeekAtLastError()); } template void func(T* A, T* B, T value, long n) { int threads = 512; int blocks = n / threads; blocks = n % threads == 0 ? blocks : blocks + 1; blocks = blocks > 65535 ? 65535 : blocks; kfunc<<>>(A, B, value, n); CUDA_CHECK_RETURN(cudaPeekAtLastError()); } //============================================================== // TEMPLATE DEFINITIONS //============================================================== template void func(float* A, float* B, float value, long n); template void func(unsigned char* A, unsigned char* B, unsigned char value, long n); template void func(float* A, float* B, float value, long n); template void func(float* A, float* B, float value, long n); template void gemm_4bit_inference( int m, int n, int k, half* A, unsigned char* B, float* absmax, half* out, int lda, int ldb, int ldc, int blocksize ); template void gemm_4bit_inference_naive( int m, int n, int k, half* A, unsigned char* B, float* absmax, float* datatype, half* out, int lda, int ldb, int ldc, int blocksize, cudaStream_t stream ); template void gemm_4bit_inference_naive<__nv_bfloat16, 16>( int m, int n, int k, __nv_bfloat16* A, unsigned char* B, float* absmax, float* datatype, __nv_bfloat16* out, int lda, int ldb, int ldc, int blocksize, cudaStream_t stream ); template void gemm_4bit_inference_naive( int m, int n, int k, float* A, unsigned char* B, float* absmax, float* datatype, float* out, int lda, int ldb, int ldc, int blocksize, cudaStream_t stream ); // template void gemm_host(int m, int n, int k, float * A, float* B, float * out, int lda, int ldb, int ldc, // int bits); template void gemm_host(int m, int n, int k, half* A, half* B, half* out, int lda, int ldb, int ldc, int bits); template void spmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, half* B, half* out, float* dequant_stats, int nnz_rows, int nnz, int rowsA, int rowsB, int colsB ); template void spmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, signed char* B, half* out, float* dequant_stats, int nnz_rows, int nnz, int rowsA, int rowsB, int colsB ); template int igemmlt<32, 0>( cublasLtHandle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, cudaStream_t stream ); template int igemmlt<8, 0>( cublasLtHandle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, cudaStream_t stream ); template int igemmlt<8, 1>( cublasLtHandle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, cudaStream_t stream ); template void quantizeBlockwise( float* code, half* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void quantizeBlockwise( float* code, half* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void quantizeBlockwise( float* code, half* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void quantizeBlockwise( float* code, half* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void quantizeBlockwise( float* code, float* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void quantizeBlockwise( float* code, float* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void quantizeBlockwise( float* code, float* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void quantizeBlockwise( float* code, float* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void quantizeBlockwise<__nv_bfloat16, 1, General8bit>( float* code, __nv_bfloat16* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void quantizeBlockwise<__nv_bfloat16, 0, General8bit>( float* code, __nv_bfloat16* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void quantizeBlockwise<__nv_bfloat16, 0, FP4>( float* code, __nv_bfloat16* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void quantizeBlockwise<__nv_bfloat16, 0, NF4>( float* code, __nv_bfloat16* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, cudaStream_t stream ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, cudaStream_t stream ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, cudaStream_t stream ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, half* out, int blocksize, const int n, cudaStream_t stream ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, half* out, int blocksize, const int n, cudaStream_t stream ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, half* out, int blocksize, const int n, cudaStream_t stream ); template void dequantizeBlockwise<__nv_bfloat16, General8bit>( float* code, unsigned char* A, float* absmax, __nv_bfloat16* out, int blocksize, const int n, cudaStream_t stream ); template void dequantizeBlockwise<__nv_bfloat16, FP4>( float* code, unsigned char* A, float* absmax, __nv_bfloat16* out, int blocksize, const int n, cudaStream_t stream ); template void dequantizeBlockwise<__nv_bfloat16, NF4>( float* code, unsigned char* A, float* absmax, __nv_bfloat16* out, int blocksize, const int n, cudaStream_t stream ); #define MAKE_optimizer32bit(name, gtype) \ template void optimizer32bit( \ gtype * g, gtype * p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, \ const float beta1, const float beta2, const float beta3, const float alpha, const float eps, \ const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, \ const int n \ ); MAKE_optimizer32bit(ADAM, half) MAKE_optimizer32bit(ADAM, float) MAKE_optimizer32bit(ADAM, __nv_bfloat16) MAKE_optimizer32bit(MOMENTUM, half) MAKE_optimizer32bit(MOMENTUM, float) MAKE_optimizer32bit( MOMENTUM, __nv_bfloat16 ) MAKE_optimizer32bit(RMSPROP, half) MAKE_optimizer32bit(RMSPROP, float) MAKE_optimizer32bit(RMSPROP, __nv_bfloat16) MAKE_optimizer32bit(LION, half) MAKE_optimizer32bit(LION, float) MAKE_optimizer32bit(LION, __nv_bfloat16) MAKE_optimizer32bit(ADAGRAD, half) MAKE_optimizer32bit(ADAGRAD, float) MAKE_optimizer32bit(ADAGRAD, __nv_bfloat16) MAKE_optimizer32bit(ADEMAMIX, half) MAKE_optimizer32bit(ADEMAMIX, __nv_bfloat16) MAKE_optimizer32bit(ADEMAMIX, float) #define MAKE_optimizerStatic8bit(name, gtype) \ template void optimizerStatic8bit( \ gtype * p, gtype * g, unsigned char* state1, unsigned char* state2, float* unorm, float max_unorm, \ float param_norm, float beta1, float beta2, float eps, int step, float lr, float* quantiles1, \ float* quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, float weight_decay, \ const float gnorm_scale, int n \ ); MAKE_optimizerStatic8bit(ADAM, half) MAKE_optimizerStatic8bit(ADAM, float) MAKE_optimizerStatic8bit(MOMENTUM, half) MAKE_optimizerStatic8bit(MOMENTUM, float) MAKE_optimizerStatic8bit( RMSPROP, half ) MAKE_optimizerStatic8bit(RMSPROP, float) MAKE_optimizerStatic8bit(LION, half) MAKE_optimizerStatic8bit(LION, float) MAKE_optimizerStatic8bit(ADAGRAD, half) MAKE_optimizerStatic8bit(ADAGRAD, float) #define MAKE_optimizerStatic8bitBlockwise(gtype, optim_name) \ template void optimizerStatic8bitBlockwise( \ gtype * p, gtype * g, unsigned char* state1, unsigned char* state2, float beta1, float beta2, float beta3, \ float alpha, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* absmax1, \ float* absmax2, float weight_decay, const float gnorm_scale, bool skip_zeros, int n \ ); MAKE_optimizerStatic8bitBlockwise(half, ADAM); MAKE_optimizerStatic8bitBlockwise(float, ADAM); MAKE_optimizerStatic8bitBlockwise(__nv_bfloat16, ADAM); MAKE_optimizerStatic8bitBlockwise(half, MOMENTUM); MAKE_optimizerStatic8bitBlockwise(float, MOMENTUM); MAKE_optimizerStatic8bitBlockwise(__nv_bfloat16, MOMENTUM); MAKE_optimizerStatic8bitBlockwise(half, RMSPROP); MAKE_optimizerStatic8bitBlockwise(float, RMSPROP); MAKE_optimizerStatic8bitBlockwise(__nv_bfloat16, RMSPROP); MAKE_optimizerStatic8bitBlockwise(half, LION); MAKE_optimizerStatic8bitBlockwise(float, LION); MAKE_optimizerStatic8bitBlockwise(__nv_bfloat16, LION); MAKE_optimizerStatic8bitBlockwise(half, ADAGRAD); MAKE_optimizerStatic8bitBlockwise(float, ADAGRAD); MAKE_optimizerStatic8bitBlockwise(__nv_bfloat16, ADAGRAD); MAKE_optimizerStatic8bitBlockwise(half, ADEMAMIX); MAKE_optimizerStatic8bitBlockwise(__nv_bfloat16, ADEMAMIX); MAKE_optimizerStatic8bitBlockwise(float, ADEMAMIX); template void percentileClipping(float* g, float* gnorm_vec, int step, const int n); template void percentileClipping(half* g, float* gnorm_vec, int step, const int n); template int get_leading_dim(int dim1, int dim2); template int get_leading_dim(int dim1, int dim2); template int get_leading_dim(int dim1, int dim2); bitsandbytes/csrc/ops.cuh000066400000000000000000000173641506752665600160420ustar00rootroot00000000000000// Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #ifndef ops_H #define ops_H #include #include #include #include #include #include #include #include #include #include #include #define CUDA_CHECK_RETURN(value) \ { \ cudaError_t _m_cudaStat = value; \ if (_m_cudaStat != cudaSuccess) { \ fprintf(stderr, "Error %s at line %d in file %s\n", cudaGetErrorString(_m_cudaStat), __LINE__, __FILE__); \ exit(1); \ } \ } #define CHECK_CUSPARSE(value) \ { \ cusparseStatus_t _m_cudaStat = value; \ if (_m_cudaStat != CUSPARSE_STATUS_SUCCESS) { \ fprintf( \ stderr, "Error %s at line %d in file %s\n", cusparseGetErrorString(_m_cudaStat), __LINE__, __FILE__ \ ); \ exit(1); \ } \ } inline void checkCudaStatus(cudaError_t status) { if (status != cudaSuccess) { printf("cuda API failed with status %d: %s\n", status, cudaGetErrorString(status)); throw std::logic_error("cuda API failed"); } } inline int checkCublasStatus(cublasStatus_t status) { if (status != CUBLAS_STATUS_SUCCESS) { printf("cuBLAS API failed with status %d\n", status); // throw std::logic_error("cuBLAS API failed"); return 1; } return 0; } typedef enum Operations_t { ksmul = 0, } Operations_t; typedef enum Optimizer_t { ADAM = 0, MOMENTUM = 1, RMSPROP = 2, LARS = 3, ADAGRAD = 4, LION = 5, ADEMAMIX = 6 } Optimizer_t; typedef enum Transform_t { ROW = 0, COL = 1, COL32 = 2, COL_TURING = 3, COL_AMPERE = 4, } Transform_t; typedef enum DataType_t { General8bit = 0, FP4 = 1, NF4 = 2, } DataType_t; typedef enum Funcs_t { FILL = 0, ARANGE = 1, _MUL = 2, } Funcs_t; class Context { public: cublasHandle_t m_handle; Context() { cublasHandle_t handle; cublasCreate_v2(&handle); m_handle = handle; } }; class ContextLt { public: cublasLtHandle_t m_handle; ContextLt() { cublasLtHandle_t handle; cublasLtCreate(&handle); m_handle = handle; } }; class ContextCusparse { public: cusparseHandle_t m_handle; ContextCusparse() { cusparseHandle_t handle; cusparseCreate(&handle); m_handle = handle; } }; void quantize(float* code, float* A, unsigned char* out, int n); void dequantize(float* code, unsigned char* A, float* out, int n, cudaStream_t stream); template void quantizeBlockwise( float* code, T* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, T* out, int block_size, const int n, cudaStream_t stream ); template void optimizer32bit( T* g, T* p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, float beta1, float beta2, float beta3, float alpha, float eps, float weight_decay, int step, float lr, const float gnorm_scale, bool skip_zeros, int n ); template void optimizerStatic8bit( T* p, T* g, unsigned char* state1, unsigned char* state2, float* unorm, float max_unorm, float param_norm, float beta1, float beta2, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, float weight_decay, const float gnorm_scale, int n ); template void optimizerStatic8bitBlockwise( T* p, T* g, unsigned char* state1, unsigned char* state2, float beta1, float beta2, float beta3, float alpha, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* absmax1, float* absmax2, float weight_decay, const float gnorm_scale, bool skip_zeros, int n ); template void percentileClipping(T* g, float* gnorm_vec, int step, const int n); void gemmex( Context* context, bool transposeA, bool transposeB, int m, int n, int k, void* A, void* B, void* C, int lda, int ldb, int ldc ); void strided_gemmex( Context* context, bool transposeA, bool transposeB, int m, int n, int k, void* A, void* B, void* C, int lda, int ldb, int ldc, long long int strideA, long long int strideB, long long int strideC, int batchCount ); template int igemmlt( cublasLtHandle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, cudaStream_t stream ); void cutlass_igemm( bool transposeA, bool transposeB, int m, int n, int k, void* A, void* B, void* C, int lda, int ldb, int ldc ); void dequant_mm_int32_fp16( int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, cudaStream_t stream ); void getRowStats(half* A, float* rowStats, float threshold, int rows, int cols, cudaStream_t stream); void int8VectorQuant( half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, cudaStream_t stream ); void spmm_coo( cusparseHandle_t handle, int* A_rowidx, int* A_colidx, half* A_vals, int A_nnz, int A_rows, int A_cols, int B_cols, int ldb, half* B, int ldc, half* C, bool transposed_B ); template void spmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, T* B, half* out, float* dequant_stats, int nnz_rows, int nnz, int rowsA, int rowsB, int colsB ); void matmul4bite(half* A, unsigned char* B, half* out, int lda, int ldb, int rowsA, int colsA, int colsB); template void gemm_host(int m, int n, int k, T* A, T* B, T* out, int lda, int ldb, int ldc, int bits); template void gemm_4bit_inference( int m, int n, int k, T* A, unsigned char* B, float* absmax, T* out, int lda, int ldb, int ldc, int blocksize ); template void gemm_4bit_inference_naive( int m, int n, int k, T* A, unsigned char* B, float* absmax, float* datatype, T* out, int lda, int ldb, int ldc, int blocksize, cudaStream_t stream ); template void func(T* A, T* B, T value, long n); #endif bitsandbytes/csrc/ops.hip000066400000000000000000001115551506752665600160400ustar00rootroot00000000000000// !!! This is a file automatically generated by hipify!!! #include "hip/hip_runtime.h" // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #include #include #include #include #include #ifndef NO_HIPBLASLT #include #endif #include #include #include #include #define ERR_NOT_IMPLEMENTED 100 #if defined(__GFX9__) #define WARP_SIZE 64 #else #define WARP_SIZE 32 #endif using namespace BinSearch; using std::cout; using std::endl; void quantize(float *code, float *A, unsigned char *out, int n) { int num_blocks = n/1024; num_blocks = n % 1024 == 0 ? num_blocks : num_blocks + 1; hipLaunchKernelGGL(( kQuantize), dim3(num_blocks), dim3(1024), 0, 0, code, A, out, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); } void dequantize(float *code, unsigned char *A, float *out, int n, hipStream_t stream) { int num_blocks = n/1024; num_blocks = n % 1024 == 0 ? num_blocks : num_blocks + 1; hipLaunchKernelGGL(( kDequantize), dim3(num_blocks), dim3(1024), 0, stream, code, A, out, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); } template void quantizeBlockwise(float * code, T *A, float *absmax, unsigned char *out, float *rand, int rand_offset, int blocksize, const int n) { int num_blocks = n/blocksize; num_blocks = n % blocksize == 0 ? num_blocks : num_blocks + 1; if(blocksize == 4096) hipLaunchKernelGGL(( kQuantizeBlockwise), dim3(num_blocks), dim3(1024), 0, 0, code, A, absmax, out, rand, rand_offset, n); else if(blocksize == 2048) hipLaunchKernelGGL(( kQuantizeBlockwise), dim3(num_blocks), dim3(512), 0, 0, code, A, absmax, out, rand, rand_offset, n); else if(blocksize == 1024) hipLaunchKernelGGL(( kQuantizeBlockwise), dim3(num_blocks), dim3(256), 0, 0, code, A, absmax, out, rand, rand_offset, n); else if(blocksize == 512) hipLaunchKernelGGL(( kQuantizeBlockwise), dim3(num_blocks), dim3(256), 0, 0, code, A, absmax, out, rand, rand_offset, n); else if(blocksize == 256) hipLaunchKernelGGL(( kQuantizeBlockwise), dim3(num_blocks), dim3(128), 0, 0, code, A, absmax, out, rand, rand_offset, n); else if(blocksize == 128) hipLaunchKernelGGL(( kQuantizeBlockwise), dim3(num_blocks), dim3(64), 0, 0, code, A, absmax, out, rand, rand_offset, n); //else if(blocksize == 64) // hipLaunchKernelGGL(( kQuantizeBlockwise), dim3(num_blocks), dim3(32), 0, 0, code, A, absmax, out, rand, rand_offset, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); } template void dequantizeBlockwise(float *code, unsigned char *A, float *absmax, T *out, int blocksize, const int n, hipStream_t stream) { int num_blocks = n/blocksize; num_blocks = n % blocksize == 0 ? num_blocks : num_blocks + 1; int tile_size = (DATA_TYPE > 0) ? 1024 : 512; if(DATA_TYPE > 0) hipLaunchKernelGGL(( kDequantizeBlockwise), dim3((n+tile_size-1)/tile_size), dim3(64), 0, stream, code, A, absmax, out, blocksize/2, n); else hipLaunchKernelGGL(( kDequantizeBlockwise), dim3((n+tile_size-1)/tile_size), dim3(64), 0, stream, code, A, absmax, out, blocksize, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); } template void optimizer32bit(T* g, T* p, float* state1, float* state2, float *unorm, float max_unorm, float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, const int step, const float lr, const float gnorm_scale, bool skip_zeros, const int n) { int num_blocks = n/4096; num_blocks = n % 4096 == 0 ? num_blocks : num_blocks + 1; switch(OPTIMIZER) { case ADAM: case ADEMAMIX: if(max_unorm > 0.0f) { CUDA_CHECK_RETURN(hipMemset(unorm, 0, 1*sizeof(float))); hipLaunchKernelGGL(( kPreconditionOptimizer32bit2State), dim3(num_blocks), dim3(512), 0, 0, g, p, state1, state2, unorm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); } hipLaunchKernelGGL(( kOptimizer32bit2State), dim3(num_blocks), dim3(1024), 0, 0, g, p, state1, state2, unorm, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, lr, gnorm_scale, skip_zeros, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); break; case MOMENTUM: case RMSPROP: case ADAGRAD: if(max_unorm > 0.0f) { CUDA_CHECK_RETURN(hipMemset(unorm, 0, 1*sizeof(float))); hipLaunchKernelGGL(( kPreconditionOptimizer32bit1State), dim3(num_blocks), dim3(512), 0, 0, g, p, state1, unorm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); } hipLaunchKernelGGL(( kOptimizer32bit1State), dim3(num_blocks), dim3(1024), 0, 0, g, p, state1, unorm, max_unorm, param_norm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, skip_zeros, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); break; case LION: // in lion, the momentum update after the parameter update hipLaunchKernelGGL(( kOptimizer32bit1State), dim3(num_blocks), dim3(1024), 0, 0, g, p, state1, unorm, max_unorm, param_norm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, skip_zeros, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); if(max_unorm > 0.0f) { CUDA_CHECK_RETURN(hipMemset(unorm, 0, 1*sizeof(float))); hipLaunchKernelGGL(( kPreconditionOptimizer32bit1State), dim3(num_blocks), dim3(512), 0, 0, g, p, state1, unorm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); } break; } } template void optimizerStatic8bit(T* p, T* g, unsigned char* state1, unsigned char* state2, float *unorm, float max_unorm, float param_norm, float beta1, float beta2, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, float weight_decay, const float gnorm_scale, int n) { int num_blocks = n/4096; num_blocks = n % 4096 == 0 ? num_blocks : num_blocks + 1; if(max_unorm > 0.0f){ CUDA_CHECK_RETURN(hipMemset(unorm, 0, 1*sizeof(float))); } switch(OPTIMIZER) { case ADAM: CUDA_CHECK_RETURN(hipMemset(new_max1, 0, 1*sizeof(float))); CUDA_CHECK_RETURN(hipMemset(new_max2, 0, 1*sizeof(float))); hipLaunchKernelGGL(( kPreconditionOptimizerStatic8bit2State), dim3(num_blocks), dim3(256), 0, 0, p, g, state1, state2, unorm, beta1, beta2, eps, step, quantiles1, quantiles2, max1, max2, new_max1, new_max2, gnorm_scale, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); hipLaunchKernelGGL(( kOptimizerStatic8bit2State), dim3(num_blocks), dim3(1024), 0, 0, p, g, state1, state2, unorm, max_unorm, param_norm, beta1, beta2, eps, step, lr, quantiles1, quantiles2, max1, max2, new_max1, new_max2, weight_decay, gnorm_scale, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); break; case MOMENTUM: case RMSPROP: case ADAGRAD: CUDA_CHECK_RETURN(hipMemset(new_max1, 0, 1*sizeof(float))); hipLaunchKernelGGL(( kPreconditionOptimizerStatic8bit1State), dim3(num_blocks), dim3(256), 0, 0, p, g, state1, unorm, beta1, beta2, eps, step, quantiles1, max1, new_max1, weight_decay, gnorm_scale, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); hipLaunchKernelGGL(( kOptimizerStatic8bit1State), dim3(num_blocks), dim3(1024), 0, 0, p, g, state1, unorm, max_unorm, param_norm, beta1, beta2, eps, step, lr, quantiles1, max1, new_max1, weight_decay, gnorm_scale, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); break; case LION: // in lion, the momentum update happens after the parameter update hipLaunchKernelGGL(( kOptimizerStatic8bit1State), dim3(num_blocks), dim3(1024), 0, 0, p, g, state1, unorm, max_unorm, param_norm, beta1, beta2, eps, step, lr, quantiles1, max1, new_max1, weight_decay, gnorm_scale, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); CUDA_CHECK_RETURN(hipMemset(new_max1, 0, 1*sizeof(float))); hipLaunchKernelGGL(( kPreconditionOptimizerStatic8bit1State), dim3(num_blocks), dim3(256), 0, 0, p, g, state1, unorm, beta1, beta2, eps, step, quantiles1, max1, new_max1, weight_decay, gnorm_scale, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); break; default: break; } } #define BLOCKSIZE_2STATE 256 #define NUM_2STATE 1 #define BLOCKSIZE_1STATE 256 #define NUM_1STATE 1 template void optimizerStatic8bitBlockwise( T* p, T* g, unsigned char* state1, unsigned char* state2, float beta1, float beta2, float beta3, float alpha, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* absmax1, float* absmax2, float weight_decay, const float gnorm_scale, bool skip_zeros, int n ) { int num_blocks = 0; switch(OPTIMIZER) { case ADAM: case ADEMAMIX: num_blocks = n/BLOCKSIZE_2STATE; num_blocks = n % BLOCKSIZE_2STATE == 0 ? num_blocks : num_blocks + 1; hipLaunchKernelGGL(( kOptimizerStatic8bit2StateBlockwise), dim3(num_blocks), dim3(BLOCKSIZE_2STATE/NUM_2STATE), 0, 0, p, g, state1, state2, beta1, beta2, beta3, alpha, eps, step, lr, quantiles1, quantiles2, absmax1, absmax2, weight_decay, gnorm_scale, skip_zeros, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); break; case MOMENTUM: case RMSPROP: case ADAGRAD: case LION: num_blocks = n/BLOCKSIZE_1STATE; num_blocks = n % BLOCKSIZE_1STATE == 0 ? num_blocks : num_blocks + 1; hipLaunchKernelGGL(( kOptimizerStatic8bit1StateBlockwise), dim3(num_blocks), dim3(BLOCKSIZE_1STATE/NUM_1STATE), 0, 0, p, g, state1, beta1, beta2, eps, step, lr, quantiles1, absmax1, weight_decay, gnorm_scale, skip_zeros, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); break; } } template void percentileClipping(T * g, float *gnorm_vec, int step, const int n) { int num_blocks = n/2048; num_blocks = n % 2048 == 0 ? num_blocks : num_blocks + 1; CUDA_CHECK_RETURN(hipMemset(&gnorm_vec[step % 100], 0, 1*sizeof(float))); hipLaunchKernelGGL(( kPercentileClipping), dim3(num_blocks), dim3(512), 0, 0, g, gnorm_vec, step, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); } void gemmex(Context *context, bool transposeA, bool transposeB, int m, int n, int k, void *A, void *B, void *C, int lda, int ldb, int ldc) { const int falpha = 1; const int fbeta = 0; const void * alpha = &falpha; const void * beta = &fbeta; hipblasStatus_t status; #if hipblasVersionMajor >= 3 status = hipblasGemmEx(context->m_handle, transposeA ? HIPBLAS_OP_T : HIPBLAS_OP_N, transposeB ? HIPBLAS_OP_T : HIPBLAS_OP_N, m, n, k, alpha, A, HIP_R_8I, lda, B, HIP_R_8I, ldb, beta, C, HIP_R_32I, ldc, HIPBLAS_COMPUTE_32I, HIPBLAS_GEMM_DEFAULT); #else status = hipblasGemmEx(context->m_handle, transposeA ? HIPBLAS_OP_T : HIPBLAS_OP_N, transposeB ? HIPBLAS_OP_T : HIPBLAS_OP_N, m, n, k, alpha, A, HIPBLAS_R_8I, lda, B, HIPBLAS_R_8I, ldb, beta, C, HIPBLAS_R_32I, ldc, HIPBLAS_R_32I, HIPBLAS_GEMM_DEFAULT); #endif if (status != HIPBLAS_STATUS_SUCCESS) { std::cout << "HIPBLAS ERROR: Status " << status << std::endl; } } void strided_gemmex(Context *context, bool transposeA, bool transposeB, int m, int n, int k, void *A, void *B, void *C, int lda, int ldb, int ldc, long long int strideA, long long int strideB, long long int strideC, int batchCount) { const int falpha = 1; const int fbeta = 0; const void * alpha = &falpha; const void * beta = &fbeta; hipblasStatus_t status; //cout << transposeA << transposeB << endl; //printf("%i %i %i\n", m,n,k); //printf("%i %i %i\n", lda,ldb,ldc); //printf("%i %i %i\n", strideA, strideB, strideC); //printf("%i\n", batchCount); #if hipblasVersionMajor >= 3 status = hipblasGemmStridedBatchedEx(context->m_handle, transposeA ? HIPBLAS_OP_T : HIPBLAS_OP_N, transposeB ? HIPBLAS_OP_T : HIPBLAS_OP_N, m, n, k, alpha, A, HIP_R_8I, lda, (long long int)strideA, B, HIP_R_8I, ldb, (long long int)strideB, beta, C, HIP_R_32I, ldc, (long long int)strideC, batchCount, HIPBLAS_COMPUTE_32I, HIPBLAS_GEMM_DEFAULT); #else status = hipblasGemmStridedBatchedEx(context->m_handle, transposeA ? HIPBLAS_OP_T : HIPBLAS_OP_N, transposeB ? HIPBLAS_OP_T : HIPBLAS_OP_N, m, n, k, alpha, A, HIPBLAS_R_8I, lda, (long long int)strideA, B, HIPBLAS_R_8I, ldb, (long long int)strideB, beta, C, HIPBLAS_R_32I, ldc, (long long int)strideC, batchCount, HIPBLAS_R_32I, HIPBLAS_GEMM_DEFAULT); #endif if (status != HIPBLAS_STATUS_SUCCESS) { std::cout << "HIPBLAS ERROR: Status " << status << std::endl; } } int roundoff(int v, int d) { return (v + d - 1) / d * d; } #ifdef NO_HIPBLASLT #else template hipblasLtOrder_t get_order() { switch(ORDER) { case ROW: return HIPBLASLT_ORDER_ROW; break; case COL: return HIPBLASLT_ORDER_COL; break; case COL32: //return HIPBLASLT_ORDER_COL32; return HIPBLASLT_ORDER_COL; break; case COL_TURING: //return HIPBLASLT_ORDER_COL4_4R2_8C; return HIPBLASLT_ORDER_COL; break; case COL_AMPERE: //return HIPBLASLT_ORDER_COL32_2R_4R4; return HIPBLASLT_ORDER_COL; break; default: break; } return HIPBLASLT_ORDER_ROW; } template hipblasLtOrder_t get_order(); template hipblasLtOrder_t get_order(); template hipblasLtOrder_t get_order(); //template hipblasLtOrder_t get_order(); //template hipblasLtOrder_t get_order(); #endif template int get_leading_dim(int dim1, int dim2) { switch(ORDER) { case ROW: return dim2; break; case COL: return dim1; break; default: return dim1; break; /*case COL32: // 32*row tiles return dim1*32; break; case COL_TURING: return 32*roundoff(dim1, 8); break; case COL_AMPERE: // 32*32 tiles return 32*roundoff(dim1, 32); break; default: return 0; break; */ } } static std::string hipError_to_string(const hipError_t ret) { switch(ret) { case hipSuccess: return "hipSuccess"; case hipErrorInvalidContext: return "hipErrorInvalidContext"; case hipErrorInvalidKernelFile: return "hipErrorInvalidKernelFile"; case hipErrorMemoryAllocation: return "hipErrorMemoryAllocation"; case hipErrorInitializationError: return "hipErrorInitializationError"; case hipErrorLaunchFailure: return "hipErrorLaunchFailure"; case hipErrorLaunchOutOfResources: return "hipErrorLaunchOutOfResources"; case hipErrorInvalidDevice: return "hipErrorInvalidDevice"; case hipErrorInvalidValue: return "hipErrorInvalidValue"; case hipErrorInvalidDevicePointer: return "hipErrorInvalidDevicePointer"; case hipErrorInvalidMemcpyDirection: return "hipErrorInvalidMemcpyDirection"; case hipErrorUnknown: return "hipErrorUnknown"; case hipErrorInvalidResourceHandle: return "hipErrorInvalidResourceHandle"; case hipErrorNotReady: return "hipErrorNotReady"; case hipErrorNoDevice: return "hipErrorNoDevice"; case hipErrorPeerAccessAlreadyEnabled: return "hipErrorPeerAccessAlreadyEnabled"; case hipErrorPeerAccessNotEnabled: return "hipErrorPeerAccessNotEnabled"; case hipErrorRuntimeMemory: return "hipErrorRuntimeMemory"; case hipErrorRuntimeOther: return "hipErrorRuntimeOther"; case hipErrorHostMemoryAlreadyRegistered: return "hipErrorHostMemoryAlreadyRegistered"; case hipErrorHostMemoryNotRegistered: return "hipErrorHostMemoryNotRegistered"; case hipErrorMapBufferObjectFailed: return "hipErrorMapBufferObjectFailed"; case hipErrorTbd: return "hipErrorTbd"; default: throw std::runtime_error("unknown hipError"); } } template int igemmlt( hipblasLtHandle_t ltHandle, int m, int n, int k, const int8_t *A, const int8_t *B, void *C, float *row_scale, int lda, int ldb, int ldc, hipStream_t stream ) { #ifdef NO_HIPBLASLT return ERR_NOT_IMPLEMENTED; #else // Calculate C = A^T @ B, in col-major layout. // // Use the IMMA kernels requires: // * A must be transposed and B must be non-transposed. // * Dimensions m and k must be multiples of 4. // * All pointers must be 4-byte aligned; 16-byte alignment preferred. int has_error = 0; const int64_t max_workspace_size = 0;//set to 0 to avoid choosing GSU kernel hipblasLtMatmulDesc_t matmulDesc; hipblasLtMatrixLayout_t aDesc, bDesc, cDesc; hipblasOperation_t opT = HIPBLAS_OP_T; hipDataType outType = DTYPE_OUT == 32 ? HIP_R_32I : HIP_R_8I; hipDataType scaleType = DTYPE_OUT == 32 ? HIP_R_32I : HIP_R_32F; hipblasLtPointerMode_t pointerMode = HIPBLASLT_POINTER_MODE_ALPHA_DEVICE_VECTOR_BETA_HOST; has_error |= checkHipblasStatus(hipblasLtMatrixLayoutCreate(&aDesc, HIP_R_8I, m, k, lda)); has_error |= checkHipblasStatus(hipblasLtMatrixLayoutCreate(&bDesc, HIP_R_8I, m, n, ldb)); has_error |= checkHipblasStatus(hipblasLtMatrixLayoutCreate(&cDesc, outType, k, n, ldc)); // Default layout order is col major has_error |= checkHipblasStatus(hipblasLtMatmulDescCreate(&matmulDesc, HIPBLAS_COMPUTE_32I, scaleType)); has_error |= checkHipblasStatus(hipblasLtMatmulDescSetAttribute(matmulDesc, HIPBLASLT_MATMUL_DESC_TRANSA, &opT, sizeof(opT))); if (DTYPE_OUT == 32) { /* Algo and workspace TODO: need to rework to not be duplicated */ // Set User Preference attributes hipblasLtMatmulPreference_t pref; checkHipblasStatus(hipblasLtMatmulPreferenceCreate(&pref)); checkHipblasStatus( hipblasLtMatmulPreferenceSetAttribute(pref, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &max_workspace_size, sizeof(max_workspace_size))); const int request_solutions = 1; hipblasLtMatmulHeuristicResult_t heuristicResult[request_solutions]; int returnedAlgoCount = 0; checkHipblasStatus(hipblasLtMatmulAlgoGetHeuristic(ltHandle, matmulDesc, aDesc, bDesc, cDesc, cDesc, pref, request_solutions, heuristicResult, &returnedAlgoCount)); if (returnedAlgoCount == 0) { has_error = 1; fprintf(stderr, "Error: Matmul Algo Heuristic didn't return algorithms\n"); } else { int alpha = 1, beta = 0; has_error |= checkHipblasStatus(hipblasLtMatmul( ltHandle, matmulDesc, &alpha, A, aDesc, B, bDesc, &beta, (int32_t*)C, cDesc, (int32_t*)C, cDesc, &heuristicResult[0].algo, NULL, 0, stream )); } } else { // This path is unlikely to be used, as 8-bit accumulation can lead to likely overflows. if (!SCALE_ROWS) { float alpha = 1.0f, beta = 0.0f; has_error |= checkHipblasStatus(hipblasLtMatmul( ltHandle, matmulDesc, &alpha, A, aDesc, B, bDesc, &beta, (int8_t*)C, cDesc, (int8_t*)C, cDesc, NULL, NULL, 0, stream )); } else { hipblasLtPointerMode_t alphaVec = HIPBLASLT_POINTER_MODE_ALPHA_DEVICE_VECTOR_BETA_HOST; float beta = 0.0f; has_error |= checkHipblasStatus(hipblasLtMatmulDescSetAttribute( matmulDesc, HIPBLASLT_MATMUL_DESC_POINTER_MODE, &pointerMode, sizeof(alphaVec) )); has_error |= checkHipblasStatus(hipblasLtMatmul( ltHandle, matmulDesc, row_scale, A, aDesc, B, bDesc, &beta, (int8_t*)C, cDesc, (int8_t*)C, cDesc, NULL, NULL, 0, stream )); } } has_error |= checkHipblasStatus(hipblasLtMatrixLayoutDestroy(cDesc)); has_error |= checkHipblasStatus(hipblasLtMatrixLayoutDestroy(bDesc)); has_error |= checkHipblasStatus(hipblasLtMatrixLayoutDestroy(aDesc)); has_error |= checkHipblasStatus(hipblasLtMatmulDescDestroy(matmulDesc)); if(has_error == 1) printf("error detected"); return has_error; #endif // NO_HIPBLASLT } int fill_up_to_nearest_multiple(int value, int multiple) { return value + (value % multiple == 0 ? 0 : (multiple - (value % multiple))); } void dequant_mm_int32_fp16(int *A, float *rowStats, float *colStats, half *out, half *bias, int numRows, int numCols, hipStream_t stream) { const int threads = 512; const int num_per_thread = 4; const int num_per_block = threads * num_per_thread; const int n = numRows*numCols; const int num_blocks = (n + num_per_block - 1) / num_per_block; hipLaunchKernelGGL(( kdequant_mm_int32_fp16), dim3(num_blocks), dim3(threads), 0, stream, A, rowStats, colStats, out, bias, numRows, numCols, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); } void int8VectorQuant(half * __restrict__ A, int8_t *out, float *rowStats, float threshold, int rows, int cols, hipStream_t stream) { if (threshold == 0.0) { kInt8VectorQuant<<>>(A, out, rowStats, threshold, rows, cols); } else { kInt8VectorQuant<<>>(A, out, rowStats, threshold, rows, cols); } CUDA_CHECK_RETURN(hipPeekAtLastError()); } void getRowStats(half *A, float *rowStats, float threshold, int rows, int cols, hipStream_t stream) { if (threshold == 0.0) kgetRowStats<<>>(A, rowStats, threshold, rows, cols); else kgetRowStats<<>>(A, rowStats, threshold, rows, cols); CUDA_CHECK_RETURN(hipPeekAtLastError()); } void spmm_coo(hipsparseHandle_t handle, int *A_rowidx, int *A_colidx, half *A_vals, int A_nnz, int A_rows, int A_cols, int B_cols, int ldb, half *B, int ldc, half* C, bool transposed_B) { #ifdef NO_HIPBLASLT #else hipsparseSpMatDescr_t descA; hipsparseDnMatDescr_t descB, descC; float alpha = 1.0f; float beta = 0.0f; void *dBuffer = NULL; size_t bufferSize = 0; CHECK_HIPSPARSE( hipsparseCreateCoo(&descA, A_rows, A_cols, A_nnz, A_rowidx, A_colidx, A_vals, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_BASE_ZERO, HIP_R_16F) ); // Create dense matrix C CHECK_HIPSPARSE( hipsparseCreateDnMat(&descC, A_rows, B_cols, ldc, C, HIP_R_16F, HIPSPARSE_ORDER_ROW) ); // Create dense matrix B if(transposed_B) { int tmp = A_cols; A_cols = B_cols; B_cols = tmp; } CHECK_HIPSPARSE( hipsparseCreateDnMat(&descB, A_cols, B_cols, ldb, B, HIP_R_16F, HIPSPARSE_ORDER_ROW) ); // allocate an external buffer if needed CHECK_HIPSPARSE( hipsparseSpMM_bufferSize( handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, transposed_B ? HIPSPARSE_OPERATION_TRANSPOSE : HIPSPARSE_OPERATION_NON_TRANSPOSE, &alpha, descA, descB, &beta, descC, HIP_R_32F, HIPSPARSE_SPMM_ALG_DEFAULT, &bufferSize) ); CUDA_CHECK_RETURN( hipMalloc(&dBuffer, bufferSize) ); // execute SpMM CHECK_HIPSPARSE( hipsparseSpMM(handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, transposed_B ? HIPSPARSE_OPERATION_TRANSPOSE : HIPSPARSE_OPERATION_NON_TRANSPOSE, &alpha, descA, descB, &beta, descC, HIP_R_32F, HIPSPARSE_SPMM_ALG_DEFAULT, dBuffer)); // destroy matrix/vector descriptors CHECK_HIPSPARSE( hipsparseDestroySpMat(descA) ); CHECK_HIPSPARSE( hipsparseDestroyDnMat(descB) ); CHECK_HIPSPARSE( hipsparseDestroyDnMat(descC) ); CUDA_CHECK_RETURN( hipFree(dBuffer) ); #endif } template void spmm_coo_very_sparse_naive(int *max_count, int *max_idx, int *offset_rowidx, int *rowidx, int *colidx, half *values, T *B, half *out, float *dequant_stats, int nnz_rows, int nnz, int rowsA, int rowsB, int colsB) { hipLaunchKernelGGL(( kspmm_coo_very_sparse_naive), dim3(nnz_rows), dim3(256), 0, 0, max_count, max_idx, offset_rowidx, rowidx, colidx, values, B, out, dequant_stats, nnz, rowsA, rowsB, colsB); CUDA_CHECK_RETURN(hipPeekAtLastError()); } template void gemm_host(int m, int n, int k, T * A, T* B, T * out, int lda, int ldb, int ldc, int bits) { int num_blocks = (m+31)/32; if(bits == 32) hipLaunchKernelGGL(( gemm_device), dim3(num_blocks), dim3(32), 0, 0, m, n, k, A, B, out, lda, ldb, ldc); if(bits == 16) hipLaunchKernelGGL(( gemm_device), dim3(num_blocks), dim3(160), 0, 0, m, n, k, A, B, out, lda, ldb, ldc); } template void gemm_4bit_inference(int m, int n, int k, T * A, unsigned char* B, float *absmax, T * out, int lda, int ldb, int ldc, int blocksize) { int num_blocks = (m+31)/32; hipLaunchKernelGGL(( kgemm_4bit_inference), dim3(num_blocks), dim3(96), 0, 0, m, n, k, A, B, absmax, out, lda, ldb, ldc, blocksize); } template void gemm_4bit_inference_naive(int m, int n, int k, T * A, unsigned char* B, float *absmax, float *datatype, T * out, int lda, int ldb, int ldc, int blocksize, hipStream_t stream) { //warpsize - 32 int num_blocks = (m+3)/4; //warpsize - 64 if (WARP_SIZE == 64) { num_blocks = (m+1)/2; } hipLaunchKernelGGL(( kgemm_4bit_inference_naive), dim3(num_blocks), dim3(128), 0, stream, m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize); CUDA_CHECK_RETURN(hipPeekAtLastError()); } template void func(T *A, T *B, T value, long n) { int threads = 512; int blocks = n/threads; blocks = n % threads == 0 ? blocks : blocks + 1; blocks = blocks > 65535 ? 65535 : blocks; hipLaunchKernelGGL(( kfunc), dim3(blocks), dim3(512), 0, 0, A, B, value, n); CUDA_CHECK_RETURN(hipPeekAtLastError()); } //============================================================== // TEMPLATE DEFINITIONS //============================================================== template void func(float *A, float *B, float value, long n); template void func(unsigned char *A, unsigned char *B, unsigned char value, long n); template void func(float *A, float *B, float value, long n); template void func(float *A, float *B, float value, long n); template void gemm_4bit_inference(int m, int n, int k, half * A, unsigned char* B, float *absmax, half * out, int lda, int ldb, int ldc, int blocksize); template void gemm_4bit_inference_naive(int m, int n, int k, half * A, unsigned char* B, float *absmax, float *datatype, half * out, int lda, int ldb, int ldc, int blocksize, hipStream_t stream); template void gemm_4bit_inference_naive(int m, int n, int k, hip_bfloat16 * A, unsigned char* B, float *absmax, float *datatype, hip_bfloat16 * out, int lda, int ldb, int ldc, int blocksize, hipStream_t stream); template void gemm_4bit_inference_naive(int m, int n, int k, float * A, unsigned char* B, float *absmax, float *datatype, float * out, int lda, int ldb, int ldc, int blocksize, hipStream_t stream); //template void gemm_host(int m, int n, int k, float * A, float* B, float * out, int lda, int ldb, int ldc, int bits); template void gemm_host(int m, int n, int k, half * A, half* B, half * out, int lda, int ldb, int ldc, int bits); template void spmm_coo_very_sparse_naive(int *max_count, int *max_idx, int *offset_rowidx, int *rowidx, int *colidx, half *values, half *B, half *out, float *dequant_stats, int nnz_rows, int nnz, int rowsA, int rowsB, int colsB); template void spmm_coo_very_sparse_naive(int *max_count, int *max_idx, int *offset_rowidx, int *rowidx, int *colidx, half *values, signed char *B, half *out, float *dequant_stats, int nnz_rows, int nnz, int rowsA, int rowsB, int colsB); template int igemmlt<32, 0>(hipblasLtHandle_t ltHandle, int m, int n, int k, const int8_t *A, const int8_t *B, void *C, float *row_scale, int lda, int ldb, int ldc, hipStream_t stream); template int igemmlt<8, 0>(hipblasLtHandle_t ltHandle, int m, int n, int k, const int8_t *A, const int8_t *B, void *C, float *row_scale, int lda, int ldb, int ldc, hipStream_t stream); template int igemmlt<8, 1>(hipblasLtHandle_t ltHandle, int m, int n, int k, const int8_t *A, const int8_t *B, void *C, float *row_scale, int lda, int ldb, int ldc, hipStream_t stream); template void quantizeBlockwise(float * code, half *A, float *absmax, unsigned char *out, float* rand, int rand_offset, int blocksize, const int n); template void quantizeBlockwise(float * code, half *A, float *absmax, unsigned char *out, float* rand, int rand_offset, int blocksize, const int n); template void quantizeBlockwise(float * code, half *A, float *absmax, unsigned char *out, float* rand, int rand_offset, int blocksize, const int n); template void quantizeBlockwise(float * code, half *A, float *absmax, unsigned char *out, float* rand, int rand_offset, int blocksize, const int n); template void quantizeBlockwise(float * code, float *A, float *absmax, unsigned char *out, float* rand, int rand_offset, int blocksize, const int n); template void quantizeBlockwise(float * code, float *A, float *absmax, unsigned char *out, float* rand, int rand_offset, int blocksize, const int n); template void quantizeBlockwise(float * code, float *A, float *absmax, unsigned char *out, float* rand, int rand_offset, int blocksize, const int n); template void quantizeBlockwise(float * code, float *A, float *absmax, unsigned char *out, float* rand, int rand_offset, int blocksize, const int n); template void quantizeBlockwise(float * code, hip_bfloat16 *A, float *absmax, unsigned char *out, float* rand, int rand_offset, int blocksize, const int n); template void quantizeBlockwise(float * code, hip_bfloat16 *A, float *absmax, unsigned char *out, float* rand, int rand_offset, int blocksize, const int n); template void quantizeBlockwise(float * code, hip_bfloat16 *A, float *absmax, unsigned char *out, float* rand, int rand_offset, int blocksize, const int n); template void quantizeBlockwise(float * code, hip_bfloat16 *A, float *absmax, unsigned char *out, float* rand, int rand_offset, int blocksize, const int n); template void dequantizeBlockwise(float *code, unsigned char *A, float *absmax, float *out, int blocksize, const int n, hipStream_t stream); template void dequantizeBlockwise(float *code, unsigned char *A, float *absmax, float *out, int blocksize, const int n, hipStream_t stream); template void dequantizeBlockwise(float *code, unsigned char *A, float *absmax, float *out, int blocksize, const int n, hipStream_t stream); template void dequantizeBlockwise(float *code, unsigned char *A, float *absmax, half *out, int blocksize, const int n, hipStream_t stream); template void dequantizeBlockwise(float *code, unsigned char *A, float *absmax, half *out, int blocksize, const int n, hipStream_t stream); template void dequantizeBlockwise(float *code, unsigned char *A, float *absmax, half *out, int blocksize, const int n, hipStream_t stream); template void dequantizeBlockwise(float *code, unsigned char *A, float *absmax, hip_bfloat16 *out, int blocksize, const int n, hipStream_t stream); template void dequantizeBlockwise(float *code, unsigned char *A, float *absmax, hip_bfloat16 *out, int blocksize, const int n, hipStream_t stream); template void dequantizeBlockwise(float *code, unsigned char *A, float *absmax, hip_bfloat16 *out, int blocksize, const int n, hipStream_t stream); #define MAKE_optimizer32bit(name, gtype) \ template void optimizer32bit(gtype* g, gtype* p, \ float* state1, float* state2, float* unorm, float max_unorm, float param_norm, \ const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, \ const int step, const float lr, const float gnorm_scale, const bool skip_zeros, const int n); MAKE_optimizer32bit(ADAM, half) MAKE_optimizer32bit(ADAM, float) MAKE_optimizer32bit(ADAM, hip_bfloat16) MAKE_optimizer32bit(MOMENTUM, half) MAKE_optimizer32bit(MOMENTUM, float) MAKE_optimizer32bit(MOMENTUM, hip_bfloat16) MAKE_optimizer32bit(RMSPROP, half) MAKE_optimizer32bit(RMSPROP, float) MAKE_optimizer32bit(RMSPROP, hip_bfloat16) MAKE_optimizer32bit(LION, half) MAKE_optimizer32bit(LION, float) MAKE_optimizer32bit(LION, hip_bfloat16) MAKE_optimizer32bit(ADAGRAD, half) MAKE_optimizer32bit(ADAGRAD, float) MAKE_optimizer32bit(ADAGRAD, hip_bfloat16) MAKE_optimizer32bit(ADEMAMIX, half) MAKE_optimizer32bit(ADEMAMIX, hip_bfloat16) MAKE_optimizer32bit(ADEMAMIX, float) #define MAKE_optimizerStatic8bit(name, gtype) \ template void optimizerStatic8bit(gtype* p, gtype* g, unsigned char* state1, unsigned char* state2, \ float *unorm, float max_unorm, float param_norm, \ float beta1, float beta2, \ float eps, int step, float lr, \ float* quantiles1, float* quantiles2, \ float* max1, float* max2, float* new_max1, float* new_max2, \ float weight_decay, \ const float gnorm_scale, int n); \ MAKE_optimizerStatic8bit(ADAM, half) MAKE_optimizerStatic8bit(ADAM, float) MAKE_optimizerStatic8bit(MOMENTUM, half) MAKE_optimizerStatic8bit(MOMENTUM, float) MAKE_optimizerStatic8bit(RMSPROP, half) MAKE_optimizerStatic8bit(RMSPROP, float) MAKE_optimizerStatic8bit(LION, half) MAKE_optimizerStatic8bit(LION, float) MAKE_optimizerStatic8bit(ADAGRAD, half) MAKE_optimizerStatic8bit(ADAGRAD, float) #define MAKE_optimizerStatic8bitBlockwise(gtype, optim_name) \ template void optimizerStatic8bitBlockwise(gtype* p, gtype* g, \ unsigned char* state1, unsigned char* state2, float beta1, float beta2, float beta3, float alpha, float eps, int step, float lr, \ float* quantiles1, float* quantiles2, float* absmax1, float* absmax2, float weight_decay, const float gnorm_scale, bool skip_zeros, int n); \ MAKE_optimizerStatic8bitBlockwise(half, ADAM); MAKE_optimizerStatic8bitBlockwise(float, ADAM); MAKE_optimizerStatic8bitBlockwise(hip_bfloat16, ADAM); MAKE_optimizerStatic8bitBlockwise(half, MOMENTUM); MAKE_optimizerStatic8bitBlockwise(float, MOMENTUM); MAKE_optimizerStatic8bitBlockwise(hip_bfloat16, MOMENTUM); MAKE_optimizerStatic8bitBlockwise(half, RMSPROP); MAKE_optimizerStatic8bitBlockwise(float, RMSPROP); MAKE_optimizerStatic8bitBlockwise(hip_bfloat16, RMSPROP); MAKE_optimizerStatic8bitBlockwise(half, LION); MAKE_optimizerStatic8bitBlockwise(float, LION); MAKE_optimizerStatic8bitBlockwise(hip_bfloat16, LION); MAKE_optimizerStatic8bitBlockwise(half, ADAGRAD); MAKE_optimizerStatic8bitBlockwise(float, ADAGRAD); MAKE_optimizerStatic8bitBlockwise(hip_bfloat16, ADAGRAD); MAKE_optimizerStatic8bitBlockwise(half, ADEMAMIX); MAKE_optimizerStatic8bitBlockwise(hip_bfloat16, ADEMAMIX); MAKE_optimizerStatic8bitBlockwise(float, ADEMAMIX); template void percentileClipping(float * g, float *gnorm_vec, int step, const int n); template void percentileClipping(half * g, float *gnorm_vec, int step, const int n); template int get_leading_dim(int dim1, int dim2); template int get_leading_dim(int dim1, int dim2); template int get_leading_dim(int dim1, int dim2); bitsandbytes/csrc/ops_hip.cuh000066400000000000000000000175551506752665600167040ustar00rootroot00000000000000// !!! This is a file automatically generated by hipify!!! // Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #ifndef ops_H #define ops_H #include #include #include #include #include #include #include #include #include #include #include #include #define CUDA_CHECK_RETURN(value) \ { \ hipError_t _m_cudaStat = value; \ if (_m_cudaStat != hipSuccess) { \ fprintf(stderr, "Error %s at line %d in file %s\n", hipGetErrorString(_m_cudaStat), __LINE__, __FILE__); \ exit(1); \ } \ } #define CHECK_HIPSPARSE(value) \ { \ hipsparseStatus_t _m_hipStat = value; \ if (_m_hipStat != HIPSPARSE_STATUS_SUCCESS) { \ fprintf( \ stderr, "Error %s at line %d in file %s\n", hipsparseGetErrorString(_m_hipStat), __LINE__, __FILE__ \ ); \ exit(1); \ } \ } inline void checkHipStatus(hipError_t status) { if (status != hipSuccess) { printf("hip API failed with status %d: %s\n", status, hipGetErrorString(status)); throw std::logic_error("hip API failed"); } } inline int checkHipblasStatus(hipblasStatus_t status) { if (status != HIPBLAS_STATUS_SUCCESS) { printf("hipBLAS API failed with status %d\n", status); // throw std::logic_error("cuBLAS API failed"); return 1; } return 0; } typedef enum Operations_t { ksmul = 0, } Operations_t; typedef enum Optimizer_t { ADAM = 0, MOMENTUM = 1, RMSPROP = 2, LARS = 3, ADAGRAD = 4, LION = 5, ADEMAMIX = 6, } Optimizer_t; typedef enum Transform_t { ROW = 0, COL = 1, COL32 = 2, COL_TURING = 3, COL_AMPERE = 4, } Transform_t; typedef enum DataType_t { General8bit = 0, FP4 = 1, NF4 = 2, } DataType_t; typedef enum Funcs_t { FILL = 0, ARANGE = 1, _MUL = 2, } Funcs_t; class Context { public: rocblas_handle m_handle; Context() { rocblas_handle handle; rocblas_create_handle(&handle); m_handle = handle; } }; class ContextLt { public: hipblasLtHandle_t m_handle; ContextLt() { hipblasLtHandle_t handle; hipblasLtCreate(&handle); m_handle = handle; } }; class ContextHipsparse { public: hipsparseHandle_t m_handle; ContextHipsparse() { hipsparseHandle_t handle; hipsparseCreate(&handle); m_handle = handle; } }; void quantize(float* code, float* A, unsigned char* out, int n); void dequantize(float* code, unsigned char* A, float* out, int n, hipStream_t stream); template void quantizeBlockwise( float* code, T* A, float* absmax, unsigned char* out, float* rand, int rand_offset, int blocksize, const int n ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, T* out, int block_size, const int n, hipStream_t stream ); template void optimizer32bit( T* g, T* p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, float beta1, float beta2, float beta3, float alpha, float eps, float weight_decay, int step, float lr, const float gnorm_scale, bool skip_zeros, int n ); template void optimizerStatic8bit( T* p, T* g, unsigned char* state1, unsigned char* state2, float* unorm, float max_unorm, float param_norm, float beta1, float beta2, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, float weight_decay, const float gnorm_scale, int n ); template void optimizerStatic8bitBlockwise( T* p, T* g, unsigned char* state1, unsigned char* state2, float beta1, float beta2, float beta3, float alpha, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* absmax1, float* absmax2, float weight_decay, const float gnorm_scale, bool skip_zeros, int n ); template void percentileClipping(T* g, float* gnorm_vec, int step, const int n); void gemmex( Context* context, bool transposeA, bool transposeB, int m, int n, int k, void* A, void* B, void* C, int lda, int ldb, int ldc ); void strided_gemmex( Context* context, bool transposeA, bool transposeB, int m, int n, int k, void* A, void* B, void* C, int lda, int ldb, int ldc, long long int strideA, long long int strideB, long long int strideC, int batchCount ); template int igemmlt( hipblasLtHandle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, hipStream_t stream ); void cutlass_igemm( bool transposeA, bool transposeB, int m, int n, int k, void* A, void* B, void* C, int lda, int ldb, int ldc ); void dequant_mm_int32_fp16( int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, hipStream_t stream ); void getRowStats(half* A, float* rowStats, float threshold, int rows, int cols, hipStream_t stream); void int8VectorQuant( half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, hipStream_t stream ); void spmm_coo( hipsparseHandle_t handle, int* A_rowidx, int* A_colidx, half* A_vals, int A_nnz, int A_rows, int A_cols, int B_cols, int ldb, half* B, int ldc, half* C, bool transposed_B ); template void spmm_coo_very_sparse_naive( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, T* B, half* out, float* dequant_stats, int nnz_rows, int nnz, int rowsA, int rowsB, int colsB ); void matmul4bite(half* A, unsigned char* B, half* out, int lda, int ldb, int rowsA, int colsA, int colsB); template void gemm_host(int m, int n, int k, T* A, T* B, T* out, int lda, int ldb, int ldc, int bits); template void gemm_4bit_inference( int m, int n, int k, T* A, unsigned char* B, float* absmax, T* out, int lda, int ldb, int ldc, int blocksize ); template void gemm_4bit_inference_naive( int m, int n, int k, T* A, unsigned char* B, float* absmax, float* datatype, T* out, int lda, int ldb, int ldc, int blocksize, hipStream_t stream ); template void func(T* A, T* B, T value, long n); #endif bitsandbytes/csrc/pythonInterface.cpp000066400000000000000000001122461506752665600204010ustar00rootroot00000000000000// Copyright (c) Facebook, Inc. and its affiliates. // // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. #if BUILD_CUDA #include #include #endif #if BUILD_HIP #include #endif #if BUILD_MPS // #include #endif #if BUILD_XPU #include #endif #include // Compatibility between HIP/CUDA APIs #if BUILD_HIP #define cudaStream_t hipStream_t #define __nv_bfloat16 hip_bfloat16 #define cublasLtHandle_t hipblasLtHandle_t #define ContextCusparse ContextHipsparse #define cusparseHandle_t hipsparseHandle_t #define cudaMallocManaged hipMallocManaged #define cudaMemAttachHost hipMemAttachHost #define cudaPeekAtLastError hipPeekAtLastError #define cudaDeviceGetAttribute hipDeviceGetAttribute #define cudaDevAttrConcurrentManagedAccess hipDeviceAttributeConcurrentManagedAccess #define cudaMemPrefetchAsync hipMemPrefetchAsync #endif // We cannot call templated code from C, so we wrap the template in a C compatible call here if necessary. // We use macro functions to expand all the different optimizers. Looks ugly, and is ugly, but its better than to // maintain all that boilerplate //=================================================================================== // UNMANGLED CALLS //=================================================================================== #if BUILD_CUDA || BUILD_HIP // void gemm_host_fp32(int M, int N, int K, float * A, float* B, float * out, int lda, int ldb, int ldc) //{ gemm_host(M, N, K, A, B, out, lda, ldb, ldc, 32); } void gemm_host_fp16(int M, int N, int K, half* A, half* B, half* out, int lda, int ldb, int ldc) { gemm_host(M, N, K, A, B, out, lda, ldb, ldc, 16); } void gemm_4bit_inference( int m, int n, int k, half* A, unsigned char* B, float* absmax, half* out, int lda, int ldb, int ldc, int blocksize ) { gemm_4bit_inference(m, n, k, A, B, absmax, out, lda, ldb, ldc, blocksize); } void gemm_4bit_inference_naive_fp16( int m, int n, int k, half* A, unsigned char* B, float* absmax, float* datatype, half* out, int lda, int ldb, int ldc, int blocksize, cudaStream_t stream ) { gemm_4bit_inference_naive(m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize, stream); } void gemm_4bit_inference_naive_bf16( int m, int n, int k, __nv_bfloat16* A, unsigned char* B, float* absmax, float* datatype, __nv_bfloat16* out, int lda, int ldb, int ldc, int blocksize, cudaStream_t stream ) { gemm_4bit_inference_naive<__nv_bfloat16, 16>( m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize, stream ); } void gemm_4bit_inference_naive_fp32( int m, int n, int k, float* A, unsigned char* B, float* absmax, float* datatype, float* out, int lda, int ldb, int ldc, int blocksize, cudaStream_t stream ) { gemm_4bit_inference_naive(m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize, stream); } #define MAKE_ELEMENTWISE_FUNC(fname, type_name, ctype, FUNC) \ void fname##_##type_name(ctype* A, ctype* B, ctype value, long n) { func(A, B, value, n); } MAKE_ELEMENTWISE_FUNC(fill, fp32, float, FILL) MAKE_ELEMENTWISE_FUNC(fill, uint8, unsigned char, FILL) MAKE_ELEMENTWISE_FUNC(arange, fp32, float, ARANGE) MAKE_ELEMENTWISE_FUNC(_mul, fp32, float, _MUL) #define MAKE_FUNC32(fname, oname, gtype, gbits) \ void fname##32bit_grad_##gbits( \ gtype* g, gtype* p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, \ const float beta1, const float beta2, const float beta3, const float alpha, const float eps, \ const float weight_decay, const int step, const float lr, float gnorm_scale, bool skip_zeros, const int n \ ) { \ optimizer32bit( \ g, p, state1, state2, unorm, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, \ lr, gnorm_scale, skip_zeros, n \ ); \ } MAKE_FUNC32(momentum, MOMENTUM, float, 32) MAKE_FUNC32(momentum, MOMENTUM, half, 16) MAKE_FUNC32(adam, ADAM, float, fp32) MAKE_FUNC32(adam, ADAM, half, fp16) MAKE_FUNC32(adam, ADAM, __nv_bfloat16, bf16) MAKE_FUNC32(rmsprop, RMSPROP, float, 32) MAKE_FUNC32(rmsprop, RMSPROP, half, 16) MAKE_FUNC32(lion, LION, float, fp32) MAKE_FUNC32(lion, LION, half, fp16) MAKE_FUNC32(lion, LION, __nv_bfloat16, bf16) MAKE_FUNC32(adagrad, ADAGRAD, float, 32) MAKE_FUNC32(adagrad, ADAGRAD, half, 16) MAKE_FUNC32(ademamix, ADEMAMIX, float, fp32) MAKE_FUNC32(ademamix, ADEMAMIX, half, fp16) MAKE_FUNC32(ademamix, ADEMAMIX, __nv_bfloat16, bf16) #define MAKE_FUNC8(fname, oname, gtype, gbits) \ void fname##_static_8bit_grad_##gbits( \ gtype* p, gtype* g, unsigned char* state1, unsigned char* state2, float* unorm, float max_unorm, \ float param_norm, float beta1, float beta2, float eps, int step, float lr, float* quantiles1, \ float* quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, float weight_decay, \ float gnorm_scale, int n \ ) { \ optimizerStatic8bit( \ g, p, state1, state2, unorm, max_unorm, param_norm, beta1, beta2, eps, step, lr, quantiles1, quantiles2, \ max1, max2, new_max1, new_max2, weight_decay, gnorm_scale, n \ ); \ } MAKE_FUNC8(adam, ADAM, float, 32) MAKE_FUNC8(adam, ADAM, half, 16) MAKE_FUNC8(momentum, MOMENTUM, float, 32) MAKE_FUNC8(momentum, MOMENTUM, half, 16) MAKE_FUNC8(rmsprop, RMSPROP, float, 32) MAKE_FUNC8(rmsprop, RMSPROP, half, 16) MAKE_FUNC8(lion, LION, float, 32) MAKE_FUNC8(lion, LION, half, 16) #define MAKE_BLOCKWISE8(fname, optim_name, gtype, gbits) \ void fname##_8bit_blockwise_grad_##gbits( \ gtype* p, gtype* g, unsigned char* state1, unsigned char* state2, float beta1, float beta2, float beta3, \ float alpha, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* absmax1, \ float* absmax2, float weight_decay, const float gnorm_scale, bool skip_zeros, int n \ ) { \ optimizerStatic8bitBlockwise( \ p, g, state1, state2, beta1, beta2, beta3, alpha, eps, step, lr, quantiles1, quantiles2, absmax1, absmax2, \ weight_decay, gnorm_scale, skip_zeros, n \ ); \ } MAKE_BLOCKWISE8(adam, ADAM, half, fp16) MAKE_BLOCKWISE8(adam, ADAM, __nv_bfloat16, bf16) MAKE_BLOCKWISE8(adam, ADAM, float, fp32) MAKE_BLOCKWISE8(momentum, MOMENTUM, half, fp16) MAKE_BLOCKWISE8(momentum, MOMENTUM, __nv_bfloat16, bf16) MAKE_BLOCKWISE8(momentum, MOMENTUM, float, fp32) MAKE_BLOCKWISE8(rmsprop, RMSPROP, half, fp16) MAKE_BLOCKWISE8(rmsprop, RMSPROP, __nv_bfloat16, bf16) MAKE_BLOCKWISE8(rmsprop, RMSPROP, float, fp32) MAKE_BLOCKWISE8(adagrad, ADAGRAD, half, fp16) MAKE_BLOCKWISE8(adagrad, ADAGRAD, __nv_bfloat16, bf16) MAKE_BLOCKWISE8(adagrad, ADAGRAD, float, fp32) MAKE_BLOCKWISE8(lion, LION, half, fp16) MAKE_BLOCKWISE8(lion, LION, __nv_bfloat16, bf16) MAKE_BLOCKWISE8(lion, LION, float, fp32) MAKE_BLOCKWISE8(ademamix, ADEMAMIX, half, fp16) MAKE_BLOCKWISE8(ademamix, ADEMAMIX, __nv_bfloat16, bf16) MAKE_BLOCKWISE8(ademamix, ADEMAMIX, float, fp32) void percentileClipping_g32(float* g, float* gnorm_vec, int step, const int n) { percentileClipping(g, gnorm_vec, step, n); } void percentileClipping_g16(half* g, float* gnorm_vec, int step, const int n) { percentileClipping(g, gnorm_vec, step, n); } void quantizeBlockwise_fp16(float* code, half* A, float* absmax, unsigned char* out, int blocksize, const int n) { quantizeBlockwise(code, A, absmax, out, NULL, 0, blocksize, n); } void quantizeBlockwise_fp16_fp4(float* code, half* A, float* absmax, unsigned char* out, int blocksize, const int n) { quantizeBlockwise(NULL, A, absmax, out, NULL, 0, blocksize, n); } void quantizeBlockwise_fp16_nf4(float* code, half* A, float* absmax, unsigned char* out, int blocksize, const int n) { quantizeBlockwise(NULL, A, absmax, out, NULL, 0, blocksize, n); } void quantizeBlockwise_bf16( float* code, __nv_bfloat16* A, float* absmax, unsigned char* out, int blocksize, const int n ) { quantizeBlockwise<__nv_bfloat16, 0, General8bit>(code, A, absmax, out, NULL, 0, blocksize, n); } void quantizeBlockwise_bf16_fp4( float* code, __nv_bfloat16* A, float* absmax, unsigned char* out, int blocksize, const int n ) { quantizeBlockwise<__nv_bfloat16, 0, FP4>(NULL, A, absmax, out, NULL, 0, blocksize, n); } void quantizeBlockwise_bf16_nf4( float* code, __nv_bfloat16* A, float* absmax, unsigned char* out, int blocksize, const int n ) { quantizeBlockwise<__nv_bfloat16, 0, NF4>(NULL, A, absmax, out, NULL, 0, blocksize, n); } void quantizeBlockwise_fp32(float* code, float* A, float* absmax, unsigned char* out, int blocksize, const int n) { quantizeBlockwise(code, A, absmax, out, NULL, 0, blocksize, n); } void quantizeBlockwise_fp32_fp4(float* code, float* A, float* absmax, unsigned char* out, int blocksize, const int n) { quantizeBlockwise(NULL, A, absmax, out, NULL, 0, blocksize, n); } void quantizeBlockwise_fp32_nf4(float* code, float* A, float* absmax, unsigned char* out, int blocksize, const int n) { quantizeBlockwise(NULL, A, absmax, out, NULL, 0, blocksize, n); } void dequantizeBlockwise_fp16( float* code, unsigned char* A, float* absmax, half* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise(code, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_fp16_fp4( float* code, unsigned char* A, float* absmax, half* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise(NULL, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_fp16_nf4( float* code, unsigned char* A, float* absmax, half* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise(NULL, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_fp32( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise(code, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_fp32_fp4( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise(NULL, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_fp32_nf4( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise(NULL, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_bf16( float* code, unsigned char* A, float* absmax, __nv_bfloat16* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise<__nv_bfloat16, General8bit>(code, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_bf16_fp4( float* code, unsigned char* A, float* absmax, __nv_bfloat16* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise<__nv_bfloat16, FP4>(NULL, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_bf16_nf4( float* code, unsigned char* A, float* absmax, __nv_bfloat16* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise<__nv_bfloat16, NF4>(NULL, A, absmax, out, blocksize, n, stream); } int igemmlt_32( cublasLtHandle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, cudaStream_t stream ) { return igemmlt<32, 0>(ltHandle, m, n, k, A, B, C, row_scale, lda, ldb, ldc, stream); } int igemmlt_8( cublasLtHandle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, cudaStream_t stream ) { return igemmlt<8, 0>(ltHandle, m, n, k, A, B, C, row_scale, lda, ldb, ldc, stream); } int igemmlt_8_rowscale( cublasLtHandle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, cudaStream_t stream ) { return igemmlt<8, 1>(ltHandle, m, n, k, A, B, C, row_scale, lda, ldb, ldc, stream); } void spmm_coo_very_sparse_naive_fp16( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, half* B, half* out, float* dequant_stats, int nnz_rows, int nnz, int rowsA, int rowsB, int colsB ) { spmm_coo_very_sparse_naive( max_count, max_idx, offset_rowidx, rowidx, colidx, values, B, out, dequant_stats, nnz_rows, nnz, rowsA, rowsB, colsB ); } void spmm_coo_very_sparse_naive_int8( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, signed char* B, half* out, float* dequant_stats, int nnz_rows, int nnz, int rowsA, int rowsB, int colsB ) { spmm_coo_very_sparse_naive( max_count, max_idx, offset_rowidx, rowidx, colidx, values, B, out, dequant_stats, nnz_rows, nnz, rowsA, rowsB, colsB ); } #endif #if BUILD_XPU void dequantizeBlockwise_fp16( float* code, unsigned char* A, float* absmax, sycl::half* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise(code, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_fp16_fp4( float* code, unsigned char* A, float* absmax, sycl::half* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise(NULL, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_fp16_nf4( float* code, unsigned char* A, float* absmax, sycl::half* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise(NULL, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_fp32( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise(code, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_fp32_fp4( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise(NULL, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_fp32_nf4( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise(NULL, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_bf16( float* code, unsigned char* A, float* absmax, sycl::ext::oneapi::bfloat16* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise(code, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_bf16_fp4( float* code, unsigned char* A, float* absmax, sycl::ext::oneapi::bfloat16* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise(NULL, A, absmax, out, blocksize, n, stream); } void dequantizeBlockwise_bf16_nf4( float* code, unsigned char* A, float* absmax, sycl::ext::oneapi::bfloat16* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise(NULL, A, absmax, out, blocksize, n, stream); } void gemv_4bit_inference_fp16( int m, int n, int k, sycl::half* A, unsigned char* B, float* absmax, float* datatype, sycl::half* out, int lda, int ldb, int ldc, int blocksize, sycl::queue* stream ) { gemv_4bit_inference(m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize, stream); } void gemv_4bit_inference_bf16( int m, int n, int k, sycl::ext::oneapi::bfloat16* A, unsigned char* B, float* absmax, float* datatype, sycl::ext::oneapi::bfloat16* out, int lda, int ldb, int ldc, int blocksize, sycl::queue* stream ) { gemv_4bit_inference( m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize, stream ); } void gemv_4bit_inference_fp32( int m, int n, int k, float* A, unsigned char* B, float* absmax, float* datatype, float* out, int lda, int ldb, int ldc, int blocksize, sycl::queue* stream ) { gemv_4bit_inference(m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize, stream); } #endif extern "C" { #if BUILD_CUDA || BUILD_HIP void cquantize(float* code, float* A, unsigned char* out, int n) { quantize(code, A, out, n); } void cdequantize(float* code, unsigned char* A, float* out, int n, cudaStream_t stream) { dequantize(code, A, out, n, stream); } void cdequantize_blockwise_fp16_fp4( float* code, unsigned char* A, float* absmax, half* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise_fp16_fp4(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_fp16( float* code, unsigned char* A, float* absmax, half* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise_fp16(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_fp16_nf4( float* code, unsigned char* A, float* absmax, half* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise_fp16_nf4(code, A, absmax, out, blocksize, n, stream); } void cquantize_blockwise_fp16(float* code, half* A, float* absmax, unsigned char* out, int blocksize, const int n) { quantizeBlockwise_fp16(code, A, absmax, out, blocksize, n); } void cquantize_blockwise_fp16_fp4(float* code, half* A, float* absmax, unsigned char* out, int blocksize, const int n) { quantizeBlockwise_fp16_fp4(code, A, absmax, out, blocksize, n); } void cquantize_blockwise_fp16_nf4(float* code, half* A, float* absmax, unsigned char* out, int blocksize, const int n) { quantizeBlockwise_fp16_nf4(code, A, absmax, out, blocksize, n); } void cquantize_blockwise_fp32(float* code, float* A, float* absmax, unsigned char* out, int blocksize, const int n) { quantizeBlockwise_fp32(code, A, absmax, out, blocksize, n); } void cquantize_blockwise_fp32_fp4( float* code, float* A, float* absmax, unsigned char* out, int blocksize, const int n ) { quantizeBlockwise_fp32_fp4(code, A, absmax, out, blocksize, n); } void cquantize_blockwise_fp32_nf4( float* code, float* A, float* absmax, unsigned char* out, int blocksize, const int n ) { quantizeBlockwise_fp32_nf4(code, A, absmax, out, blocksize, n); } void cdequantize_blockwise_fp32( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise_fp32(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_fp32_fp4( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise_fp32_fp4(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_fp32_nf4( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise_fp32_nf4(code, A, absmax, out, blocksize, n, stream); } void cquantize_blockwise_bf16( float* code, __nv_bfloat16* A, float* absmax, unsigned char* out, int blocksize, const int n ) { quantizeBlockwise_bf16(code, A, absmax, out, blocksize, n); } void cquantize_blockwise_bf16_fp4( float* code, __nv_bfloat16* A, float* absmax, unsigned char* out, int blocksize, const int n ) { quantizeBlockwise_bf16_fp4(code, A, absmax, out, blocksize, n); } void cquantize_blockwise_bf16_nf4( float* code, __nv_bfloat16* A, float* absmax, unsigned char* out, int blocksize, const int n ) { quantizeBlockwise_bf16_nf4(code, A, absmax, out, blocksize, n); } void cdequantize_blockwise_bf16( float* code, unsigned char* A, float* absmax, __nv_bfloat16* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise_bf16(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_bf16_fp4( float* code, unsigned char* A, float* absmax, __nv_bfloat16* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise_bf16_fp4(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_bf16_nf4( float* code, unsigned char* A, float* absmax, __nv_bfloat16* out, int blocksize, const int n, cudaStream_t stream ) { dequantizeBlockwise_bf16_nf4(code, A, absmax, out, blocksize, n, stream); } #define MAKE_CFUNC32(name, gtype, gbits) \ void c##name##32bit_grad_##gbits( \ gtype* g, gtype* p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, \ const float beta1, const float beta2, const float beta3, const float alpha, const float eps, \ const float weight_decay, const int step, const float lr, const float gnorm_scale, bool skip_zeros, \ const int n \ ) { \ name##32bit_grad_##gbits( \ g, p, state1, state2, unorm, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, \ lr, gnorm_scale, skip_zeros, n \ ); \ } MAKE_CFUNC32(adam, float, fp32) MAKE_CFUNC32(adam, half, fp16) MAKE_CFUNC32(adam, __nv_bfloat16, bf16) MAKE_CFUNC32(momentum, float, 32) MAKE_CFUNC32(momentum, half, 16) MAKE_CFUNC32(rmsprop, float, 32) MAKE_CFUNC32(rmsprop, half, 16) MAKE_CFUNC32(lion, float, fp32) MAKE_CFUNC32(lion, half, fp16) MAKE_CFUNC32(lion, __nv_bfloat16, bf16) MAKE_CFUNC32(adagrad, float, 32) MAKE_CFUNC32(adagrad, half, 16) MAKE_CFUNC32(ademamix, float, fp32) MAKE_CFUNC32(ademamix, half, fp16) MAKE_CFUNC32(ademamix, __nv_bfloat16, bf16) #define MAKE_CFUNC8(name, gtype, gbits) \ void c##name##_static_8bit_grad_##gbits( \ gtype* p, gtype* g, unsigned char* state1, unsigned char* state2, float* unorm, float max_unorm, \ float param_norm, float beta1, float beta2, float eps, int step, float lr, float* quantiles1, \ float* quantiles2, float* max1, float* max2, float* new_max1, float* new_max2, float weight_decay, \ float gnorm_scale, int n \ ) { \ name##_static_8bit_grad_##gbits( \ g, p, state1, state2, unorm, max_unorm, param_norm, beta1, beta2, eps, step, lr, quantiles1, quantiles2, \ max1, max2, new_max1, new_max2, weight_decay, gnorm_scale, n \ ); \ } MAKE_CFUNC8(adam, float, 32) MAKE_CFUNC8(adam, half, 16) MAKE_CFUNC8(momentum, float, 32) MAKE_CFUNC8(momentum, half, 16) MAKE_CFUNC8(rmsprop, float, 32) MAKE_CFUNC8(rmsprop, half, 16) MAKE_CFUNC8(lion, float, 32) MAKE_CFUNC8(lion, half, 16) #define MAKE_CBLOCKWISE8(fname, optim_name, gtype, gbits) \ void c##fname##_8bit_blockwise_grad_##gbits( \ gtype* p, gtype* g, unsigned char* state1, unsigned char* state2, float beta1, float beta2, float beta3, \ float alpha, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* absmax1, \ float* absmax2, float weight_decay, const float gnorm_scale, bool skip_zeros, int n \ ) { \ fname##_8bit_blockwise_grad_##gbits( \ p, g, state1, state2, beta1, beta2, beta3, alpha, eps, step, lr, quantiles1, quantiles2, absmax1, absmax2, \ weight_decay, gnorm_scale, skip_zeros, n \ ); \ } MAKE_CBLOCKWISE8(adam, ADAM, half, fp16) MAKE_CBLOCKWISE8(adam, ADAM, float, fp32) MAKE_CBLOCKWISE8(adam, ADAM, __nv_bfloat16, bf16) MAKE_CBLOCKWISE8(momentum, MOMENTUM, half, fp16) MAKE_CBLOCKWISE8(momentum, MOMENTUM, float, fp32) MAKE_CBLOCKWISE8(momentum, MOMENTUM, __nv_bfloat16, bf16) MAKE_CBLOCKWISE8(rmsprop, RMSPROP, half, fp16) MAKE_CBLOCKWISE8(rmsprop, RMSPROP, float, fp32) MAKE_CBLOCKWISE8(rmsprop, RMSPROP, __nv_bfloat16, bf16) MAKE_CBLOCKWISE8(adagrad, ADAGRAD, half, fp16) MAKE_CBLOCKWISE8(adagrad, ADAGRAD, float, fp32) MAKE_CBLOCKWISE8(adagrad, ADAGRAD, __nv_bfloat16, bf16) MAKE_CBLOCKWISE8(lion, LION, half, fp16) MAKE_CBLOCKWISE8(lion, LION, float, fp32) MAKE_CBLOCKWISE8(lion, LION, __nv_bfloat16, bf16) MAKE_CBLOCKWISE8(ademamix, ADEMAMIX, half, fp16) MAKE_CBLOCKWISE8(ademamix, ADEMAMIX, float, fp32) MAKE_CBLOCKWISE8(ademamix, ADEMAMIX, __nv_bfloat16, bf16) void cpercentile_clipping_g32(float* g, float* gnorm_vec, int step, const int n) { percentileClipping_g32(g, gnorm_vec, step, n); } void cpercentile_clipping_g16(half* g, float* gnorm_vec, int step, const int n) { percentileClipping_g16(g, gnorm_vec, step, n); } void cigemm( Context* context, bool transposeA, bool transposeB, int m, int n, int k, void* A, void* B, void* C, int lda, int ldb, int ldc ) { gemmex(context, transposeA, transposeB, m, n, k, A, B, C, lda, ldb, ldc); } void cbatched_igemm( Context* context, bool transposeA, bool transposeB, int m, int n, int k, void* A, void* B, void* C, int lda, int ldb, int ldc, long strideA, long strideB, long strideC, int batchCount ) { strided_gemmex( context, transposeA, transposeB, m, n, k, A, B, C, lda, ldb, ldc, strideA, strideB, strideC, batchCount ); } Context* get_context() { return new Context(); } ContextCusparse* get_cusparse() { return new ContextCusparse(); } int cigemmlt_32( Context* context, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, cudaStream_t stream ) { return igemmlt_32((cublasLtHandle_t)context->m_handle, m, n, k, A, B, C, row_scale, lda, ldb, ldc, stream); } int cigemmlt_8( Context* context, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, cudaStream_t stream ) { return igemmlt_8((cublasLtHandle_t)context->m_handle, m, n, k, A, B, C, row_scale, lda, ldb, ldc, stream); } int cigemmlt_8_rowscale( Context* context, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, cudaStream_t stream ) { return igemmlt_8_rowscale((cublasLtHandle_t)context->m_handle, m, n, k, A, B, C, row_scale, lda, ldb, ldc, stream); } void cdequant_mm_int32_fp16( int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, cudaStream_t stream ) { dequant_mm_int32_fp16(A, rowStats, colStats, out, bias, numRows, numCols, stream); } void cget_row_stats(half* A, float* rowStats, float threshold, int rows, int cols, cudaStream_t stream) { getRowStats(A, rowStats, threshold, rows, cols, stream); } void cint8_vector_quant( half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, cudaStream_t stream ) { int8VectorQuant(A, out, rowStats, threshold, rows, cols, stream); } void cspmm_coo( ContextCusparse* context, int* A_rowidx, int* A_colidx, half* A_vals, int A_nnz, int A_rows, int A_cols, int B_cols, int ldb, half* B, int ldc, half* C, bool transposed_B ) { spmm_coo( (cusparseHandle_t)context->m_handle, A_rowidx, A_colidx, A_vals, A_nnz, A_rows, A_cols, B_cols, ldb, B, ldc, C, transposed_B ); } void cspmm_coo_very_sparse_naive_fp16( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, half* B, half* out, float* dequant_stats, int nnz_rows, int nnz, int rowsA, int rowsB, int colsB ) { spmm_coo_very_sparse_naive_fp16( max_count, max_idx, offset_rowidx, rowidx, colidx, values, B, out, dequant_stats, nnz_rows, nnz, rowsA, rowsB, colsB ); } void cspmm_coo_very_sparse_naive_int8( int* max_count, int* max_idx, int* offset_rowidx, int* rowidx, int* colidx, half* values, signed char* B, half* out, float* dequant_stats, int nnz_rows, int nnz, int rowsA, int rowsB, int colsB ) { spmm_coo_very_sparse_naive_int8( max_count, max_idx, offset_rowidx, rowidx, colidx, values, B, out, dequant_stats, nnz_rows, nnz, rowsA, rowsB, colsB ); } // void cgemm_host_fp32(int M, int N, int K, float * A, float* B, float * out, int lda, int ldb, int ldc) //{ gemm_host_fp32(M, N, K, A, B, out, lda, ldb, ldc); } void cgemm_host_fp16(int M, int N, int K, half* A, half* B, half* out, int lda, int ldb, int ldc) { gemm_host_fp16(M, N, K, A, B, out, lda, ldb, ldc); } void cgemm_4bit_inference( int m, int n, int k, half* A, unsigned char* B, float* absmax, half* out, int lda, int ldb, int ldc, int blocksize ) { gemm_4bit_inference(m, n, k, A, B, absmax, out, lda, ldb, ldc, blocksize); } void* cget_managed_ptr(size_t bytes) { void* ptr; CUDA_CHECK_RETURN(cudaMallocManaged(&ptr, bytes, cudaMemAttachHost)); CUDA_CHECK_RETURN(cudaPeekAtLastError()); return ptr; } void cprefetch(void* ptr, size_t bytes, int device) { int hasPrefetch = 0; CUDA_CHECK_RETURN( cudaDeviceGetAttribute(&hasPrefetch, cudaDevAttrConcurrentManagedAccess, device) ); // 40ns overhead if (hasPrefetch == 0) return; #if CUDART_VERSION >= 13000 cudaMemLocation loc{}; loc.type = cudaMemLocationTypeDevice; loc.id = device; CUDA_CHECK_RETURN(cudaMemPrefetchAsync(ptr, bytes, loc, 0u, 0)); #else CUDA_CHECK_RETURN(cudaMemPrefetchAsync(ptr, bytes, device, 0)); #endif CUDA_CHECK_RETURN(cudaPeekAtLastError()); } #define CMAKE_ELEMENTWISE_FUNC(fname, type_name, ctype, FUNC) \ void c##fname##_##type_name(ctype* A, ctype* B, ctype value, long n) { fname##_##type_name(A, B, value, n); } CMAKE_ELEMENTWISE_FUNC(fill, fp32, float, FILL) CMAKE_ELEMENTWISE_FUNC(fill, uint8, unsigned char, FILL) CMAKE_ELEMENTWISE_FUNC(arange, fp32, float, ARANGE) CMAKE_ELEMENTWISE_FUNC(_mul, fp32, float, _MUL) void cgemm_4bit_inference_naive_fp16( int m, int n, int k, half* A, unsigned char* B, float* absmax, float* datatype, half* out, int lda, int ldb, int ldc, int blocksize, cudaStream_t stream ) { gemm_4bit_inference_naive_fp16(m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize, stream); } void cgemm_4bit_inference_naive_bf16( int m, int n, int k, __nv_bfloat16* A, unsigned char* B, float* absmax, float* datatype, __nv_bfloat16* out, int lda, int ldb, int ldc, int blocksize, cudaStream_t stream ) { gemm_4bit_inference_naive_bf16(m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize, stream); } void cgemm_4bit_inference_naive_fp32( int m, int n, int k, float* A, unsigned char* B, float* absmax, float* datatype, float* out, int lda, int ldb, int ldc, int blocksize, cudaStream_t stream ) { gemm_4bit_inference_naive_fp32(m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize, stream); } #endif #if BUILD_XPU void cdequantize_blockwise_fp16_fp4( float* code, unsigned char* A, float* absmax, sycl::half* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise_fp16_fp4(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_fp16( float* code, unsigned char* A, float* absmax, sycl::half* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise_fp16(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_fp16_nf4( float* code, unsigned char* A, float* absmax, sycl::half* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise_fp16_nf4(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_fp32( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise_fp32(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_fp32_fp4( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise_fp32_fp4(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_fp32_nf4( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise_fp32_nf4(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_bf16( float* code, unsigned char* A, float* absmax, sycl::ext::oneapi::bfloat16* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise_bf16(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_bf16_fp4( float* code, unsigned char* A, float* absmax, sycl::ext::oneapi::bfloat16* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise_bf16_fp4(code, A, absmax, out, blocksize, n, stream); } void cdequantize_blockwise_bf16_nf4( float* code, unsigned char* A, float* absmax, sycl::ext::oneapi::bfloat16* out, int blocksize, const int n, sycl::queue* stream ) { dequantizeBlockwise_bf16_nf4(code, A, absmax, out, blocksize, n, stream); } void cgemv_4bit_inference_fp16( int m, int n, int k, sycl::half* A, unsigned char* B, float* absmax, float* datatype, sycl::half* out, int lda, int ldb, int ldc, int blocksize, sycl::queue* stream ) { gemv_4bit_inference_fp16(m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize, stream); } void cgemv_4bit_inference_bf16( int m, int n, int k, sycl::ext::oneapi::bfloat16* A, unsigned char* B, float* absmax, float* datatype, sycl::ext::oneapi::bfloat16* out, int lda, int ldb, int ldc, int blocksize, sycl::queue* stream ) { gemv_4bit_inference_bf16(m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize, stream); } void cgemv_4bit_inference_fp32( int m, int n, int k, float* A, unsigned char* B, float* absmax, float* datatype, float* out, int lda, int ldb, int ldc, int blocksize, sycl::queue* stream ) { gemv_4bit_inference_fp32(m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize, stream); } #endif void cquantize_blockwise_cpu_fp32( float* code, float* A, float* absmax, unsigned char* out, long long blocksize, long long n ) { quantize_cpu(code, A, absmax, out, blocksize, n); } void cdequantize_blockwise_cpu_fp32( float* code, unsigned char* A, float* absmax, float* out, long long blocksize, long long n ) { dequantize_cpu(code, A, absmax, out, blocksize, n); } } bitsandbytes/csrc/xpu_kernels.cpp000066400000000000000000000254011506752665600175720ustar00rootroot00000000000000#include "xpu_kernels.h" #include #include #include #include inline float dDequantizeFP4(unsigned char val) { if ((val & 0b1000) == 8) if ((val & 0b0100) == 4) if ((val & 0b0010) == 2) if ((val & 0b0001) == 1) return -0.25000000f; else return -0.16666667f; else if ((val & 0b0001) == 1) return -0.50000000f; else return -0.33333333f; else if ((val & 0b0010) == 2) if ((val & 0b0001) == 1) return -1.00000000f; else return -0.66666667f; else if ((val & 0b0001) == 1) return -5.208333333e-03f; else return 0.00000000f; else if ((val & 0b0100) == 4) if ((val & 0b0010) == 2) if ((val & 0b0001) == 1) return 0.25000000f; else return 0.16666667f; else if ((val & 0b0001) == 1) return 0.50000000f; else return 0.33333333f; else if ((val & 0b0010) == 2) if ((val & 0b0001) == 1) return 1.00000000f; else return 0.66666667f; else if ((val & 0b0001) == 1) return 5.208333333e-03f; else return 0.00000000f; } inline float dDequantizeNF4(unsigned char val) { // the values for this tree was generated by test_normal_map_tree // in the file tests/test_functional.py if ((val & 0b1000) == 8) if ((val & 0b0100) == 4) // 1 if ((val & 0b0010) == 2) // 11 if ((val & 0b0001) == 1) // 111 return 1.0f; //*1111 else return 0.7229568362236023f; //*1110 else if ((val & 0b0001) == 1) // 110 return 0.5626170039176941f; //*1101 else return 0.44070982933044434f; //*1100 else if ((val & 0b0010) == 2) // 10 if ((val & 0b0001) == 1) // 101 return 0.33791524171829224f; //*1011 else return 0.24611230194568634f; //*1010 else if ((val & 0b0001) == 1) // 100 return 0.16093020141124725f; //*1001 else return 0.07958029955625534f; //*1000 else if ((val & 0b0100) == 4) // 0 if ((val & 0b0010) == 2) // 01 if ((val & 0b0001) == 1) // 011 return 0.0f; //*0111 else return -0.09105003625154495f; //*0110 else if ((val & 0b0001) == 1) // 010 return -0.18477343022823334f; //*0101 else return -0.28444138169288635f; //*0100 else if ((val & 0b0010) == 2) // 00 if ((val & 0b0001) == 1) // 001 return -0.39491748809814453f; //*0011 else return -0.5250730514526367f; //*0010 else if ((val & 0b0001) == 1) // 000 return -0.6961928009986877f; //*0001 else return -1.0f; //*0000 } template SYCL_EXTERNAL void kDequantizeBlockwise::operator()(sycl::nd_item<1> item) const { const int base_idx = item.get_group(0) * TILE_SIZE; size_t local_idx = item.get_local_id(0) * NUM_PER_TH; float local_abs_max = -FLT_MAX; int local_load_idx = 0; int local_store_idx = 0; uint8_t qvals[NUM_PER_TH]; T vals[NUM_PER_TH * ((DATA_TYPE > 0) ? 2 : 1)]; if (DATA_TYPE > 0) { local_load_idx = sycl::min(TILE_SIZE, (n + 1) / 2 - base_idx); local_store_idx = sycl::min(TILE_SIZE * 2, n - base_idx * 2); } else { local_load_idx = sycl::min(TILE_SIZE, n - base_idx); local_store_idx = local_load_idx; } // Avoid expensive division by the blocksize (as blocksize will always be a // power-of-2) local_abs_max = absmax[(base_idx + local_idx) >> (31 - std::countl_zero(blocksize))]; if (local_idx + NUM_PER_TH < local_load_idx) { reinterpret_cast(&)[NUM_PER_TH]>(qvals)[0] = reinterpret_cast*>(A)[(base_idx + local_idx) / NUM_PER_TH]; } else { #pragma unroll NUM_PER_TH for (int i = 0; i < NUM_PER_TH; i++) { if (local_idx + i < local_load_idx) { qvals[i] = A[base_idx + local_idx + i]; } else { qvals[i] = (uint8_t)0; } } } switch (DATA_TYPE) { case General8bit: #pragma unroll NUM_PER_TH for (int j = 0; j < NUM_PER_TH; j++) vals[j] = code[qvals[j]] * local_abs_max; break; case FP4: #pragma unroll NUM_PER_TH for (int j = 0; j < NUM_PER_TH; j++) { vals[j * 2] = dDequantizeFP4(qvals[j] >> 4) * local_abs_max; vals[j * 2 + 1] = dDequantizeFP4(qvals[j] & 0x0F) * local_abs_max; } break; case NF4: #pragma unroll NUM_PER_TH for (int j = 0; j < NUM_PER_TH; j++) { vals[j * 2] = dDequantizeNF4(qvals[j] >> 4) * local_abs_max; vals[j * 2 + 1] = dDequantizeNF4(qvals[j] & 0x0F) * local_abs_max; } break; } const int local_dst_size = (DATA_TYPE > 0) ? NUM_PER_TH * 2 : NUM_PER_TH; int local_dst_idx = (DATA_TYPE > 0) ? local_idx * 2 : local_idx; if (local_dst_idx + local_dst_size < local_store_idx) { reinterpret_cast*>( out )[(((DATA_TYPE > 0) ? base_idx * 2 : base_idx) + local_dst_idx) / local_dst_size] = reinterpret_cast(&)[local_dst_size]>(vals)[0]; } else { #pragma unroll NUM_PER_TH for (int i = 0; i < local_dst_size; i++) { if (local_dst_idx + i < local_store_idx) { out[((DATA_TYPE > 0) ? base_idx * 2 : base_idx) + local_dst_idx + i] = vals[i]; } } } } template SYCL_EXTERNAL void kgemv_4bit_inference::operator()(sycl::nd_item<1> item) const { size_t idx = item.get_local_id(); const int sg_idx = idx / SUBG_SIZE; const int sg_lane = idx % SUBG_SIZE; const int num_values_4bit = SUBG_SIZE; const int row_B = NUM_PER_THREAD * item.get_group().get_group_id() + sg_idx; const int offset_B = ldb * row_B; const int num_values_8bit = num_values_4bit / 2; float local_C = 0.0f; unsigned char local_B_4bit[num_values_8bit]; T local_B[num_values_4bit / 4]; T local_A[num_values_4bit / 4]; T local_absmax = T(0.0f); if (idx < 16) { quant_map[idx] = T(datatype[idx]); } item.barrier(sycl::access::fence_space::local_space); for (int inner_idx = sg_lane * num_values_4bit; inner_idx < K; inner_idx += SUBG_SIZE * num_values_4bit) { const int inner_idx_halved = inner_idx / 2; // Avoid expensive division by the blocksize (as blocksize will always be a // power-of-2) const int absidx = ((2 * offset_B) + inner_idx) >> (31 - std::countl_zero((unsigned int)blocksize)); local_absmax = absmax[absidx]; if (row_B < N) { if ((inner_idx_halved + num_values_8bit) < (K / 2)) { reinterpret_cast(&)[num_values_8bit]>(local_B_4bit)[0] = reinterpret_cast*>(B)[(offset_B + (inner_idx_halved)) / (num_values_8bit)]; } else { #pragma unroll for (int j = 0; j < (num_values_8bit); j++) if ((inner_idx_halved) + j < (K / 2)) local_B_4bit[j] = B[offset_B + inner_idx_halved + j]; else local_B_4bit[j] = 0b01110111; } } else { #pragma unroll for (int j = 0; j < (num_values_8bit); j++) local_B_4bit[j] = 0b01110111; } for (int i = 0; i < 4; i++) { #pragma unroll for (int k = 0; k < num_values_8bit / 4; k++) { local_B[k * 2] = quant_map[local_B_4bit[(i * num_values_8bit / 4) + k] >> 4] * local_absmax; local_B[k * 2 + 1] = quant_map[local_B_4bit[(i * num_values_8bit / 4) + k] & 0x0F] * local_absmax; } if (inner_idx + (num_values_4bit / 4) + (i * num_values_4bit / 4) < K) { if (BITS == 16) { reinterpret_cast(&)[num_values_4bit / 4]>(local_A)[0] = reinterpret_cast*>(A)[inner_idx / (num_values_4bit / 4) + i]; } else { reinterpret_cast(&)[num_values_4bit / 4]>(local_A)[0] = reinterpret_cast*>(A)[inner_idx / (num_values_4bit / 8) + (2 * i) + 0]; reinterpret_cast(&)[num_values_4bit / 4]>(local_A)[1] = reinterpret_cast*>(A)[inner_idx / (num_values_4bit / 8) + (2 * i) + 1]; } } else { #pragma unroll for (int k = 0; k < num_values_4bit / 4; k++) if (inner_idx + (i * num_values_4bit / 4) + k < K) local_A[k] = A[inner_idx + k + (i * num_values_4bit / 4)]; else local_A[k] = T(0.0f); } // accumulate in float for accuracy; #pragma unroll for (int k = 0; k < num_values_4bit / 4; k++) { local_C += (float)(local_A[k] * local_B[k]); } } } local_C = sycl::reduce_over_group(item.get_sub_group(), local_C, sycl::plus<>()); if (row_B < N && sg_lane == 0) out[row_B] = T(local_C); } //============================================================== // TEMPLATE DEFINITIONS //============================================================== template class kDequantizeBlockwise; template class kDequantizeBlockwise; template class kDequantizeBlockwise; template class kDequantizeBlockwise; template class kDequantizeBlockwise; template class kDequantizeBlockwise; template class kDequantizeBlockwise; template class kDequantizeBlockwise; template class kDequantizeBlockwise; template class kgemv_4bit_inference; template class kgemv_4bit_inference; template class kgemv_4bit_inference; bitsandbytes/csrc/xpu_kernels.h000066400000000000000000000027301506752665600172370ustar00rootroot00000000000000#include #include #ifndef xpu_kernels #define xpu_kernels template class kDequantizeBlockwise { public: SYCL_EXTERNAL void operator()(sycl::nd_item<1> item) const; kDequantizeBlockwise(float* code_, uint8_t* A_, float* absmax_, T* out_, const int blocksize_, const int n_) : code(code_), A(A_), absmax(absmax_), out(out_), blocksize(blocksize_), n(n_) {} private: float* code; uint8_t* A; float* absmax; T* out; const int blocksize; const int n; }; template class kgemv_4bit_inference { public: SYCL_EXTERNAL void operator()(sycl::nd_item<1> item) const; kgemv_4bit_inference( int M_, int N_, int K_, T* A_, unsigned char* B_, float* absmax_, const float* datatype_, T* out_, int lda_, int ldb_, int ldc_, int blocksize_ ) : M(M_), N(N_), K(K_), A(A_), B(B_), absmax(absmax_), datatype(datatype_), out(out_), lda(lda_), ldb(ldb_), ldc(ldc_), blocksize(blocksize_), quant_map() {} void sycl_ker_local_memory_creation(sycl::handler& cgh) { quant_map = sycl::local_accessor(16, cgh); } private: int M; int N; int K; T* A; unsigned char* B; float* absmax; const float* datatype; T* out; int lda; int ldb; int ldc; int blocksize; sycl::local_accessor quant_map; }; #endif bitsandbytes/csrc/xpu_ops.cpp000066400000000000000000000112251506752665600167270ustar00rootroot00000000000000#include #include #include template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, T* out, int blocksize, const int n, sycl::queue* stream ) { auto& queue = *stream; const int workgroup_size = 128; const int num_per_th = 4; const int tile_size = workgroup_size * num_per_th; if (DATA_TYPE > 0) { const int workgroup_num = (n + tile_size * 2 - 1) / (tile_size * 2); sycl::range<1> local_range{(size_t)workgroup_size}; sycl::range<1> global_range{(size_t)workgroup_num * (size_t)workgroup_size}; kDequantizeBlockwise kfn(code, A, absmax, out, blocksize / 2, n); sycl_kernel_submit( sycl::nd_range<1>(sycl::range<1>(global_range), sycl::range<1>(local_range)), queue, kfn ); } else { const int workgroup_num = (n + tile_size - 1) / tile_size; sycl::range<1> local_range{(size_t)workgroup_size}; sycl::range<1> global_range{(size_t)workgroup_num * (size_t)workgroup_size}; kDequantizeBlockwise kfn(code, A, absmax, out, blocksize, n); sycl_kernel_submit( sycl::nd_range<1>(sycl::range<1>(global_range), sycl::range<1>(local_range)), queue, kfn ); } } template void gemv_4bit_inference( int m, int n, int k, T* A, unsigned char* B, float* absmax, float* datatype, T* out, int lda, int ldb, int ldc, int blocksize, sycl::queue* stream ) { auto& queue = *stream; const size_t GROUP_SIZE = 128; // workgroup_size const size_t SUBG_SIZE = 32; // subgroup_size const size_t NUM_PER_THREAD = GROUP_SIZE / SUBG_SIZE; size_t workgroup_num = (n + NUM_PER_THREAD - 1) / NUM_PER_THREAD; kgemv_4bit_inference kfn( m, n, k, A, B, absmax, datatype, out, lda, ldb, ldc, blocksize ); sycl_comp_kernel_submit( sycl::nd_range<1>(sycl::range<1>(GROUP_SIZE * workgroup_num), sycl::range<1>(GROUP_SIZE)), queue, kfn ); } //============================================================== // TEMPLATE DEFINITIONS //============================================================== template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, sycl::queue* stream ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, sycl::queue* stream ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, float* out, int blocksize, const int n, sycl::queue* stream ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, sycl::half* out, int blocksize, const int n, sycl::queue* stream ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, sycl::half* out, int blocksize, const int n, sycl::queue* stream ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, sycl::half* out, int blocksize, const int n, sycl::queue* stream ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, sycl::ext::oneapi::bfloat16* out, int blocksize, const int n, sycl::queue* stream ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, sycl::ext::oneapi::bfloat16* out, int blocksize, const int n, sycl::queue* stream ); template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, sycl::ext::oneapi::bfloat16* out, int blocksize, const int n, sycl::queue* stream ); template void gemv_4bit_inference( int m, int n, int k, sycl::half* A, unsigned char* B, float* absmax, float* datatype, sycl::half* out, int lda, int ldb, int ldc, int blocksize, sycl::queue* stream ); template void gemv_4bit_inference( int m, int n, int k, sycl::ext::oneapi::bfloat16* A, unsigned char* B, float* absmax, float* datatype, sycl::ext::oneapi::bfloat16* out, int lda, int ldb, int ldc, int blocksize, sycl::queue* stream ); template void gemv_4bit_inference( int m, int n, int k, float* A, unsigned char* B, float* absmax, float* datatype, float* out, int lda, int ldb, int ldc, int blocksize, sycl::queue* stream ); bitsandbytes/csrc/xpu_ops.h000066400000000000000000000025071506752665600163770ustar00rootroot00000000000000#ifndef xpu_ops_H #define xpu_ops_H #include #include #include #include #include #include #include template static inline void sycl_kernel_submit(sycl::nd_range range, sycl::queue q, ker_t ker) { auto cgf = [&](::sycl::handler& cgh) [[sycl::reqd_sub_group_size(subgroup_size)]] { cgh.parallel_for(range, ker); }; q.submit(cgf); } template static inline void sycl_comp_kernel_submit(sycl::nd_range range, sycl::queue q, ker_t ker) { auto cgf = [&](::sycl::handler& cgh) [[sycl::reqd_sub_group_size(subgroup_size)]] { ker.sycl_ker_local_memory_creation(cgh); cgh.parallel_for(range, ker); }; q.submit(cgf); } typedef enum DataType_t { General8bit = 0, FP4 = 1, NF4 = 2, } DataType_t; template void dequantizeBlockwise( float* code, unsigned char* A, float* absmax, T* out, int workgroup_size, const int n, sycl::queue* stream ); template void gemv_4bit_inference( int m, int n, int k, T* A, unsigned char* B, float* absmax, float* datatype, T* out, int lda, int ldb, int ldc, int blocksize, sycl::queue* stream ); #endif bitsandbytes/docs/000077500000000000000000000000001506752665600145235ustar00rootroot00000000000000bitsandbytes/docs/source/000077500000000000000000000000001506752665600160235ustar00rootroot00000000000000bitsandbytes/docs/source/_toctree.yml000066400000000000000000000027771506752665600203670ustar00rootroot00000000000000- title: Get started sections: - local: index title: bitsandbytes - local: installation title: Installation - local: quickstart title: Quickstart - title: Usage Guides sections: - local: optimizers title: 8-bit optimizers - local: fsdp_qlora title: FSDP-QLoRA - local: integrations title: Integrations - local: errors title: Troubleshoot - local: contributing title: Contribute - local: faqs title: FAQs - title: Explanation sections: - local: explanations/optimizers title: 8-bit optimizers - local: explanations/resources title: Papers, resources & how to cite - title: API reference sections: - title: Functional local: reference/functional - title: Optimizers sections: - local: reference/optim/optim_overview title: Overview - local: reference/optim/adagrad title: AdaGrad - local: reference/optim/adam title: Adam - local: reference/optim/adamw title: AdamW - local: reference/optim/ademamix title: AdEMAMix - local: reference/optim/lamb title: LAMB - local: reference/optim/lars title: LARS - local: reference/optim/lion title: Lion - local: reference/optim/rmsprop title: RMSprop - local: reference/optim/sgd title: SGD - title: Modules sections: - local: reference/nn/linear8bit title: LLM.int8() - local: reference/nn/linear4bit title: 4-bit quantizer - local: reference/nn/embeddings title: Embedding bitsandbytes/docs/source/contributing.mdx000066400000000000000000000030011506752665600212360ustar00rootroot00000000000000# Contribution Guide ## Setup ### Setup pre-commit hooks - Install pre-commit hooks with `pip install pre-commit`. - Run `pre-commit install` once to install the hooks, so they will be run on every commit. - If the hooks introduce changes, they'll be visible with `git diff`. Review them and `git add` them if everything is fine, then re-execute the before commit, it should pass now. - If you want to manually trigger the hooks, you may do `pre-commit run --all-files` Now all the pre-commit hooks will be automatically run when you try to commit and if they introduce some changes, you need to re-add the changed files before being able to commit and push. ### Ignore formatting revs - Run `git config blame.ignoreRevsFile .git-blame-ignore-revs`. This will make it so that `git blame` is aware of commits that were logged to be solely formatting-related. ## Doc-string syntax We're following NumPy doc-string conventions with the only notable difference being that we use Markdown instead of Rich text format (RTF) for markup within the doc-strings. Please see the existing documentation to see how to generate autodocs. ## Documentation - [guideline for documentation syntax](https://github.com/huggingface/doc-builder#readme) - images shall be uploaded via PR in the `bitsandbytes/` directory [here](https://huggingface.co/datasets/huggingface/documentation-images) - find the documentation builds for each PR in a link posted to the PR, such as https://moon-ci-docs.huggingface.co/docs/bitsandbytes/pr_1012/en/introduction bitsandbytes/docs/source/errors.mdx000066400000000000000000000030071506752665600200510ustar00rootroot00000000000000# Troubleshoot ## No kernel image available This problem arises with the cuda version loaded by bitsandbytes is not supported by your GPU, or if you pytorch CUDA version mismatches. To solve this problem you need to debug ``$LD_LIBRARY_PATH``, ``$CUDA_HOME`` as well as ``$PATH``. You can print these via ``echo $PATH``. You should look for multiple paths to different CUDA versions. This can include versions in your anaconda path, for example ``$HOME/anaconda3/lib``. You can check those versions via ``ls -l $HOME/anaconda3/lib/*cuda*`` or equivalent paths. Look at the CUDA versions of files in these paths. Does it match with ``nvidia-smi``? If you are feeling lucky, you can also try to compile the library from source. This can be still problematic if your PATH variables have multiple cuda versions. As such, it is recommended to figure out path conflicts before you proceed with compilation. ## `fatbinwrap` This error occurs if there is a mismatch between CUDA versions in the C++ library and the CUDA part. Make sure you have right CUDA in your `$PATH` and `$LD_LIBRARY_PATH` variable. In the conda base environment you can find the library under: ```bash ls $CONDA_PREFIX/lib/*cudart* ``` Make sure this path is appended to the `LD_LIBRARY_PATH` so bnb can find the CUDA runtime environment library (cudart). If this does not fix the issue, please try compilation from source next. If this does not work, please open an issue and paste the printed environment if you call `make` and the associated error when running bnb. bitsandbytes/docs/source/explanations/000077500000000000000000000000001506752665600205305ustar00rootroot00000000000000bitsandbytes/docs/source/explanations/optimizers.mdx000066400000000000000000000123461506752665600234550ustar00rootroot00000000000000# 8-bit optimizers Stateful optimizers maintain gradient statistics over time, for example, the exponentially smoothed sum (SGD with momentum) or squared sum (Adam) of past gradient values. This state can be used to accelerate optimization compared to plain stochastic gradient descent, but uses memory that might otherwise be allocated to model parameters. As a result, this limits the maximum size of models that can be trained in practice. Now take a look at the biggest models that can be trained with 8-bit optimizers.
Depending on your GPU size, you can train a much larger model with a 8-bit optimizer.
bitsandbytes optimizers use 8-bit statistics, while maintaining the performance levels of using 32-bit optimizer states. To overcome the resulting computational, quantization and stability challenges, 8-bit optimizers have three components: 1. Block-wise quantization: divides input tensors into smaller blocks that are independently quantized, isolating outliers and distributing the error more equally over all bits. Each block is processed in parallel across cores, yielding faster optimization and high precision quantization. 2. Dynamic quantization: quantizes both small and large values with high precision. 3. Stable embedding layer: improves stability during optimization for models with word embeddings. With these components, performing an optimizer update with 8-bit states is straightforward. The 8-bit optimizer states are dequantized to 32-bit before you perform the update, and then the states are quantized back to 8-bit for storage. The 8-bit to 32-bit conversion happens element-by-element in registers, meaning no slow copies to GPU memory or additional temporary memory are needed to perform quantization and dequantization. For GPUs, this makes 8-bit optimizers much faster than regular 32-bit optimizers.
A comparison of memory and time saved using 8-bit and 32-bit optimizers.
## Stable embedding layer The stable embedding layer improves the training stability of the standard word embedding layer for NLP tasks. It addresses the challenge of non-uniform input distributions and mitigates extreme gradient variations. This means the stable embedding layer can support more aggressive quantization strategies without compromising training stability, and it can help achieve stable training outcomes, which is particularly important for models dealing with diverse and complex language data. There are three features of the stable embedding layer: - Initialization: utilizes Xavier uniform initialization to maintain consistent variance, reducing the likelihood of large gradients. - Normalization: incorporates layer normalization before adding positional embeddings, aiding in output stability. - Optimizer states: employs 32-bit optimizer states exclusively for this layer to enhance stability, while the rest of the model may use standard 16-bit precision. ## Paged optimizers Paged optimizers are built on top of the [unified memory](https://developer.nvidia.com/blog/unified-memory-cuda-beginners/) feature of CUDA. Unified memory provides a single memory space the GPU and CPU can easily access. While this feature is not supported by PyTorch, it has been added to bitsandbytes. Paged optimizers works like regular CPU paging, which means that it *only becomes active if you run out of GPU memory*. When that happens, memory is transferred page-by-page from GPU to CPU. The memory is mapped, meaning that pages are pre-allocated on the CPU but they are not updated automatically. Pages are only updated if the memory is accessed or a swapping operation is launched. The unified memory feature is less efficient than regular asynchronous memory transfers, and you usually won't be able to get full PCIe memory bandwidth utilization. If you do a manual prefetch, transfer speeds can be high but still only about half or worse than the full PCIe memory bandwidth (tested on 16x lanes PCIe 3.0). This means performance depends highly on the particular use-case. For example, if you evict 1 GB of memory per forward-backward-optimizer loop, then you can expect about 50% of the PCIe bandwidth as time in the best case. So, 1 GB for PCIe 3.0 with 16x lanes would run at 16 GB/s, which is `1/(16*0.5) = 1/8 = 125ms` of overhead per optimizer step. Other overhead can be estimated for the particular use-case given a PCIe interface, lanes, and the memory evicted in each iteration. Compared to CPU offloading, a paged optimizer has zero overhead if all the memory fits onto the device and only some overhead if some of memory needs to be evicted. For offloading, you usually offload fixed parts of the model and need to off and onload all this memory with each iteration through the model (sometimes twice for both forward and backward pass). bitsandbytes/docs/source/explanations/resources.mdx000066400000000000000000000074441506752665600232650ustar00rootroot00000000000000# Papers, related resources & how to cite The below academic work is ordered in reverse chronological order. ## [SpQR: A Sparse-Quantized Representation for Near-Lossless LLM Weight Compression (Jun 2023)](https://arxiv.org/abs/2306.03078) Authors: Tim Dettmers, Ruslan Svirschevski, Vage Egiazarian, Denis Kuznedelev, Elias Frantar, Saleh Ashkboos, Alexander Borzunov, Torsten Hoefler, Dan Alistarh - [Twitter summary thread](https://twitter.com/Tim_Dettmers/status/1666076553665744896) ``` @article{dettmers2023spqr, title={SpQR: A Sparse-Quantized Representation for Near-Lossless LLM Weight Compression}, author={Dettmers, Tim and Svirschevski, Ruslan and Egiazarian, Vage and Kuznedelev, Denis and Frantar, Elias and Ashkboos, Saleh and Borzunov, Alexander and Hoefler, Torsten and Alistarh, Dan}, journal={arXiv preprint arXiv:2306.03078}, year={2023} } ``` ## [QLoRA: Efficient Finetuning of Quantized LLMs (May 2023)](https://arxiv.org/abs/2305.14314) Authors: Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, Luke Zettlemoyer - [Video](https://www.youtube.com/watch?v=y9PHWGOa8HA&ab_channel=LondonMachineLearningMeetup) - [Twitter summary thread](https://twitter.com/Tim_Dettmers/status/1661379354507476994) ``` @article{dettmers2023qlora, title={Qlora: Efficient finetuning of quantized llms}, author={Dettmers, Tim and Pagnoni, Artidoro and Holtzman, Ari and Zettlemoyer, Luke}, journal={arXiv preprint arXiv:2305.14314}, year={2023} } ``` ## [The case for 4-bit precision: k-bit Inference Scaling Laws (Dec 2022)](https://arxiv.org/abs/2212.09720) Authors: Tim Dettmers, Luke Zettlemoyer - [Video](https://www.youtube.com/watch?v=odlQa6AE1gY&ab_channel=TheInsideView) - [Twitter summary thread](https://twitter.com/Tim_Dettmers/status/1605209171758284805) ``` @inproceedings{dettmers2023case, title={The case for 4-bit precision: k-bit inference scaling laws}, author={Dettmers, Tim and Zettlemoyer, Luke}, booktitle={International Conference on Machine Learning}, pages={7750--7774}, year={2023}, organization={PMLR} } ``` ## [LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale (Nov 2022)](https://arxiv.org/abs/2208.07339) [[llm-int8]] Authors: Tim Dettmers, Mike Lewis, Younes Belkada, Luke Zettlemoyer - [LLM.int8() Blog Post](https://huggingface.co/blog/hf-bitsandbytes-integration) - [LLM.int8() Emergent Features Blog Post](https://timdettmers.com/2022/08/17/llm-int8-and-emergent-features/) - [Introduction to Weight Quantization](https://towardsdatascience.com/introduction-to-weight-quantization-2494701b9c0c) - [Poster](https://twitter.com/Tim_Dettmers/status/1598351301942951937) ``` @article{dettmers2022llm, title={Llm. int8 (): 8-bit matrix multiplication for transformers at scale}, author={Dettmers, Tim and Lewis, Mike and Belkada, Younes and Zettlemoyer, Luke}, journal={arXiv preprint arXiv:2208.07339}, year={2022} } ``` ## [8-bit Optimizers via Block-wise Quantization (Oct 2021)](https://arxiv.org/abs/2110.02861) Authors: Tim Dettmers, Mike Lewis, Sam Shleifer, Luke Zettlemoyer - [Video](https://www.youtube.com/watch?v=IxrlHAJtqKE) - [Twitter summary thread](https://twitter.com/Tim_Dettmers/status/1446472128979562499) ``` @article{DBLP:journals/corr/abs-2110-02861, author = {Tim Dettmers and Mike Lewis and Sam Shleifer and Luke Zettlemoyer}, title = {8-bit Optimizers via Block-wise Quantization}, journal = {CoRR}, volume = {abs/2110.02861}, year = {2021}, url = {https://arxiv.org/abs/2110.02861}, eprinttype = {arXiv}, eprint = {2110.02861}, timestamp = {Thu, 21 Oct 2021 16:20:08 +0200}, biburl = {https://dblp.org/rec/journals/corr/abs-2110-02861.bib}, bibsource = {dblp computer science bibliography, https://dblp.org} } ``` bitsandbytes/docs/source/faqs.mdx000066400000000000000000000006701506752665600174720ustar00rootroot00000000000000# FAQs Please submit your questions in [this Github Discussion thread](https://github.com/bitsandbytes-foundation/bitsandbytes/discussions/1013) if you feel that they will likely affect a lot of other users and that they haven't been sufficiently covered in the documentation. We'll pick the most generally applicable ones and post the QAs here or integrate them into the general documentation (also feel free to submit doc PRs, please). bitsandbytes/docs/source/fsdp_qlora.md000066400000000000000000000154761506752665600205140ustar00rootroot00000000000000# FSDP-QLoRA FSDP-QLoRA combines data parallelism (FSDP enables sharding model parameters, optimizer states, and gradients across GPUs), 4-bit quantization, and LoRA to train LLMs up to 70B parameters on a dual 24GB GPU system. This technique was released by [Answer.AI](https://www.answer.ai/posts/2024-03-06-fsdp-qlora) in collaboration with bitsandbytes to make training LLMs more efficient and accessible for everyone. This guide provides a brief guide on how bitsandbytes supports storing quantized weights to enable FSDP-QLoRA, and how to run training with the Hugging Face libraries. > [!TIP] > Other changes required for bitsandbytes to support FSDP-QLoRA, such as reconstructing the weights from the quantization metadata and preventing quantizing already quantized weights when they're moved from a CPU to GPU, are documented in this [Pull Request](https://github.com/bitsandbytes-foundation/bitsandbytes/pull/970) and described in the [Enabling 70B Finetuning on Consumer GPUs](https://www.answer.ai/posts/2024-03-14-fsdp-qlora-deep-dive) blog post. We highly recommend reading these resources for a better understanding of FSDP-QLoRA! ## Quantized data storage FSDP only supports sharding float data types which can be problematic because quantized weights are typically stored as integer data types (uint8). bitsandbytes doesn't have this problem because it uses `StoreChar` to read and write quantized weights regardless of the data type storage. This makes it simple to add a `quant_storage` parameter to the [`~nn.Linear4bit`] and [`~nn.Params4bit`] classes and set it to `torch.uint8` to maintain backward compatibility with the codebase. With the `quant_storage` parameter, you can select any of the FSDP supported data types to shard [`~nn.Linear4bit`] with such as bfloat16, float16 or float32. You'll typically access and configure this option from [`transformers.BitsAndBytesConfig`] by setting the `bnb_4bit_quant_storage` parameter. It is very **important** the `quant_storage` data type matches the data types used throughout the model because FSDP can only wrap layers and modules that have the *same floating data type*. Making sure the data types are aligned will ensure the model is correctly sharded. > [!TIP] > The `compute_dtype` is the data type used for computation inside the CUDA kernel, where the 4-bit quantized weights are unpacked from the data type in `quant_storage` and dequantized to `compute_dtype`. We recommend using torch.bfloat16 (if available on your hardware) for better numerical stability. ```py from transformers import BitsAndBytesConfig, AutoModelForCausalLM bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_quant_storage=torch.bfloat16, ) model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-2-70b", quantization_config=bnb_config, torch_dtype=torch.bfloat16, ) ``` Check out this [section](https://hf.co/docs/peft/main/en/accelerate/fsdp#use-peft-qlora-and-fsdp-for-finetuning-large-models-on-multiple-gpus) of the PEFT documentation for the config file and training code to run FSDP-QLoRA training. ## Training > [!TIP] > FSDP is a distributed training framework that needs to be launched as a distributed training job with a library like [Accelerate](https://hf.co/docs/accelerate/index) or [torchrun](https://pytorch.org/docs/stable/elastic/run.html). The launch command provided in this section uses Accelerate to launch the training script. bitsandbytes is deeply integrated with the Hugging Face ecosystem, making it easy to use with libraries like [Transformers](https://hf.co/docs/transformers), [PEFT](https://hf.co/docs/peft), and [TRL](https://hf.co/docs/trl). PEFT provides a configuration file ([fsdp_config_qlora.yaml](https://github.com/huggingface/peft/blob/main/examples/sft/configs/fsdp_config_qlora.yaml)), launch command ([run_peft_qlora_fsdp.sh](https://github.com/huggingface/peft/blob/main/examples/sft/run_peft_qlora_fsdp.sh)), and training script ([train.py](https://github.com/huggingface/peft/blob/main/examples/sft/train.py)) for running FSDP-QLoRA. To learn more, check out the [Use PEFT QLoRA and FSDP for finetuning large models on multiple GPUs](https://huggingface.co/docs/peft/main/en/accelerate/fsdp#use-peft-qlora-and-fsdp-for-finetuning-large-models-on-multiple-gpus) documentation. This section briefly covers the steps to run FSDP-QLoRA training. Before you begin, make sure you have the latest libraries installed. ```bash pip install -U bitsandbytes accelerate transformers peft trl ``` The important change that enables FSDP-QLoRA training is the `bnb_4bit_quant_storage` parameter in the [`~transformers.BitsAndBytesConfig`] class. This allows you to set the storage data type of the quantized weights to a float data type. ```py from transformers import BitsAndBytesConfig bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_storage=torch.bfloat16, ) ``` Pass the [`~transformers.BitsAndBytesConfig`] to a model to set it up for FSDP-QLoRA. You should set the `torch_dtype` parameter to match `bnb_4bit_quant_storage` so that the [`~nn.Linear4bit`] layers are wrapped identically to the `Linear` layers. If the storage types do not match, then each [`~nn.Linear4bit`] layer is wrapped individually. ```py from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-2-70b", quantization_config=bnb_config, torch_dtype=torch.bfloat16, ) ``` Configure the [`~peft.LoraConfig`] class for QLoRA training by setting `target_modules="all-linear"`. ```py from peft import LoraConfig peft_config = LoraConfig( lora_alpha=16, lora_dropout=0.1, r=64, bias="none", task_type="CAUSAL_LM", target_modules="all-linear", ) ``` Now you can pass everything to the [`~trl.SFTTrainer`] for training. ```py from trl import SFTTrainer trainer = SFTTrainer( model=model, train_dataset=dataset, peft_config=peft_config, processing_class=tokenizer, args=training_arguments, ) trainer.train() ``` ## Resources To learn more about FSDP and QLoRA, check out the following resources: - The [AnswerDotAI/fsdp_qlora](https://github.com/AnswerDotAI/fsdp_qlora) repository. - The introductory [You can now train a 70b language model at home](https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html) blog post by Answer.AI. - For an introduction to FSDP, read the [Introducing PyTorch Fully Sharded Data Parallel (FSDP) API](https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api) blog post. - For more details about QLoRA, take a look at the [Making LLMs even more accessible with bitsandbytes, 4-bit quantization and QLoRA](https://huggingface.co/blog/4bit-transformers-bitsandbytes) blog post. bitsandbytes/docs/source/index.mdx000066400000000000000000000021201506752665600176370ustar00rootroot00000000000000# bitsandbytes bitsandbytes enables accessible large language models via k-bit quantization for PyTorch. bitsandbytes provides three main features for dramatically reducing memory consumption for inference and training: * 8-bit optimizers uses block-wise quantization to maintain 32-bit performance at a small fraction of the memory cost. * LLM.int8() or 8-bit quantization enables large language model inference with only half the required memory and without any performance degradation. This method is based on vector-wise quantization to quantize most features to 8-bits and separately treating outliers with 16-bit matrix multiplication. * QLoRA or 4-bit quantization enables large language model training with several memory-saving techniques that don't compromise performance. This method quantizes a model to 4-bits and inserts a small set of trainable low-rank adaptation (LoRA) weights to allow training. # License bitsandbytes is MIT licensed. We thank Fabio Cannizzo for his work on [FastBinarySearch](https://github.com/fabiocannizzo/FastBinarySearch) which we use for CPU quantization. bitsandbytes/docs/source/installation.mdx000066400000000000000000000256451506752665600212520ustar00rootroot00000000000000# Installation Guide Welcome to the installation guide for the `bitsandbytes` library! This document provides step-by-step instructions to install `bitsandbytes` across various platforms and hardware configurations. We provide official support for NVIDIA GPUs, CPUs, Intel XPUs, and Intel Gaudi platforms. We also have experimental support for additional platforms such as AMD ROCm. ## Table of Contents - [System Requirements](#requirements) - [NVIDIA CUDA](#cuda) - [Installation via PyPI](#cuda-pip) - [Compile from Source](#cuda-compile) - [Intel XPU](#xpu) - [Installation via PyPI](#xpu-pip) - [Intel Gaudi](#gaudi) - [Installation via PyPI](#gaudi-pip) - [CPU](#cpu) - [Installation via PyPI](#cpu-pip) - [Compile from Source](#cpu-compile) - [AMD ROCm (Preview)](#rocm-preview) - [Preview Wheels](#preview-wheels) ## System Requirements[[requirements]] These are the minimum requirements for `bitsandbytes` across all platforms. Please be aware that some compute platforms may impose more strict requirements. * Python >= 3.9 * PyTorch >= 2.3 ## NVIDIA CUDA[[cuda]] `bitsandbytes` is currently supported on NVIDIA GPUs with [Compute Capability](https://developer.nvidia.com/cuda-gpus) 6.0+. The library can be built using CUDA Toolkit versions as old as **11.8**. | **Feature** | **CC Required** | **Example Hardware Requirement** | |---------------------------------|-----------------|---------------------------------------------| | LLM.int8() | 7.5+ | Turing (RTX 20 series, T4) or newer GPUs | | 8-bit optimizers/quantization | 6.0+ | Pascal (GTX 10X0 series, P100) or newer GPUs| | NF4/FP4 quantization | 6.0+ | Pascal (GTX 10X0 series, P100) or newer GPUs| > [!WARNING] > Support for Maxwell GPUs is deprecated and will be removed in a future release. > Maxwell support is not included in PyPI distributions from `v0.48.0` on and must be built from source. > For the best results, a Turing generation device or newer is recommended. ### Installation via PyPI[[cuda-pip]] This is the most straightforward and recommended installation option. The currently distributed `bitsandbytes` packages are built with the following configurations: | **OS** | **CUDA Toolkit** | **Host Compiler** | **Targets** |--------------------|------------------|----------------------|-------------- | **Linux x86-64** | 11.8 - 12.6 | GCC 11.2 | sm60, sm70, sm75, sm80, sm86, sm89, sm90 | **Linux x86-64** | 12.8 - 12.9 | GCC 11.2 | sm70, sm75, sm80, sm86, sm89, sm90, sm100, sm120 | **Linux x86-64** | 13.0 | GCC 11.2 | sm75, sm80, sm86, sm89, sm90, sm100, sm110, sm120 | **Linux aarch64** | 11.8 - 12.6 | GCC 11.2 | sm75, sm80, sm90 | **Linux aarch64** | 12.8 - 13.0 | GCC 11.2 | sm75, sm80, sm90, sm100, sm120 | **Windows x86-64** | 11.8 - 12.6 | MSVC 19.43+ (VS2022) | sm50, sm60, sm75, sm80, sm86, sm89, sm90 | **Windows x86-64** | 12.8 - 12.9 | MSVC 19.43+ (VS2022) | sm70, sm75, sm80, sm86, sm89, sm90, sm100, sm120 | **Windows x86-64** | 13.0 | MSVC 19.43+ (VS2022) | sm75, sm80, sm86, sm89, sm90, sm100, sm120 The Linux build has a minimum glibc version of 2.24. Use `pip` or `uv` to install the latest release: ```bash pip install bitsandbytes ``` ### Compile from Source[[cuda-compile]] > [!TIP] > Don't hesitate to compile from source! The process is pretty straight forward and resilient. This might be needed for older CUDA Toolkit versions or Linux distributions, or other less common configurations. For Linux and Windows systems, compiling from source allows you to customize the build configurations. See below for detailed platform-specific instructions (see the `CMakeLists.txt` if you want to check the specifics and explore some additional options): To compile from source, you need CMake >= **3.22.1** and Python >= **3.9** installed. Make sure you have a compiler installed to compile C++ (`gcc`, `make`, headers, etc.). It is recommended to use GCC 9 or newer. For example, to install a compiler and CMake on Ubuntu: ```bash apt-get install -y build-essential cmake ``` You should also install CUDA Toolkit by following the [NVIDIA CUDA Installation Guide for Linux](https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html) guide. The current minimum supported CUDA Toolkit version that we support is **11.8**. ```bash git clone https://github.com/bitsandbytes-foundation/bitsandbytes.git && cd bitsandbytes/ cmake -DCOMPUTE_BACKEND=cuda -S . make pip install -e . # `-e` for "editable" install, when developing BNB (otherwise leave that out) ``` > [!TIP] > If you have multiple versions of the CUDA Toolkit installed or it is in a non-standard location, please refer to CMake CUDA documentation for how to configure the CUDA compiler. Compilation from source on Windows systems require Visual Studio with C++ support as well as an installation of the CUDA Toolkit. To compile from source, you need CMake >= **3.22.1** and Python >= **3.9** installed. You should also install CUDA Toolkit by following the [CUDA Installation Guide for Windows](https://docs.nvidia.com/cuda/cuda-installation-guide-microsoft-windows/index.html) guide from NVIDIA. The current minimum supported CUDA Toolkit version that we support is **11.8**. ```bash git clone https://github.com/bitsandbytes-foundation/bitsandbytes.git && cd bitsandbytes/ cmake -DCOMPUTE_BACKEND=cuda -S . cmake --build . --config Release pip install -e . # `-e` for "editable" install, when developing BNB (otherwise leave that out) ``` Big thanks to [wkpark](https://github.com/wkpark), [Jamezo97](https://github.com/Jamezo97), [rickardp](https://github.com/rickardp), [akx](https://github.com/akx) for their amazing contributions to make bitsandbytes compatible with Windows. ## Intel XPU[[xpu]] * A compatible PyTorch version with Intel XPU support is required. The current minimum is **PyTorch 2.6.0**. It is recommended to use the latest stable release. See [Getting Started on Intel GPU](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html) for guidance. ### Installation via PyPI[[xpu-pip]] This is the most straightforward and recommended installation option. The currently distributed `bitsandbytes` packages are built with the following configurations: | **OS** | **oneAPI Toolkit** | **Kernel Implementation** | |--------------------|------------------|----------------------| | **Linux x86-64** | 2025.1.3 | SYCL + Triton | | **Windows x86-64** | N/A | SYCL | The Linux build has a minimum glibc version of 2.34. Use `pip` or `uv` to install the latest release: ```bash pip install bitsandbytes ``` ## Intel Gaudi[[gaudi]] * A compatible PyTorch version with Intel Gaudi support is required. The current minimum is **Gaudi v1.21** with **PyTorch 2.6.0**. It is recommended to use the latest stable release. See the Gaudi software [installation guide](https://docs.habana.ai/en/latest/Installation_Guide/index.html) for guidance. ### Installation from PyPI[[gaudi-pip]] Use `pip` or `uv` to install the latest release: ```bash pip install bitsandbytes ``` ## CPU[[cpu]] ### Installation from PyPI[[cpu-pip]] This is the most straightforward and recommended installation option. The currently distributed `bitsandbytes` packages are built with the following configurations: | **OS** | **Host Compiler** | Hardware Minimum |--------------------|----------------------|----------------------| | **Linux x86-64** | GCC 11.4 | AVX2 | | **Linux aarch64** | GCC 11.4 | | | **Windows x86-64** | MSVC 19.43+ (VS2022) | AVX2 | The Linux build has a minimum glibc version of 2.24. Use `pip` or `uv` to install the latest release: ```bash pip install bitsandbytes ``` ### Compile from Source[[cpu-compile]] To compile from source, simply install the package from source using `pip`. The package will be built for CPU only at this time. ```bash git clone https://github.com/bitsandbytes-foundation/bitsandbytes.git && cd bitsandbytes/ pip install -e . ``` ## AMD ROCm (Preview)[[rocm]] * A compatible PyTorch version with AMD ROCm support is required. It is recommended to use the latest stable release. See [PyTorch on ROCm](https://rocm.docs.amd.com/projects/install-on-linux/en/latest/install/3rd-party/pytorch-install.html) for guidance. * ROCm support is currently only available in our preview wheels or when building from source. ### Preview Wheels from `main`[[rocm-preview]] The currently distributed preview `bitsandbytes` are built with the following configurations: | **OS** | **ROCm** | **Targets** |--------------------|----------|---------------------------| | **Linux x86-64** | 6.1.2 | gfx90a / gfx942 / gfx1100 | **Linux x86-64** | 6.2.4 | gfx90a / gfx942 / gfx1100 | **Linux x86-64** | 6.3.4 | gfx90a / gfx942 / gfx1100 | **Linux x86-64** | 6.4.4 | gfx90a / gfx942 / gfx1100 | **Linux x86-64** | 7.0.0 | gfx90a / gfx942 / gfx1100 **Windows is not currently supported.** Please see [Preview Wheels](#preview-wheels) for installation instructions. ### Compile from Source[[rocm-compile]] bitsandbytes can be compiled from ROCm 6.1 - ROCm 7.0. ```bash # Install bitsandbytes from source # Clone bitsandbytes repo git clone https://github.com/bitsandbytes-foundation/bitsandbytes.git && cd bitsandbytes/ # Compile & install apt-get install -y build-essential cmake # install build tools dependencies, unless present cmake -DCOMPUTE_BACKEND=hip -S . # Use -DBNB_ROCM_ARCH="gfx90a;gfx942" to target specific gpu arch make pip install -e . # `-e` for "editable" install, when developing BNB (otherwise leave that out) ``` ## Preview Wheels[[preview-wheels]] If you would like to use new features even before they are officially released and help us test them, feel free to install the wheel directly from our CI (*the wheel links will remain stable!*): ```bash # Note: if you don't want to reinstall our dependencies, append the `--no-deps` flag! # x86_64 (most users) pip install --force-reinstall https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl # ARM/aarch64 pip install --force-reinstall https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_aarch64.whl ``` ```bash # Note: if you don't want to reinstall our dependencies, append the `--no-deps` flag! pip install --force-reinstall https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl ``` bitsandbytes/docs/source/integrations.mdx000066400000000000000000000136251506752665600212520ustar00rootroot00000000000000# Integrations bitsandbytes is widely integrated with many of the libraries in the Hugging Face and wider PyTorch ecosystem. This guide provides a brief overview of the integrations and how to use bitsandbytes with them. For more details, you should refer to the linked documentation for each library. ## Transformers > [!TIP] > Learn more in the bitsandbytes Transformers integration [guide](https://huggingface.co/docs/transformers/quantization#bitsandbytes). With Transformers, it's very easy to load any model in 4 or 8-bit and quantize them on the fly. To configure the quantization parameters, specify them in the [`~transformers.BitsAndBytesConfig`] class. For example, to load and quantize a model to 4-bits and use the bfloat16 data type for compute: > [!WARNING] > bfloat16 is the ideal `compute_dtype` if your hardware supports it. While the default `compute_dtype`, float32, ensures backward compatibility (due to wide-ranging hardware support) and numerical stability, it is large and slows down computations. In contrast, float16 is smaller and faster but can lead to numerical instabilities. bfloat16 combines the best aspects of both; it offers the numerical stability of float32 and the reduced memory footprint and speed of a 16-bit data type. Check if your hardware supports bfloat16 and configure it using the `bnb_4bit_compute_dtype` parameter in [`~transformers.BitsAndBytesConfig`]! ```py from transformers import AutoModelForCausalLM, BitsAndBytesConfig quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16) model_4bit = AutoModelForCausalLM.from_pretrained( "bigscience/bloom-1b7", device_map=device_map, quantization_config=quantization_config, ) ``` ### 8-bit optimizers You can use any of the 8-bit or paged optimizers with Transformers by passing them to the [`~transformers.Trainer`] class on initialization. All bitsandbytes optimizers are supported by passing the correct string in the [`~transformers.TrainingArguments`] `optim` parameter. For example, to load a [`~bitsandbytes.optim.PagedAdamW32bit`] optimizer: ```py from transformers import TrainingArguments, Trainer training_args = TrainingArguments( ..., optim="paged_adamw_32bit", ) trainer = Trainer(model, training_args, ...) trainer.train() ``` ## PEFT > [!TIP] > Learn more in the bitsandbytes PEFT integration [guide](https://huggingface.co/docs/peft/developer_guides/quantization#quantization). PEFT builds on the bitsandbytes Transformers integration, and extends it for training with a few more steps. Let's prepare the 4-bit model from the section above for training. Call the [`~peft.prepare_model_for_kbit_training`] method to prepare the model for training. This only works for Transformers models! ```py from peft import prepare_model_for_kbit_training model_4bit = prepare_model_for_kbit_training(model_4bit) ``` Setup a [`~peft.LoraConfig`] to use QLoRA: ```py from peft import LoraConfig config = LoraConfig( r=16, lora_alpha=8, target_modules="all-linear", lora_dropout=0.05 bias="none", task_type="CAUSAL_LM" ) ``` Now call the [`~peft.get_peft_model`] function on your model and config to create a trainable [`PeftModel`]. ```py from peft import get_peft_model model = get_peft_model(model_4bit, config) ``` ## Accelerate > [!TIP] > Learn more in the bitsandbytes Accelerate integration [guide](https://huggingface.co/docs/accelerate/usage_guides/quantization). bitsandbytes is also easily usable from Accelerate and you can quantize any PyTorch model by passing a [`~accelerate.utils.BnbQuantizationConfig`] with your desired settings, and then calling the [`~accelerate.utils.load_and_quantize_model`] function to quantize it. ```py from accelerate import init_empty_weights from accelerate.utils import BnbQuantizationConfig, load_and_quantize_model from mingpt.model import GPT model_config = GPT.get_default_config() model_config.model_type = 'gpt2-xl' model_config.vocab_size = 50257 model_config.block_size = 1024 with init_empty_weights(): empty_model = GPT(model_config) bnb_quantization_config = BnbQuantizationConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, # optional bnb_4bit_use_double_quant=True, # optional bnb_4bit_quant_type="nf4" # optional ) quantized_model = load_and_quantize_model( empty_model, weights_location=weights_location, bnb_quantization_config=bnb_quantization_config, device_map = "auto" ) ``` ## PyTorch Lightning and Lightning Fabric bitsandbytes is available from: - [PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/), a deep learning framework for professional AI researchers and machine learning engineers who need maximal flexibility without sacrificing performance at scale. - [Lightning Fabric](https://lightning.ai/docs/fabric/stable/), a fast and lightweight way to scale PyTorch models without boilerplate. Learn more in the bitsandbytes PyTorch Lightning integration [guide](https://lightning.ai/docs/pytorch/stable/common/precision_intermediate.html#quantization-via-bitsandbytes). ## Lit-GPT bitsandbytes is integrated with [Lit-GPT](https://github.com/Lightning-AI/lit-gpt), a hackable implementation of state-of-the-art open-source large language models. Lit-GPT is based on Lightning Fabric, and it can be used for quantization during training, finetuning, and inference. Learn more in the bitsandbytes Lit-GPT integration [guide](https://github.com/Lightning-AI/lit-gpt/blob/main/tutorials/quantize.md). ## Blog posts To learn in more detail about some of bitsandbytes integrations, take a look at the following blog posts: - [Making LLMs even more accessible with bitsandbytes, 4-bit quantization and QLoRA](https://huggingface.co/blog/4bit-transformers-bitsandbytes) - [A Gentle Introduction to 8-bit Matrix Multiplication for transformers at scale using Hugging Face Transformers, Accelerate and bitsandbytes](https://huggingface.co/blog/hf-bitsandbytes-integration) bitsandbytes/docs/source/optimizers.mdx000066400000000000000000000107651506752665600207530ustar00rootroot00000000000000# 8-bit optimizers With 8-bit optimizers, large models can be finetuned with 75% less GPU memory without losing any accuracy compared to training with standard 32-bit optimizers. The reduced memory requirements means 8-bit optimizers are 4x faster than a standard optimizer, and no hyperparameter tuning is required. This guide will show you how to use 8-bit optimizers. > [!WARNING] > 8-bit optimizers reduce memory usage and accelerate optimization on a wide range of tasks. However, since 8-bit optimizers only reduce memory proportional to the number of parameters, models that use large amounts of activation memory, such as convolutional networks, don't really benefit from 8-bit optimizers. 8-bit optimizers are most beneficial for training or finetuning models with many parameters on highly memory-constrained GPUs. 8-bit optimizers are a drop-in replacement for regular optimizers which means they also accept the same arguments as a regular optimizer. For NLP models, it is recommended to use the [`~nn.StableEmbedding`] class to improve stability and results. ```diff import bitsandbytes as bnb - adam = torch.optim.Adam(...) + adam = bnb.optim.Adam8bit(...) # recommended for NLP models - before: torch.nn.Embedding(...) + bnb.nn.StableEmbedding(...) ``` By default, all parameter tensors with less than 4096 elements are kept at 32-bits even if you initialize those parameters with 8-bit optimizers. This is done because small tensors do not save much memory and often contain highly variable parameters (biases) or parameters that require high precision (batch norm, layer norm). You can change this value with the `min_8bit_size` parameter. For example, if you want to optimize parameters to 8-bits only if the minimum size is 16384 values (it is recommended to use multiples of 4096): ```py import bitsandbytes as bnb adam = bnb.optim.Adam8bit(model.parameters(), min_8bit_size=16384) ``` Other parameters you can configure include the learning rate (`lr`), the decay rates (`betas`), the number of bits of the optimizer state (`optim_bits`), and percentile clipping (`percentile_clipping`) which can increase stability. For example, to initialize a 32-bit [`~bitsandbytes.optim.Adam`] optimizer with 5th percentile clipping: ```py import bitsandbytes as bnb adam = bnb.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.995), optim_bits=32, percentile_clipping=5) ``` ## Optimize unstable parameters To optimize some unstable parameters with 32-bit Adam and others with 8-bit Adam, use the [`~bitsandbytes.optim.GlobalOptimManager`] class to override the specific hyperparameters for a particular layer. You'll need to: 1. Register the parameters while they're on the CPU. ```py import torch import bitsandbytes as bnb mng = bnb.optim.GlobalOptimManager.get_instance() model = MyModel() mng.register_parameters(model.parameters()) ``` 2. Override the config with the new desired hyperparameters. For example, let's override the `model.fc1.weight` layer to use 32-bit Adam. > [!TIP] > Check the optimizer API documentation for more information about other hyperparameters you can override. ```py model = model.cuda() # use 8-bit optimizer states for all parameters adam = bnb.optim.Adam(model.parameters(), lr=0.001, optim_bits=8) # override the parameter model.fc1.weight now uses 32-bit Adam mng.override_config(model.fc1.weight, "optim_bits", 32) ``` You can also override multiple layers at once by passing them as a list and the new hyperparameters as a dictionary. For example, let's override the `model.special.weight` and `model.also_special.weight` layers to use sparse optimization and a lower learning and decay rate. ```py mng.override_config([model.special.weight, model.also_special.weight], key_value_dict ={'is_sparse': True, 'lr': 1e-5, 'betas'=(0.9, 0.98)}) ``` For a specific layer, we recommend overriding locally in each module. Pass the module, the parameter, and its attribute name to the [`~bitsandbytes.optim.GlobalOptimManager`]: ```py class MyModule(torch.nn.Module): def __init__(d_in, d_out): super(MyModule, self).__init__() self.linear = torch.nn.Linear(d_in, d_out) # optimization will happen in 32-bit and # learning rate will be set to 0.0001 independent of the main learning rate config = {'optim_bits': 32, 'lr' : 0.0001} GlobalOptimManager.get_instance().register_module_override(self, 'weight', config) ``` ## Next steps For more conceptual details and explanation about 8-bit optimizers, take a look at the [8-bit optimizers](./explanations/optimizers) guide. bitsandbytes/docs/source/quickstart.mdx000066400000000000000000000003371506752665600207320ustar00rootroot00000000000000# Quickstart ## How does it work? ... work in progress ... (Community contributions would we very welcome!) ## Minimal examples The following code illustrates the steps above. ```py code examples will soon follow ``` bitsandbytes/docs/source/reference/000077500000000000000000000000001506752665600177615ustar00rootroot00000000000000bitsandbytes/docs/source/reference/functional.mdx000066400000000000000000000024021506752665600226330ustar00rootroot00000000000000# Overview The `bitsandbytes.functional` API provides the low-level building blocks for the library's features. ## When to Use `bitsandbytes.functional` * When you need direct control over quantized operations and their parameters. * To build custom layers or operations leveraging low-bit arithmetic. * To integrate with other ecosystem tooling. * For experimental or research purposes requiring non-standard quantization or performance optimizations. ## LLM.int8() [[autodoc]] functional.int8_linear_matmul [[autodoc]] functional.int8_mm_dequant [[autodoc]] functional.int8_vectorwise_dequant [[autodoc]] functional.int8_vectorwise_quant ## 4-bit [[autodoc]] functional.dequantize_4bit [[autodoc]] functional.dequantize_fp4 [[autodoc]] functional.dequantize_nf4 [[autodoc]] functional.gemv_4bit [[autodoc]] functional.quantize_4bit [[autodoc]] functional.quantize_fp4 [[autodoc]] functional.quantize_nf4 [[autodoc]] functional.QuantState ## Dynamic 8-bit Quantization Primitives used in the 8-bit optimizer quantization. For more details see [8-Bit Approximations for Parallelism in Deep Learning](https://arxiv.org/abs/1511.04561) [[autodoc]] functional.dequantize_blockwise [[autodoc]] functional.quantize_blockwise ## Utility [[autodoc]] functional.get_ptr bitsandbytes/docs/source/reference/nn/000077500000000000000000000000001506752665600203745ustar00rootroot00000000000000bitsandbytes/docs/source/reference/nn/embeddings.mdx000066400000000000000000000012121506752665600232030ustar00rootroot00000000000000# Embedding The embedding class is used to store and retrieve word embeddings from their indices. There are two types of embeddings in bitsandbytes, the standard PyTorch [`Embedding`] class and the [`StableEmbedding`] class. The [`StableEmbedding`] class was introduced in the [8-bit Optimizers via Block-wise Quantization](https://hf.co/papers/2110.02861) paper to reduce gradient variance as a result of the non-uniform distribution of input tokens. This class is designed to support quantization. ## Embedding [[autodoc]] bitsandbytes.nn.Embedding - __init__ ## StableEmbedding [[autodoc]] bitsandbytes.nn.StableEmbedding - __init__ bitsandbytes/docs/source/reference/nn/linear4bit.mdx000066400000000000000000000013501506752665600231420ustar00rootroot00000000000000# 4-bit quantization [QLoRA](https://hf.co/papers/2305.14314) is a finetuning method that quantizes a model to 4-bits and adds a set of low-rank adaptation (LoRA) weights to the model and tuning them through the quantized weights. This method also introduces a new data type, 4-bit NormalFloat (`LinearNF4`) in addition to the standard Float4 data type (`LinearFP4`). `LinearNF4` is a quantization data type for normally distributed data and can improve performance. ## Linear4bit [[autodoc]] bitsandbytes.nn.Linear4bit - __init__ ## LinearFP4 [[autodoc]] bitsandbytes.nn.LinearFP4 - __init__ ## LinearNF4 [[autodoc]] bitsandbytes.nn.LinearNF4 - __init__ ## Params4bit [[autodoc]] bitsandbytes.nn.Params4bit - __init__ bitsandbytes/docs/source/reference/nn/linear8bit.mdx000066400000000000000000000015451506752665600231540ustar00rootroot00000000000000# LLM.int8() [LLM.int8()](https://hf.co/papers/2208.07339) is a quantization method that aims to make large language model inference more accessible without significant degradation. Unlike naive 8-bit quantization, which can result in loss of critical information and accuracy, LLM.int8() dynamically adapts to ensure sensitive components of the computation retain higher precision when needed. The key is to extract the outliers from the inputs and weights and multiply them in 16-bit. All other values are multiplied in 8-bit before being dequantized back to 16-bits. The outputs from the 16-bit and 8-bit multiplication are combined to produce the final output. [Further Resources](../../explanations/resources#llm-int8) ## Linear8bitLt [[autodoc]] bitsandbytes.nn.Linear8bitLt - __init__ ## Int8Params [[autodoc]] bitsandbytes.nn.Int8Params - __init__ bitsandbytes/docs/source/reference/optim/000077500000000000000000000000001506752665600211115ustar00rootroot00000000000000bitsandbytes/docs/source/reference/optim/adagrad.mdx000066400000000000000000000011731506752665600232100ustar00rootroot00000000000000# AdaGrad [AdaGrad (Adaptive Gradient)](https://jmlr.org/papers/v12/duchi11a.html) is an adaptive learning rate optimizer. AdaGrad stores a sum of the squared past gradients for each parameter and uses it to scale their learning rate. This allows the learning rate to be automatically lower or higher depending on the magnitude of the gradient, eliminating the need to manually tune the learning rate. ## Adagrad[[api-class]] [[autodoc]] bitsandbytes.optim.Adagrad - __init__ ## Adagrad8bit [[autodoc]] bitsandbytes.optim.Adagrad8bit - __init__ ## Adagrad32bit [[autodoc]] bitsandbytes.optim.Adagrad32bit - __init__ bitsandbytes/docs/source/reference/optim/adam.mdx000066400000000000000000000017731506752665600225350ustar00rootroot00000000000000# Adam [Adam (Adaptive moment estimation)](https://hf.co/papers/1412.6980) is an adaptive learning rate optimizer, combining ideas from [`SGD`] with momentum and [`RMSprop`] to automatically scale the learning rate: - a weighted average of the past gradients to provide direction (first-moment) - a weighted average of the *squared* past gradients to adapt the learning rate to each parameter (second-moment) bitsandbytes also supports paged optimizers which take advantage of CUDAs unified memory to transfer memory from the GPU to the CPU when GPU memory is exhausted. ## Adam[[api-class]] [[autodoc]] bitsandbytes.optim.Adam - __init__ ## Adam8bit [[autodoc]] bitsandbytes.optim.Adam8bit - __init__ ## Adam32bit [[autodoc]] bitsandbytes.optim.Adam32bit - __init__ ## PagedAdam [[autodoc]] bitsandbytes.optim.PagedAdam - __init__ ## PagedAdam8bit [[autodoc]] bitsandbytes.optim.PagedAdam8bit - __init__ ## PagedAdam32bit [[autodoc]] bitsandbytes.optim.PagedAdam32bit - __init__ bitsandbytes/docs/source/reference/optim/adamw.mdx000066400000000000000000000015471506752665600227230ustar00rootroot00000000000000# AdamW [AdamW](https://hf.co/papers/1711.05101) is a variant of the [`Adam`] optimizer that separates weight decay from the gradient update based on the observation that the weight decay formulation is different when applied to [`SGD`] and [`Adam`]. bitsandbytes also supports paged optimizers which take advantage of CUDAs unified memory to transfer memory from the GPU to the CPU when GPU memory is exhausted. ## AdamW[[api-class]] [[autodoc]] bitsandbytes.optim.AdamW - __init__ ## AdamW8bit [[autodoc]] bitsandbytes.optim.AdamW8bit - __init__ ## AdamW32bit [[autodoc]] bitsandbytes.optim.AdamW32bit - __init__ ## PagedAdamW [[autodoc]] bitsandbytes.optim.PagedAdamW - __init__ ## PagedAdamW8bit [[autodoc]] bitsandbytes.optim.PagedAdamW8bit - __init__ ## PagedAdamW32bit [[autodoc]] bitsandbytes.optim.PagedAdamW32bit - __init__ bitsandbytes/docs/source/reference/optim/ademamix.mdx000066400000000000000000000013571506752665600234160ustar00rootroot00000000000000# AdEMAMix [AdEMAMix](https://hf.co/papers/2409.03137) is a variant of the [`Adam`] optimizer. bitsandbytes also supports paged optimizers which take advantage of CUDAs unified memory to transfer memory from the GPU to the CPU when GPU memory is exhausted. ## AdEMAMix[[api-class]] [[autodoc]] bitsandbytes.optim.AdEMAMix - __init__ ## AdEMAMix8bit [[autodoc]] bitsandbytes.optim.AdEMAMix8bit - __init__ ## AdEMAMix32bit [[autodoc]] bitsandbytes.optim.AdEMAMix32bit - __init__ ## PagedAdEMAMix [[autodoc]] bitsandbytes.optim.PagedAdEMAMix - __init__ ## PagedAdEMAMix8bit [[autodoc]] bitsandbytes.optim.PagedAdEMAMix8bit - __init__ ## PagedAdEMAMix32bit [[autodoc]] bitsandbytes.optim.PagedAdEMAMix32bit - __init__ bitsandbytes/docs/source/reference/optim/lamb.mdx000066400000000000000000000012651506752665600225420ustar00rootroot00000000000000# LAMB [LAMB (Layerwise adaptive large batch optimization)](https://hf.co/papers/1904.00962) is an adaptive optimizer designed for training with large batch sizes to accelerate training, combining ideas from [`LARS`] and [`Adam`] to automatically scale the learning rate for each layer: - calculates a *trust ratio* between the weight and gradient norm in a layer and clips the ratio to prevent overly large or small updates - updates weights with the first and second-moments ## LAMB[[api-class]] [[autodoc]] bitsandbytes.optim.LAMB - __init__ ## LAMB8bit [[autodoc]] bitsandbytes.optim.LAMB8bit - __init__ ## LAMB32bit [[autodoc]] bitsandbytes.optim.LAMB32bit - __init__ bitsandbytes/docs/source/reference/optim/lars.mdx000066400000000000000000000011351506752665600225640ustar00rootroot00000000000000# LARS [LARS (Layer-wise Adaptive Rate Scaling)](https:/hf.co/papers/1708.03888) is an optimizer designed for training with large batch sizes to accelerate training. LARS uses a separate learning rate for each *layer* instead of each parameter. The learning rate is calculated from a *trust ratio* between the weight and gradient norm in a layer. This helps calibrate a stable update size. ## LARS[[api-class]] [[autodoc]] bitsandbytes.optim.LARS - __init__ ## LARS8bit [[autodoc]] bitsandbytes.optim.LARS8bit - __init__ ## LARS32bit [[autodoc]] bitsandbytes.optim.LARS32bit - __init__ bitsandbytes/docs/source/reference/optim/lion.mdx000066400000000000000000000013541506752665600225670ustar00rootroot00000000000000# Lion [Lion (Evolved Sign Momentum)](https://hf.co/papers/2302.06675) is a unique optimizer that uses the sign of the gradient to determine the update direction of the momentum. This makes Lion more memory-efficient and faster than [`AdamW`] which tracks and store the first and second-order moments. ## Lion[[api-class]] [[autodoc]] bitsandbytes.optim.Lion - __init__ ## Lion8bit [[autodoc]] bitsandbytes.optim.Lion8bit - __init__ ## Lion32bit [[autodoc]] bitsandbytes.optim.Lion32bit - __init__ ## PagedLion [[autodoc]] bitsandbytes.optim.PagedLion - __init__ ## PagedLion8bit [[autodoc]] bitsandbytes.optim.PagedLion8bit - __init__ ## PagedLion32bit [[autodoc]] bitsandbytes.optim.PagedLion32bit - __init__ bitsandbytes/docs/source/reference/optim/optim_overview.mdx000066400000000000000000000020521506752665600247000ustar00rootroot00000000000000# Overview [8-bit optimizers](https://hf.co/papers/2110.02861) reduce the memory footprint of 32-bit optimizers without any performance degradation which means you can train large models with many parameters faster. At the core of 8-bit optimizers is block-wise quantization which enables quantization accuracy, computational efficiency, and stability. bitsandbytes provides 8-bit optimizers through the base [`Optimizer8bit`] class, and additionally provides [`Optimizer2State`] and [`Optimizer1State`] for 2-state (for example, [`Adam`]) and 1-state (for example, [`Adagrad`]) optimizers respectively. To provide custom optimizer hyperparameters, use the [`GlobalOptimManager`] class to configure the optimizer. ## Optimizer8bit [[autodoc]] bitsandbytes.optim.optimizer.Optimizer8bit - __init__ ## Optimizer2State [[autodoc]] bitsandbytes.optim.optimizer.Optimizer2State - __init__ ## Optimizer1State [[autodoc]] bitsandbytes.optim.optimizer.Optimizer1State - __init__ ## Utilities [[autodoc]] bitsandbytes.optim.optimizer.GlobalOptimManager bitsandbytes/docs/source/reference/optim/rmsprop.mdx000066400000000000000000000010731506752665600233260ustar00rootroot00000000000000# RMSprop RMSprop is an adaptive learning rate optimizer that is very similar to [`Adagrad`]. RMSprop stores a *weighted average* of the squared past gradients for each parameter and uses it to scale their learning rate. This allows the learning rate to be automatically lower or higher depending on the magnitude of the gradient, and it prevents the learning rate from diminishing. ## RMSprop[[api-class]] [[autodoc]] bitsandbytes.optim.RMSprop ## RMSprop8bit [[autodoc]] bitsandbytes.optim.RMSprop8bit ## RMSprop32bit [[autodoc]] bitsandbytes.optim.RMSprop32bit bitsandbytes/docs/source/reference/optim/sgd.mdx000066400000000000000000000011761506752665600224050ustar00rootroot00000000000000# SGD Stochastic gradient descent (SGD) is a basic gradient descent optimizer to minimize loss given a set of model parameters and updates the parameters in the opposite direction of the gradient. The update is performed on a randomly sampled mini-batch of data from the dataset. bitsandbytes also supports momentum and Nesterov momentum to accelerate SGD by adding a weighted average of past gradients to the current gradient. ## SGD[[api-class]] [[autodoc]] bitsandbytes.optim.SGD - __init__ ## SGD8bit [[autodoc]] bitsandbytes.optim.SGD8bit - __init__ ## SGD32bit [[autodoc]] bitsandbytes.optim.SGD32bit - __init__ bitsandbytes/examples/000077500000000000000000000000001506752665600154115ustar00rootroot00000000000000bitsandbytes/examples/compile_inference.py000066400000000000000000000016401506752665600214320ustar00rootroot00000000000000import torch import torch._dynamo from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig # torch._dynamo.config.suppress_errors = True torch.set_float32_matmul_precision("high") quantization_config = BitsAndBytesConfig(load_in_8bit=True) # torch._dynamo.config.capture_dynamic_output_shape_ops = True model_id = "google/gemma-2-2b-it" # model_id = "Qwen/Qwen2.5-7B" tokenizer = AutoTokenizer.from_pretrained(model_id) model = AutoModelForCausalLM.from_pretrained( model_id, quantization_config=quantization_config, device_map="auto", torch_dtype=torch.bfloat16, ) input_text = "Write me a poem about Machine Learning." input_ids = tokenizer(input_text, return_tensors="pt").to(model.device) # model.forward = torch.compile(model.forward, fullgraph=True) model = torch.compile(model) outputs = model.generate(**input_ids, max_new_tokens=32) print(tokenizer.decode(outputs[0])) bitsandbytes/examples/int8_inference_huggingface.py000066400000000000000000000012721506752665600232140ustar00rootroot00000000000000import torch from transformers import LlamaForCausalLM, LlamaTokenizer MAX_NEW_TOKENS = 128 model_name = "meta-llama/Llama-2-7b-hf" text = "Hamburg is in which country?\n" tokenizer = LlamaTokenizer.from_pretrained(model_name) input_ids = tokenizer(text, return_tensors="pt").input_ids max_memory = f"{int(torch.cuda.mem_get_info()[0] / 1024**3) - 2}GB" n_gpus = torch.cuda.device_count() max_memory = {i: max_memory for i in range(n_gpus)} model = LlamaForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True, max_memory=max_memory) generated_ids = model.generate(input_ids, max_length=MAX_NEW_TOKENS) print(tokenizer.decode(generated_ids[0], skip_special_tokens=True)) bitsandbytes/include/000077500000000000000000000000001506752665600152165ustar00rootroot00000000000000bitsandbytes/include/AAlloc.h000066400000000000000000000034301506752665600165220ustar00rootroot00000000000000#pragma once #include "Portable.h" namespace BinSearch { namespace Details { template bool isAligned(const T *p, size_t A) { return (reinterpret_cast(p) % A) == 0; } template struct AlignedVec { AlignedVec() : m_storage(0) , m_data(0) , m_sz(0) { } static size_t nBytes(size_t sz) { return sz * sizeof(T) + A; } static size_t shiftAmt(char *p) { return A>1? (A - (reinterpret_cast(p) % A)) % A: 0; } void setPtr(char *p, size_t sz) { m_sz = sz; m_data = reinterpret_cast(p + shiftAmt(p)); } //void setPtr(T *p, size_t sz) //{ // m_sz = sz; // if (A>1) // myassert(((reinterpret_cast(p) % A) == 0), "bad alignment"); // m_data = p; //} // internal allocation void resize(size_t sz) { m_storage = new char[nBytes(sz)]; setPtr(m_storage, sz); } // external allocation void set(char *storage, size_t sz) { setPtr(storage, sz); } ~AlignedVec() { if (m_storage) delete [] m_storage; } size_t size() const { return m_sz; } T& operator[](size_t i) { return m_data[i]; } const T& operator[](size_t i) const { return m_data[i]; } T* begin() { return m_data; } T* end() { return m_data+m_sz; } const T* begin() const { return m_data; } const T* end() const { return m_data+m_sz; } T& front() { return m_data[0]; } T& back() { return m_data[m_sz-1]; } const T& front() const { return m_data[0]; } const T& back() const { return m_data[m_sz - 1]; } private: char *m_storage; T *m_data; size_t m_sz; }; } // namespace Details } // namespace BinSearch bitsandbytes/include/Algo-Direct-Common.h000066400000000000000000000235551506752665600207210ustar00rootroot00000000000000#pragma once #include #include #include #include "AAlloc.h" namespace BinSearch { namespace Details { namespace DirectAux { #define SAFETY_MULTI_PASS true template struct HResults { HResults(T h, double ratio, size_t n) : H(h), hRatio(ratio), nInc(n) {} T H; double hRatio; size_t nInc; }; #ifdef USE_FMA template struct IsDirect { static const bool value = (A == Direct) || (A == DirectFMA); }; template struct IsDirect2 { static const bool value = (A == Direct2) || (A == Direct2FMA); }; template struct IsDirectCache { static const bool value = (A == DirectCache) || (A == DirectCacheFMA); }; #else template struct IsDirect { static const bool value = (A == Direct); }; template struct IsDirect2 { static const bool value = (A == Direct2); }; template struct IsDirectCache { static const bool value = (A == DirectCache); }; #endif // general definition template struct BucketElem { FORCE_INLINE void set( uint32 b, const T *) { m_b = b; } FORCE_INLINE uint32 index() const { return m_b; } private: uint32 m_b; }; // specialization for DirectCache methods template struct MatchingIntType; template <> struct MatchingIntType { typedef uint64 type; }; template <> struct MatchingIntType { typedef uint32 type; }; template struct BucketElem::value >::type > { typedef typename MatchingIntType::type I; void set(uint32 b, const T *xi) { u.u.x = xi[b]; u.u.b = b; } FORCE_INLINE I index() const { return u.u.b; } FORCE_INLINE T x() const { return u.u.x; } private: union { double dummy; struct { T x; I b; } u; } u; }; template struct DirectTraits { static void checkH(T scaler, T x0, T xN) { T Dn = xN - x0; T ifmax = Dn * scaler; myassert((ifmax < std::numeric_limits::max() - (Gap - 1)), "Problem unfeasible: index size exceeds uint32 capacity:" << " D[N] =" << Dn << ", H =" << scaler << ", H D[n] =" << ifmax << "\n" ); } FORCE_INLINE static uint32 f(T scaler, T x0, T z) { T tmp = scaler * (z - x0); #ifdef USE_SSE2 return ftoi(FVec1(tmp)); #else return static_cast(tmp); #endif } template FORCE_INLINE static typename FTOITraits::vec_t f(const FVec& scaler, const FVec& x0, const FVec& z) { return ftoi(scaler*(z-x0)); } static T cst0(T scaler, T x0) { return x0; } }; #ifdef USE_FMA template struct DirectTraits { typedef FVec1 fVec1; static void checkH(T scaler, T H_Times_x0, T xN) { union { typename FVec1::vec_t v; T s; } ifmax; ifmax.v = mulSub(fVec1(scaler), fVec1(xN), fVec1(H_Times_x0)); myassert((ifmax.s < std::numeric_limits::max() - (Gap - 1)), "Problem unfeasible: index size exceeds uint32 capacity:" << " H X[0] =" << H_Times_x0 << ", H =" << scaler << ", X[N] =" << xN << ", H X[N] - H X[0] =" << ifmax.s << "\n" ); } FORCE_INLINE static uint32 f(T scaler, T Hx0, T xi) { return ftoi(mulSub(fVec1(scaler), fVec1(xi), fVec1(Hx0))); } template FORCE_INLINE static typename FTOITraits::vec_t f(const FVec& scaler, const FVec& H_Times_X0, const FVec& z) { return ftoi(mulSub(scaler, z, H_Times_X0)); } static T cst0(T scaler, T x0) { return scaler*x0; } }; #endif template struct DirectInfo { static const bool UseFMA = (A == DirectFMA) || (A == Direct2FMA) || (A == DirectCacheFMA); typedef DirectTraits fun_t; typedef BucketElem bucket_t; typedef AlignedVec bucketvec_t; struct Data { Data() : buckets(0), xi(0), scaler(0), cst0(0) {} Data( const T *x // for Direct must persist if xws=NULL , uint32 n , T H , bucket_t *bws // assumed to gave size nb, as computed below , T *xws = NULL // assumed to have size (n+Gap-1). Optional for Direct, unused for DirectCache, required for DirectGap ) : buckets(bws) , scaler(H) , cst0(fun_t::cst0(H, x[0])) { myassert(((bws != NULL) && (isAligned(bws,64))), "bucket pointer not allocated or incorrectly aligned"); uint32 nb = 1 + fun_t::f(H, cst0, x[n-1]); const uint32 npad = Gap-1; const uint32 n_sz = n + npad; // size of padded vector if (xws) { myassert(isAligned(xws,8), "x pointer not allocated or incorrectly aligned"); std::fill_n(xws, npad, x[0]); // pad in front with x[0] std::copy(x, x+n, xws + npad); xi = xws; } else { myassert((Gap==1), "if Gap>1 then X workspace must be provided"); xi = x; } populateIndex(bws, nb, xi, n_sz, scaler, cst0); } const bucket_t *buckets; const T *xi; T scaler; T cst0; // could be x0 or (scaler*x0), depending if we are using FMA or not } data; static T growStep(T H) { T step; T P = next(H); while ((step = P - H) == 0) P = next(P); return step; } static HResults computeH(const T *px, uint32 nx) { myassert((nx > Gap), "Array X too small"); myassert(((Gap == 1) || (Gap == 2)), "Only tested for these values of Gap"); const T x0 = px[0]; const T xN = px[nx-1]; const T range = xN - x0; myassert((range < std::numeric_limits::max()), "range too large"); // check that D_i are strictly increasing and compute minimum value D_{i+Offset}-D_i T deltaDMin = range; for (uint32 i = Gap; i < nx; ++i) { T Dnew = px[i] - x0; T Dold = px[i - Gap] - x0; myassert((Dnew > Dold), "Problem unfeasible: D_i sequence not strictly increasing" << " X[" << 0 << "]=" << x0 << " X[" << i - Gap << "]=" << px[i - Gap] << " X[" << i << "]=" << px[i] << "\n" ); T deltaD = Dnew - Dold; if (deltaD < deltaDMin) deltaDMin = deltaD; } // initial guess for H const T H0 = T(1.0) / deltaDMin; T H = H0; T cst0 = fun_t::cst0(H, x0); fun_t::checkH(H, cst0, xN); // adjust H by trial and error until succeed size_t nInc = 0; bool modified = false; size_t npasses = 0; T step = growStep(H); uint32 seg_already_checked_from = nx; do { myassert((npasses++ < 2), "verification failed\n"); // if there has been an increase, then check only up to that point uint32 last_seg_to_be_checked = seg_already_checked_from - 1; modified = false; uint32 inew = 0; for (uint32 i = Gap; i <= last_seg_to_be_checked; ++i) { uint32 iold = fun_t::f(H, cst0, px[i-Gap]); uint32 inew = fun_t::f(H, cst0, px[i]); while (inew == iold) { seg_already_checked_from = i; last_seg_to_be_checked = nx-1; // everything needs to be checked modified = true; H = H + step; step *= 2; // recalculate all constants and indices cst0 = fun_t::cst0(H, x0); fun_t::checkH(H, cst0, xN); iold = fun_t::f(H, cst0, px[i - Gap]); inew = fun_t::f(H, cst0, px[i]); } } } while (SAFETY_MULTI_PASS && modified); return HResults(H, (((double)H) / H0) - 1.0, nInc); } static void populateIndex(BucketElem *buckets, uint32 index_size, const T *px, uint32 x_size, T scaler, T cst0) { for (uint32 i = x_size-1, b = index_size-1, j=0; ; --i) { uint32 idx = fun_t::f(scaler, cst0, px[i]); while (b > idx) { // in the 1st iteration it is j=0 but this condition is always false buckets[b].set( j, px ); --b; } if (Gap==1 || b == idx) { // if Gap==1, which is known at compile time, the check b==idx is redundant j = i - (Gap-1); // subtracting (Gap-1) points to the index of the first X-element to check buckets[b].set(j, px); if (b-- == 0) break; } } } DirectInfo(const Data& d) : data(d) { } DirectInfo(const T* px, const uint32 n) { HResults res = computeH(px, n); #ifdef PAPER_TEST nInc = res.nInc; hRatio = res.hRatio; #endif const uint32 npad = Gap-1; const uint32 n_sz = n + npad; // size of padded vector if (npad) xi.resize(n_sz); T H = res.H; T cst0 = fun_t::cst0(H, px[0]); const uint32 maxIndex = fun_t::f(H, cst0, px[n-1]); buckets.resize(maxIndex + 1); data = Data(px, n, H, buckets.begin(), (npad? xi.begin(): NULL)); } private: bucketvec_t buckets; AlignedVec xi; #ifdef PAPER_TEST public: double hRatio; size_t nInc; #endif }; } // namespace DirectAux } // namespace Details } // namespace BinSearch bitsandbytes/include/Algo-Direct2.h000066400000000000000000000261301506752665600175450ustar00rootroot00000000000000#pragma once #include "Algo-Direct-Common.h" namespace BinSearch { namespace Details { template struct AlgoScalarBase::value>::type> : DirectAux::DirectInfo<2, T, A> { private: typedef DirectAux::DirectInfo<2, T, A> base_t; static const size_t Offset=2; public: AlgoScalarBase(const T* x, const uint32 n) : base_t(x, n) { } FORCE_INLINE uint32 scalar(T z) const { const T* px = base_t::data.xi; const uint32* buckets = reinterpret_cast(base_t::data.buckets); uint32 bidx = base_t::fun_t::f(base_t::data.scaler, base_t::data.cst0, z); uint32 iidx = buckets[bidx]; px += iidx; if (z < *px) --iidx; if (z < *(px+1)) --iidx; return iidx; } }; template struct AlgoVecBase::value>::type> : AlgoScalarBase { static const uint32 nElem = sizeof(typename InstrFloatTraits::vec_t) / sizeof(T); typedef FVec fVec; typedef IVec i128; struct Constants { fVec vscaler; fVec vcst0; IVec one; }; private: typedef AlgoScalarBase base_t; #ifdef USE_SSE2 FORCE_INLINE //NO_INLINE void resolve(const FVec& vz, const IVec& bidx, uint32 *pr) const { union U { __m128i vec; uint32 ui32[4]; } u; const uint32* buckets = reinterpret_cast(base_t::data.buckets); const float *xi = base_t::data.xi; // read indices t const double *p3 = reinterpret_cast(&xi[(u.ui32[3] = buckets[bidx.get3()])]); const double *p2 = reinterpret_cast(&xi[(u.ui32[2] = buckets[bidx.get2()])]); const double *p1 = reinterpret_cast(&xi[(u.ui32[1] = buckets[bidx.get1()])]); const double *p0 = reinterpret_cast(&xi[(u.ui32[0] = buckets[bidx.get0()])]); #if 0 // read pairs ( X(t-1), X(t) ) __m128 xp3 = _mm_castpd_ps(_mm_load_sd(p3)); __m128 xp2 = _mm_castpd_ps(_mm_load_sd(p2)); __m128 xp1 = _mm_castpd_ps(_mm_load_sd(p1)); __m128 xp0 = _mm_castpd_ps(_mm_load_sd(p0)); // build: // { X(t(0)-1), X(t(1)-1), X(t(2)-1), X(t(3)-1) } // { X(t(0)), X(t(1)), X(t(2)), X(t(3)) } __m128 h13 = _mm_shuffle_ps(xp1, xp3, (1 << 2) + (1 << 6)); __m128 h02 = _mm_shuffle_ps(xp0, xp2, (1 << 2) + (1 << 6)); __m128 u01 = _mm_unpacklo_ps(h02, h13); __m128 u23 = _mm_unpackhi_ps(h02, h13); __m128 vxm = _mm_shuffle_ps(u01, u23, (0) + (1 << 2) + (0 << 4) + (1 << 6)); __m128 vxp = _mm_shuffle_ps(u01, u23, (2) + (3 << 2) + (2 << 4) + (3 << 6)); #else __m128 xp23 = _mm_castpd_ps(_mm_set_pd(*p3, *p2)); __m128 xp01 = _mm_castpd_ps(_mm_set_pd(*p1, *p0)); __m128 vxm = _mm_shuffle_ps(xp01, xp23, (0) + (2 << 2) + (0 << 4) + (2 << 6)); __m128 vxp = _mm_shuffle_ps(xp01, xp23, (1) + (3 << 2) + (1 << 4) + (3 << 6)); #endif IVec i(u.vec); IVec vlem = vz < vxm; IVec vlep = vz < vxp; i = i + vlem + vlep; i.store(pr); } FORCE_INLINE //NO_INLINE void resolve(const FVec& vz, const IVec& bidx, uint32 *pr) const { const uint32* buckets = reinterpret_cast(base_t::data.buckets); const double *xi = base_t::data.xi; uint32 b1 = buckets[bidx.get1()]; uint32 b0 = buckets[bidx.get0()]; const double *p1 = &xi[b1]; const double *p0 = &xi[b0]; // read pairs ( X(t-1), X(t) ) __m128d vx1 = _mm_loadu_pd(p1); __m128d vx0 = _mm_loadu_pd(p0); // build: // { X(t(0)-1), X(t(1)-1) } // { X(t(0)), X(t(1)) } __m128d vxm = _mm_shuffle_pd(vx0, vx1, 0); __m128d vxp = _mm_shuffle_pd(vx0, vx1, 3); IVec i(b1, b0); IVec vlem = (vz < vxm); IVec vlep = (vz < vxp); i = i + vlem + vlep; union { __m128i vec; uint32 ui32[4]; } u; u.vec = i; pr[0] = u.ui32[0]; pr[1] = u.ui32[2]; } #endif // USE_SSE2 #ifdef USE_AVX FORCE_INLINE //NO_INLINE void resolve(const FVec& vz, const IVec& bidx, uint32 *pr) const { const uint32* buckets = reinterpret_cast(base_t::data.buckets); const float *xi = base_t::data.xi; #if 0 // use gather instructions IVec idxm; idxm.setidx(buckets, bidx); __m256i z = _mm256_setzero_si256(); IVec minusone = _mm256_cmpeq_epi32(z,z); IVec idxp = idxm - minusone; FVec vxm = _mm256_i32gather_ps(xi, idxm, sizeof(float)); FVec vxp = _mm256_i32gather_ps(xi, idxp, sizeof(float)); IVec ip = idxm; #else // do not use gather instructions union U { __m256i vec; uint32 ui32[8]; } u; // read indices t const double *p7 = reinterpret_cast(&xi[(u.ui32[7] = buckets[bidx.get7()])]); const double *p6 = reinterpret_cast(&xi[(u.ui32[6] = buckets[bidx.get6()])]); const double *p5 = reinterpret_cast(&xi[(u.ui32[5] = buckets[bidx.get5()])]); const double *p4 = reinterpret_cast(&xi[(u.ui32[4] = buckets[bidx.get4()])]); const double *p3 = reinterpret_cast(&xi[(u.ui32[3] = buckets[bidx.get3()])]); const double *p2 = reinterpret_cast(&xi[(u.ui32[2] = buckets[bidx.get2()])]); const double *p1 = reinterpret_cast(&xi[(u.ui32[1] = buckets[bidx.get1()])]); const double *p0 = reinterpret_cast(&xi[(u.ui32[0] = buckets[bidx.get0()])]); #if 0 // perform 8 loads in double precision // read pairs ( X(t-1), X(t) ) __m128 xp7 = _mm_castpd_ps(_mm_load_sd(p7)); __m128 xp6 = _mm_castpd_ps(_mm_load_sd(p6)); __m128 xp5 = _mm_castpd_ps(_mm_load_sd(p5)); __m128 xp4 = _mm_castpd_ps(_mm_load_sd(p4)); __m128 xp3 = _mm_castpd_ps(_mm_load_sd(p3)); __m128 xp2 = _mm_castpd_ps(_mm_load_sd(p2)); __m128 xp1 = _mm_castpd_ps(_mm_load_sd(p1)); __m128 xp0 = _mm_castpd_ps(_mm_load_sd(p0)); // build: // { X(t(0)-1), X(t(1)-1), X(t(2)-1), X(t(3)-1) } // { X(t(0)), X(t(1)), X(t(2)), X(t(3)) } __m128 h57 = _mm_shuffle_ps(xp5, xp7, (1 << 2) + (1 << 6)); // F- F+ H- H+ __m128 h46 = _mm_shuffle_ps(xp4, xp6, (1 << 2) + (1 << 6)); // E- E+ G- G+ __m128 h13 = _mm_shuffle_ps(xp1, xp3, (1 << 2) + (1 << 6)); // B- B+ D- D+ __m128 h02 = _mm_shuffle_ps(xp0, xp2, (1 << 2) + (1 << 6)); // A- A+ C- C+ __m128 u01 = _mm_unpacklo_ps(h02, h13); // A- B- A+ B+ __m128 u23 = _mm_unpackhi_ps(h02, h13); // C- D- C+ D+ __m128 u45 = _mm_unpacklo_ps(h46, h57); // E- F- E+ F+ __m128 u67 = _mm_unpackhi_ps(h46, h57); // G- H- G+ H+ __m128 abcdm = _mm_shuffle_ps(u01, u23, (0) + (1 << 2) + (0 << 4) + (1 << 6)); // A- B- C- D- __m128 abcdp = _mm_shuffle_ps(u01, u23, (2) + (3 << 2) + (2 << 4) + (3 << 6)); // A+ B+ C+ D+ __m128 efghm = _mm_shuffle_ps(u45, u67, (0) + (1 << 2) + (0 << 4) + (1 << 6)); // E- F- G- H- __m128 efghp = _mm_shuffle_ps(u45, u67, (2) + (3 << 2) + (2 << 4) + (3 << 6)); // E+ F+ G+ H+ FVec vxp = _mm256_insertf128_ps(_mm256_castps128_ps256(abcdm), efghm, 1); FVec vxm = _mm256_insertf128_ps(_mm256_castps128_ps256(abcdp), efghp, 1); IVec ip(u.vec); #else // use __mm256_set_pd // read pairs ( X(t-1), X(t) ) __m256 x0145 = _mm256_castpd_ps(_mm256_set_pd(*p5, *p4, *p1, *p0)); // { x0(t-1), x0(t), x1(t-1), x1(t), x4(t-1), x4(t), x5(t-1), x5(t) } __m256 x2367 = _mm256_castpd_ps(_mm256_set_pd(*p7, *p6, *p3, *p2)); // { x2(t-1), x2(t), x3(t-1), x3(t), x6(t-1), x6(t), x7(t-1), x7(t) } // { x0(t-1), x1(t-1), x2(t-1), 3(t-1, x4(t-1), x5(t-1), x6(t-1), xt(t-1) } FVec vxm = _mm256_shuffle_ps(x0145, x2367, 0 + (2 << 2) + (0 << 4) + (2 << 6) ); // { x0(t), x1(t), x2(t), 3(t, x4(t), x5(t), x6(t), xt(t) } FVec vxp = _mm256_shuffle_ps(x0145, x2367, 1 + (3 << 2) + (1 << 4) + (3 << 6) ); IVec ip(u.vec); #endif #endif IVec vlem = vz < vxm; IVec vlep = vz < vxp; ip = ip + vlem + vlep; ip.store(pr); } FORCE_INLINE //NO_INLINE void resolve(const FVec& vz, const IVec& bidx, uint32 *pr) const { union { __m256i vec; uint64 ui64[4]; } u; const uint32* buckets = reinterpret_cast(base_t::data.buckets); const double *xi = base_t::data.xi; // read indices t const double *p3 = &xi[(u.ui64[3] = buckets[bidx.get3()])]; const double *p2 = &xi[(u.ui64[2] = buckets[bidx.get2()])]; const double *p1 = &xi[(u.ui64[1] = buckets[bidx.get1()])]; const double *p0 = &xi[(u.ui64[0] = buckets[bidx.get0()])]; // read pairs ( X(t-1), X(t) ) __m128d xp3 = _mm_loadu_pd(p3); __m128d xp2 = _mm_loadu_pd(p2); __m128d xp1 = _mm_loadu_pd(p1); __m128d xp0 = _mm_loadu_pd(p0); // build: // { X(t(0)-1), X(t(1)-1), X(t(2)-1), X(t(3)-1) } // { X(t(0)), X(t(1)), X(t(2)), X(t(3)) } __m256d x02 = _mm256_insertf128_pd(_mm256_castpd128_pd256(xp0), xp2, 1); __m256d x13 = _mm256_insertf128_pd(_mm256_castpd128_pd256(xp1), xp3, 1); FVec vxm = _mm256_unpacklo_pd(x02,x13); FVec vxp = _mm256_unpackhi_pd(x02,x13); // __m128d h01m = _mm_shuffle_pd(xp0, xp1, 0); // __m128d h23m = _mm_shuffle_pd(xp2, xp3, 0); // __m128d h01p = _mm_shuffle_pd(xp0, xp1, 3); // __m128d h23p = _mm_shuffle_pd(xp2, xp3, 3); // FVec vxm = _mm256_insertf128_pd(_mm256_castpd128_pd256(h01m), h23m, 1); // FVec vxp = _mm256_insertf128_pd(_mm256_castpd128_pd256(h01p), h23p, 1); IVec i(u.vec); IVec vlem = vz < vxm; IVec vlep = vz < vxp; i = i + vlem + vlep; i.extractLo32s().store(pr); } #endif public: AlgoVecBase(const T* x, const uint32 n) : base_t(x, n) {} void initConstants(Constants& cst) const { cst.vscaler.setN(base_t::data.scaler); cst.vcst0.setN(base_t::data.cst0); cst.one.setN(uint32(1)); } void vectorial(uint32 *pr, const T *pz, const Constants& cst) const { fVec vz(pz); resolve(vz, base_t::fun_t::f(cst.vscaler, cst.vcst0, vz), pr); } }; } // namespace Details } // namespace BinSearch bitsandbytes/include/AlgoXCodes.h000066400000000000000000000007731506752665600173660ustar00rootroot00000000000000ALGOENUM(DirectCacheFMA, 5) ALGOENUM(DirectFMA, 15) ALGOENUM(Direct2FMA, 25) ALGOENUM(DirectCache, 10) ALGOENUM(Direct, 20) ALGOENUM(Direct2, 30) ALGOENUM(Nonary, 40) ALGOENUM(Pentary, 50) ALGOENUM(Ternary, 60) ALGOENUM(Eytzinger, 70) ALGOENUM(BitSet, 80) ALGOENUM(ClassicOffset, 90) #ifdef PAPER_TEST ALGOENUM(MorinOffset, 100) ALGOENUM(BitSetNoPad, 110) ALGOENUM(ClassicMod, 120) ALGOENUM(MorinBranchy, 130) ALGOENUM(Classic, 140) ALGOENUM(LowerBound, 145) #ifdef USE_MKL ALGOENUM(MKL, 150) #endif #endif bitsandbytes/include/BinAlgo.h000066400000000000000000000045171506752665600167110ustar00rootroot00000000000000#pragma once #include "Type.h" #include namespace BinSearch { template struct BinAlgo : Details::BinAlgoBase { typedef Details::BinAlgoBase base_t; BinAlgo(const T* px, const uint32 n) : base_t(px, n), x0(px[0]), xN(px[n-1]), N(n) {} BinAlgo(const T* px, const uint32 n, const typename base_t::Data& d) : base_t(d), x0(px[0]), xN(px[n-1]), N(n) {} FORCE_INLINE uint32 scalar(T z) const { if (!L || z >= x0) if (!R || z < xN) return base_t::scalar(z); else return N; else return std::numeric_limits::max(); } FORCE_INLINE void vectorial(uint32 *pr, const T *pz, uint32 n) const { if (!L && !R) { Details::Loop::loop(*this, pr, pz, n); } else { const uint32 nElem = base_t::nElem; const uint32 idealbufsize = 256; const uint32 bufsize = nElem * (idealbufsize / nElem + ((idealbufsize % nElem) ? 1 : 0)); T databuf[bufsize]; uint32 resbuf[bufsize]; uint32 indexbuf[bufsize]; uint32 *prend = pr + n; while(pr != prend) { uint32 cnt = 0; uint32 niter = std::min(bufsize, (uint32)std::distance(pr,prend)); for (uint32 j = 0; j < niter; ++j) { T z = pz[j]; // FIXME: use SSE2? if (!L || z >= x0) if (!R || z < xN) { databuf[cnt] = z; indexbuf[cnt] = j; ++cnt; } else pr[j] = N; else pr[j] = std::numeric_limits::max(); } // FIXME: merge these two loops Details::Loop::loop(*this, resbuf, databuf, cnt); for (uint32 j = 0; j < cnt; ++j) pr[indexbuf[j]] = resbuf[j]; pr += niter; pz += niter; } } } Details::CondData x0; Details::CondData xN; Details::CondData N; }; } // namespace BinSearch bitsandbytes/include/BinSearch.h000066400000000000000000000002151506752665600172230ustar00rootroot00000000000000#pragma once #include "AAlloc.h" #include "BinAlgo.h" #include "SIMD.h" #include #include #include "Algo-Direct2.h" bitsandbytes/include/Portable.h000066400000000000000000000102731506752665600171420ustar00rootroot00000000000000#pragma once #include #include #include #include #if defined(__aarch64__) #ifdef __CUDACC__ #undef USE_NEON // Doesn't work with nvcc, undefined symbols #else #include #undef USE_NEON // Not yet implemented #endif #undef USE_AVX // x86_64 only #undef USE_AVX2 // x86_64 only #undef USE_SSE2 // x86_64 only #undef USE_SSE41 // x86_64 only #undef USE_SSE42 // x86_64 only #undef USE_FMA // x86_64 only #ifdef USE_NEON typedef float32x4_t __m128; typedef int32x4_t __m128i; typedef float64x2_t __m128d; #else typedef struct {float a; float b; float c; float d;} __m128; typedef struct {int a; int b; int c; int d;} __m128i; typedef struct {double a; double b;} __m128d; #endif #elif defined(__powerpc64__) #ifdef __CUDACC__ #undef USE_VSX // Doesn't work with nvcc, undefined symbols #else #include #undef USE_VSX // Not yet implemented #endif #undef USE_AVX // x86_64 only #undef USE_AVX2 // x86_64 only #undef USE_SSE2 // x86_64 only #undef USE_SSE41 // x86_64 only #undef USE_SSE42 // x86_64 only #undef USE_FMA // x86_64 only #ifdef USE_VSX typedef vector float __m128; typedef vector signed int __m128i; typedef vector double __m128d; #else typedef struct {float a; float b; float c; float d;} __m128; typedef struct {int a; int b; int c; int d;} __m128i; typedef struct {double a; double b;} __m128d; #endif #else #undef USE_NEON // ARM64 only #undef USE_VSX // PPC only #ifdef __FMA__ #define USE_FMA #endif #if !defined(__SSE2__) && !defined(_MSC_VER) #error Compiler must support SSE2 #endif #define USE_SSE2 #if defined(__aarch64__) #else #ifdef __AVX2__ #define USE_AVX2 #endif #ifdef __AVX__ #define USE_AVX #endif #ifdef __SSE4_1__ #define USE_SSE41 #endif #ifdef __SSE4_2__ #define USE_SSE42 #endif #endif #endif #ifndef _MSC_VER #include #endif namespace BinSearch { #ifndef _MSC_VER typedef int8_t int8; typedef uint8_t uint8; typedef int32_t int32; typedef uint32_t uint32; typedef int64_t int64; typedef uint64_t uint64; #else typedef __int8 int8; typedef unsigned __int8 uint8; typedef __int32 int32; typedef unsigned __int32 uint32; typedef __int64 int64; typedef unsigned __int64 uint64; #endif namespace Details { #define myassert(cond, msg) if (!cond){ std::ostringstream os; os << "\nassertion failed: " << #cond << ", " << msg << "\n"; throw std::invalid_argument(os.str()); } // log2 is not defined in VS2008 #if defined(_MSC_VER) inline uint32 log2 (uint32 val) { if (val == 1) return 0; uint32 ret = 0; do { ret++; val >>= 1; } while (val > 1); return ret; } #endif #ifdef _DEBUG #define DEBUG #endif #ifdef _MSC_VER # define FORCE_INLINE __forceinline # define NO_INLINE __declspec(noinline) #else # define NO_INLINE __attribute__((noinline)) # ifdef DEBUG # define FORCE_INLINE NO_INLINE # else # define FORCE_INLINE __attribute__((always_inline)) inline # endif #endif #ifdef USE_AVX #define COMISS "vcomiss" #define COMISD "vcomisd" #else #define COMISS "comiss" #define COMISD "comisd" #endif // nextafter is not defined in VS2008 #if defined(_MSC_VER) && (_MSC_VER <= 1500) #include inline float mynext(float x) { return _nextafterf(x, std::numeric_limits::max()); } inline double mynext(double x) { return _nextafter(x, std::numeric_limits::max()); } inline float myprev(float x) { return _nextafterf(x, -std::numeric_limits::max()); } inline double myprev(double x) { return _nextafter(x, -std::numeric_limits::max()); } #else inline float mynext(float x) { return std::nextafterf(x, std::numeric_limits::max()); } inline double mynext(double x) { return std::nextafter(x, std::numeric_limits::max()); } inline float myprev(float x) { return std::nextafterf(x, -std::numeric_limits::max()); } inline double myprev(double x) { return std::nextafter(x, -std::numeric_limits::max()); } #endif template inline T next(T x) { for (int i = 0; i < 4; ++i) x = mynext(x); return x; } template inline T prev(T x) { for (int i = 0; i < 4; ++i) x = myprev(x); return x; } } // namespace Details } // namespace BinSearch bitsandbytes/include/SIMD.h000066400000000000000000000577551506752665600161460ustar00rootroot00000000000000#pragma once #include "Portable.h" #ifdef USE_SSE2 #include #if defined(USE_AVX) || defined(USE_AVX2) #include #else #ifdef USE_SSE41 #include #endif #endif #endif namespace BinSearch { namespace Details { template struct FTOITraits{}; template struct FVec; template struct IVec; template struct FVec1; template <> struct InstrFloatTraits { typedef __m128 vec_t; }; template <> struct InstrFloatTraits { typedef __m128d vec_t; }; } } #if !defined(__aarch64__) && !defined(__powerpc64__) #ifdef USE_SSE42 #ifndef _MSC_VER #include #define popcnt32 _mm_popcnt_u32 #else #include #define popcnt32 __popcnt #endif #else // USE_SSE42 namespace BinSearch { FORCE_INLINE int popcnt32(int x32) { // strictly speaking this is not correct, as it ignores higher order bits // however, this is only used on the resuot of movemask on a 128-bit register, which is 8 at most, so it is ok // with 256-bit registers, SSE42 is defined, and we do not use this function uint8 x = static_cast(x32); x = (x & 0x55) + (x >> 1 & 0x55); x = (x & 0x33) + (x >> 2 & 0x33); x = (x & 0x0f) + (x >> 4 & 0x0f); return x; } } // namespace #endif #include "Type.h" namespace BinSearch { namespace Details { template <> struct InstrIntTraits { typedef __m128i vec_t; }; template <> struct InstrFloatTraits { typedef __m128 vec_t; }; template <> struct InstrFloatTraits { typedef __m128d vec_t; }; template <> struct FTOITraits { typedef IVec vec_t; }; #ifdef USE_AVX template <> struct FTOITraits { typedef IVec vec_t; }; template <> struct InstrIntTraits { typedef __m256i vec_t; }; template <> struct InstrFloatTraits { typedef __m256 vec_t; }; template <> struct InstrFloatTraits { typedef __m256d vec_t; }; #endif template struct VecStorage { typedef typename TR::vec_t vec_t; FORCE_INLINE operator vec_t&() { return vec; } FORCE_INLINE operator const vec_t&() const { return vec; } protected: FORCE_INLINE VecStorage() {} FORCE_INLINE VecStorage(const vec_t& v) : vec( v ) {} vec_t vec; }; template struct IVecBase; template <> struct IVecBase : VecStorage> { protected: FORCE_INLINE IVecBase() {} FORCE_INLINE IVecBase( const vec_t& v) : VecStorage>( v ) {} public: FORCE_INLINE static vec_t zero() { return _mm_setzero_si128(); } FORCE_INLINE int32 get0() const { return _mm_cvtsi128_si32( vec ); } FORCE_INLINE void assignIf( const vec_t& val, const vec_t& mask ) { #ifdef USE_SSE41 vec = _mm_blendv_epi8(vec, val, mask); #else vec = _mm_or_si128(_mm_andnot_si128(mask,vec), _mm_and_si128(mask,val)); #endif } FORCE_INLINE void orIf(const vec_t& val, const vec_t& mask) { vec = _mm_or_si128(vec, _mm_and_si128(val,mask)); } }; template <> struct IVec : IVecBase { FORCE_INLINE IVec() {} FORCE_INLINE IVec( int32 i ) : IVecBase( _mm_set1_epi32( i ) ) {} FORCE_INLINE IVec( const vec_t& v) : IVecBase( v ) {} FORCE_INLINE IVec( uint32 u3, uint32 u2, uint32 u1, uint32 u0) : IVecBase( _mm_set_epi32( u3, u2, u1, u0 ) ) {} void setN( int32 i ) { vec = _mm_set1_epi32( i ); } #ifdef USE_SSE41 FORCE_INLINE int32 get1() const { return _mm_extract_epi32(vec, 1); } FORCE_INLINE int32 get2() const { return _mm_extract_epi32(vec, 2); } FORCE_INLINE int32 get3() const { return _mm_extract_epi32(vec, 3); } #else FORCE_INLINE int32 get1() const { return _mm_cvtsi128_si32( _mm_shuffle_epi32( vec, 1 ) ); } FORCE_INLINE int32 get2() const { return _mm_cvtsi128_si32( _mm_shuffle_epi32( vec, 2 ) ); } FORCE_INLINE int32 get3() const { return _mm_cvtsi128_si32( _mm_shuffle_epi32( vec, 3 ) ); } #endif FORCE_INLINE void store( uint32 *pi ) const { _mm_storeu_si128( reinterpret_cast(pi), vec ); } FORCE_INLINE int countbit() { return popcnt32(_mm_movemask_ps(_mm_castsi128_ps(vec))); } }; template <> struct IVec : IVecBase { FORCE_INLINE IVec() {} FORCE_INLINE IVec( int32 i ) : IVecBase( _mm_set1_epi64x( i ) ) {} FORCE_INLINE IVec( const vec_t& v) : IVecBase( v ) {} FORCE_INLINE IVec( uint64 u1, uint64 u0 ) : IVecBase( _mm_set_epi64x(u1, u0) ) {} void setN( int32 i ) { vec = _mm_set1_epi64x( i ); } FORCE_INLINE int32 get1() const { #ifdef USE_SSE41 return _mm_extract_epi32(vec, 2); #else return _mm_cvtsi128_si32( _mm_shuffle_epi32( vec, 2 ) ); #endif } // extract the 2 32 bits integers no. 0, 2 and store them in a __m128i FORCE_INLINE IVec extractLo32s() const { return _mm_shuffle_epi32(vec, ((2 << 2) | 0)); } FORCE_INLINE void store( uint32 *pi ) const { pi[0] = get0(); pi[1] = get1(); } FORCE_INLINE int countbit() { #if 1 // takes 4 cycles __m128i hi = _mm_shuffle_epi32(vec, 2); // 1 cycle __m128i s = _mm_add_epi32(vec, hi); int32 x = _mm_cvtsi128_si32(s); return -x; #else // takes 6 cycles return popcnt32(_mm_movemask_pd(_mm_castsi128_pd(vec))); #endif } }; template FORCE_INLINE IVec operator>> (const IVec& a, unsigned n) { return _mm_srli_epi32(a, n); } template FORCE_INLINE IVec operator<< (const IVec& a, unsigned n) { return _mm_slli_epi32(a, n); } template FORCE_INLINE IVec operator& (const IVec& a, const IVec& b ) { return _mm_and_si128( a, b ); } template FORCE_INLINE IVec operator| (const IVec& a, const IVec& b ) { return _mm_or_si128( a, b ); } template FORCE_INLINE IVec operator^ (const IVec& a, const IVec& b ) { return _mm_xor_si128( a, b ); } template FORCE_INLINE IVec operator+ (const IVec& a, const IVec& b ) { return _mm_add_epi32( a, b ); } template FORCE_INLINE IVec operator- (const IVec& a, const IVec& b ) { return _mm_sub_epi32( a, b ); } #ifdef USE_SSE41 template FORCE_INLINE IVec min (const IVec& a, const IVec& b ) { return _mm_min_epi32( a, b ); } #endif typedef VecStorage> FVec128Float; template <> struct FVec1 : FVec128Float { FORCE_INLINE FVec1() {} FORCE_INLINE FVec1( float f ) : FVec128Float( _mm_load_ss( &f ) ) {} FORCE_INLINE FVec1( const vec_t& v ): FVec128Float( v ) {} FORCE_INLINE float get0() const { return _mm_cvtss_f32( vec ); } }; template <> struct FVec : FVec128Float { FORCE_INLINE FVec() {} FORCE_INLINE FVec( float f ) : FVec128Float( _mm_set1_ps( f ) ) {} FORCE_INLINE FVec( const float *v ) : FVec128Float( _mm_loadu_ps( v ) ) {} FORCE_INLINE FVec( const vec_t& v) : FVec128Float(v) {} FORCE_INLINE FVec( float f3, float f2, float f1, float f0 ) : FVec128Float( _mm_set_ps(f3, f2, f1, f0) ) {} void set0( float f ) { vec = _mm_load_ss( &f ); } void setN( float f ) { vec = _mm_set1_ps( f ); } FORCE_INLINE void setidx( const float *xi, const IVec& idx ) { uint32 i0 = idx.get0(); uint32 i1 = idx.get1(); uint32 i2 = idx.get2(); uint32 i3 = idx.get3(); vec = _mm_set_ps( xi[i3], xi[i2], xi[i1], xi[i0] ); } FORCE_INLINE float get0() const { return _mm_cvtss_f32( vec ); } FORCE_INLINE float get1() const { return _mm_cvtss_f32( _mm_shuffle_ps( vec, vec, 1 ) ); } FORCE_INLINE float get2() const { return _mm_cvtss_f32( _mm_shuffle_ps( vec, vec, 2 ) ); } FORCE_INLINE float get3() const { return _mm_cvtss_f32( _mm_shuffle_ps( vec, vec, 3 ) ); } }; FORCE_INLINE FVec1 operator+ (const FVec1& a, const FVec1& b) { return _mm_add_ss( a, b ); } FORCE_INLINE FVec1 operator- (const FVec1& a, const FVec1& b) { return _mm_sub_ss( a, b ); } FORCE_INLINE FVec1 operator* (const FVec1& a, const FVec1& b) { return _mm_mul_ss( a, b ); } FORCE_INLINE FVec1 operator/ (const FVec1& a, const FVec1& b) { return _mm_div_ss( a, b ); } FORCE_INLINE int ftoi (const FVec1& a) { return _mm_cvttss_si32(a); } FORCE_INLINE IVec operator> (const FVec1& a, const FVec1& b) { return _mm_castps_si128( _mm_cmpgt_ss( a, b ) ); } #ifdef USE_FMA FORCE_INLINE FVec1 mulSub(const FVec1& a, const FVec1& b, const FVec1& c) { return _mm_fmsub_ss(a, b, c); } #endif FORCE_INLINE FVec operator- (const FVec& a, const FVec& b) { return _mm_sub_ps( a, b ); } FORCE_INLINE FVec operator* (const FVec& a, const FVec& b) { return _mm_mul_ps( a, b ); } FORCE_INLINE FVec operator/ (const FVec& a, const FVec& b) { return _mm_div_ps( a, b ); } FORCE_INLINE IVec ftoi (const FVec& a) { return _mm_cvttps_epi32(a); } #ifndef __clang__ // Conflicts with builtin operator FORCE_INLINE IVec operator<= (const FVec& a, const FVec& b) { return _mm_castps_si128( _mm_cmple_ps( a, b ) ); } FORCE_INLINE IVec operator>= (const FVec& a, const FVec& b) { return _mm_castps_si128( _mm_cmpge_ps( a, b ) ); } FORCE_INLINE IVec operator< (const FVec& a, const FVec& b) { return _mm_castps_si128(_mm_cmplt_ps(a, b)); } #endif #ifdef USE_FMA FORCE_INLINE FVec mulSub(const FVec& a, const FVec& b, const FVec& c) { return _mm_fmsub_ps(a, b, c); } #endif typedef VecStorage> FVec128Double; template <> struct FVec1 : FVec128Double { FORCE_INLINE FVec1() {} FORCE_INLINE FVec1( double f ) : FVec128Double( _mm_load_sd( &f ) ) {} FORCE_INLINE FVec1( const vec_t& v ) : FVec128Double( v ) {} FORCE_INLINE double get0() const { return _mm_cvtsd_f64( vec ); } }; template <> struct FVec : FVec128Double { FORCE_INLINE FVec() {} FORCE_INLINE FVec( double d ) : FVec128Double( _mm_set1_pd( d ) ) {} FORCE_INLINE FVec( const double *v ) : FVec128Double( _mm_loadu_pd( v ) ) {} FORCE_INLINE FVec( const vec_t& v) : FVec128Double( v ) {} FORCE_INLINE FVec( double f1, double f0 ) : FVec128Double( _mm_set_pd(f1, f0) ) {} void set0( double f ) { vec = _mm_load_sd( &f ); } void setN( double f ) { vec = _mm_set1_pd( f ); } FORCE_INLINE void setidx( const double *xi, const IVec& idx ) { vec = _mm_set_pd( xi[idx.get1()], xi[idx.get0()] ); } FORCE_INLINE double get0() const { return _mm_cvtsd_f64( vec ); } FORCE_INLINE double get1() const { return _mm_cvtsd_f64( _mm_shuffle_pd( vec, vec, 1 ) ); }; }; FORCE_INLINE FVec1 operator+ (const FVec1& a, const FVec1& b) { return _mm_add_sd( a, b ); } FORCE_INLINE FVec1 operator- (const FVec1& a, const FVec1& b) { return _mm_sub_sd( a, b ); } FORCE_INLINE FVec1 operator* (const FVec1& a, const FVec1& b) { return _mm_mul_sd( a, b ); } FORCE_INLINE FVec1 operator/ (const FVec1& a, const FVec1& b) { return _mm_div_sd( a, b ); } FORCE_INLINE int ftoi (const FVec1& a) { return _mm_cvttsd_si32(a); } FORCE_INLINE IVec operator> (const FVec1& a, const FVec1& b) { return _mm_castpd_si128( _mm_cmpgt_sd( a, b ) ); } #ifdef USE_FMA FORCE_INLINE FVec1 mulSub(const FVec1& a, const FVec1& b, const FVec1& c) { return _mm_fmsub_sd(a, b, c); } #endif FORCE_INLINE FVec operator- (const FVec& a, const FVec& b) { return _mm_sub_pd( a, b ); } FORCE_INLINE FVec operator* (const FVec& a, const FVec& b) { return _mm_mul_pd( a, b ); } FORCE_INLINE FVec operator/ (const FVec& a, const FVec& b) { return _mm_div_pd( a, b ); } FORCE_INLINE IVec ftoi (const FVec& a) { return _mm_cvttpd_epi32(a); } #ifndef __clang__ // Conflicts with builtin operator FORCE_INLINE IVec operator<= (const FVec& a, const FVec& b) { return _mm_castpd_si128( _mm_cmple_pd( a, b ) ); } FORCE_INLINE IVec operator< (const FVec& a, const FVec& b) { return _mm_castpd_si128(_mm_cmplt_pd(a, b)); } FORCE_INLINE IVec operator>= (const FVec& a, const FVec& b) { return _mm_castpd_si128( _mm_cmpge_pd( a, b ) ); } #endif #ifdef USE_FMA FORCE_INLINE FVec mulSub(const FVec& a, const FVec& b, const FVec& c ) { return _mm_fmsub_pd(a, b, c); } #endif #ifdef USE_AVX template <> struct IVecBase : VecStorage> { protected: FORCE_INLINE IVecBase() {} FORCE_INLINE IVecBase( const vec_t& v) : VecStorage>( v ) {} public: FORCE_INLINE static vec_t zero() { return _mm256_setzero_si256(); } FORCE_INLINE int32 get0() const { return _mm_cvtsi128_si32(_mm256_castsi256_si128(vec)); } FORCE_INLINE void assignIf( const vec_t& val, const vec_t& mask ) { vec = _mm256_blendv_epi8(vec, val, mask); } FORCE_INLINE void orIf(const vec_t& val, const vec_t& mask) { vec = _mm256_blendv_epi8(vec, val, mask); //vec = _mm256_or_si256(vec, _mm256_and_si256(val,mask)); } FORCE_INLINE __m128i lo128() const { return _mm256_castsi256_si128(vec); } FORCE_INLINE __m128i hi128() const { return _mm256_extractf128_si256(vec, 1); } }; template <> struct IVec : IVecBase { FORCE_INLINE IVec() {} FORCE_INLINE IVec( int32 i ) : IVecBase( _mm256_set1_epi32( i ) ) {} FORCE_INLINE IVec( const vec_t& v) : IVecBase( v ) {} FORCE_INLINE IVec(uint32 u7, uint32 u6, uint32 u5, uint32 u4, uint32 u3, uint32 u2, uint32 u1, uint32 u0) : IVecBase(_mm256_set_epi32(u7, u6, u5, u4, u3, u2, u1, u0)) {} void setN( int32 i ) { vec = _mm256_set1_epi32( i ); } FORCE_INLINE int32 get1() const { return _mm256_extract_epi32(vec, 1); } FORCE_INLINE int32 get2() const { return _mm256_extract_epi32(vec, 2); } FORCE_INLINE int32 get3() const { return _mm256_extract_epi32(vec, 3); } FORCE_INLINE int32 get4() const { return _mm256_extract_epi32(vec, 4); } FORCE_INLINE int32 get5() const { return _mm256_extract_epi32(vec, 5); } FORCE_INLINE int32 get6() const { return _mm256_extract_epi32(vec, 6); } FORCE_INLINE int32 get7() const { return _mm256_extract_epi32(vec, 7); } FORCE_INLINE void setidx( const uint32 *bi, const IVec& idx ) { vec = _mm256_i32gather_epi32(reinterpret_cast(bi), idx, sizeof(uint32)); } FORCE_INLINE void store( uint32 *pi ) const { _mm256_storeu_si256( reinterpret_cast(pi), vec ); } FORCE_INLINE int countbit() { return popcnt32(_mm256_movemask_ps(_mm256_castsi256_ps(vec))); } }; template <> struct IVec : IVecBase { FORCE_INLINE IVec() {} FORCE_INLINE IVec( int32 i ) : IVecBase( _mm256_set1_epi64x( i ) ) {} FORCE_INLINE IVec( const vec_t& v) : IVecBase( v ) {} FORCE_INLINE IVec(uint64 u3, uint64 u2, uint64 u1, uint64 u0) : IVecBase(_mm256_set_epi64x(u3, u2, u1, u0)) {} void setN( int32 i ) { vec = _mm256_set1_epi64x( i ); } // extract the 4 32 bits integers no. 0, 2, 4, 6 and store them in a __m128i FORCE_INLINE IVec extractLo32s() const { union { uint32 u32[4]; __m128i u; } mask = {0,2,4,6}; //__m256 ps256 = _mm256_castsi256_ps(vec); //__m128 lo128 = _mm256_castps256_ps128(ps256); //__m128 hi128 = _mm256_extractf128_ps(ps256, 1); //__m128 blend = _mm_shuffle_ps(lo128, hi128, 0 + (2<<2) + (0<<4) + (2<<6)); __m256i blend = _mm256_permutevar8x32_epi32(vec, _mm256_castsi128_si256(mask.u)); return _mm256_castsi256_si128(blend); } //int32 get1() const { return _mm256_cvtsi256_si32( _mm256_shuffle_epi32( vec, 2 ) ); }; FORCE_INLINE int32 get1() const { return _mm256_extract_epi32(vec, 2); } FORCE_INLINE void store( uint32 *pi ) const { extractLo32s().store(pi); } FORCE_INLINE int countbit() { return popcnt32(_mm256_movemask_pd(_mm256_castsi256_pd(vec))); } }; template FORCE_INLINE IVec operator>> (const IVec& a, unsigned n) { return _mm256_srli_epi32(a, n); } template FORCE_INLINE IVec operator<< (const IVec& a, unsigned n) { return _mm256_slli_epi32(a, n); } template FORCE_INLINE IVec operator& (const IVec& a, const IVec& b ) { return _mm256_and_si256( a, b ); } template FORCE_INLINE IVec operator| (const IVec& a, const IVec& b ) { return _mm256_or_si256( a, b ); } template FORCE_INLINE IVec operator^ (const IVec& a, const IVec& b ) { return _mm256_xor_si256( a, b ); } template FORCE_INLINE IVec min (const IVec& a, const IVec& b ) { return _mm256_min_epi32( a, b ); } FORCE_INLINE IVec operator+ (const IVec& a, const IVec& b ) { return _mm256_add_epi32( a, b ); } FORCE_INLINE IVec operator- (const IVec& a, const IVec& b ) { return _mm256_sub_epi32( a, b ); } FORCE_INLINE IVec operator+ (const IVec& a, const IVec& b ) { return _mm256_add_epi64( a, b ); } FORCE_INLINE IVec operator- (const IVec& a, const IVec& b ) { return _mm256_sub_epi64( a, b ); } typedef VecStorage> FVec256Float; template <> struct FVec : FVec256Float { FORCE_INLINE FVec() {} FORCE_INLINE FVec( float f ) : FVec256Float( _mm256_set1_ps( f ) ) {} FORCE_INLINE FVec( const float *v ) : FVec256Float( _mm256_loadu_ps( v ) ) {} FORCE_INLINE FVec( const vec_t& v) : FVec256Float(v) {} FORCE_INLINE FVec(float f7, float f6, float f5, float f4, float f3, float f2, float f1, float f0) : FVec256Float(_mm256_set_ps(f7, f6, f5, f4, f3, f2, f1, f0)) {} //void set0( float f ) { vec = _mm256_load_ss( &f ); } void setN( float f ) { vec = _mm256_set1_ps( f ); } FORCE_INLINE void setidx( const float *xi, const IVec& idx ) { #if 1 // use gather primitives vec = _mm256_i32gather_ps (xi, idx, 4); #elif 0 uint32 i0 = idx.get0(); uint32 i1 = idx.get1(); uint32 i2 = idx.get2(); uint32 i3 = idx.get3(); uint32 i4 = idx.get4(); uint32 i5 = idx.get5(); uint32 i6 = idx.get6(); uint32 i7 = idx.get7(); vec = _mm256_set_ps( xi[i7], xi[i6], xi[i5], xi[i4], xi[i3], xi[i2], xi[i1], xi[i0] ); #else union { __m256i vec; uint32 ui32[8]; } i; i.vec = static_cast(idx); vec = _mm256_set_ps(xi[i.ui32[7]], xi[i.ui32[6]], xi[i.ui32[5]], xi[i.ui32[4]], xi[i.ui32[3]], xi[i.ui32[2]], xi[i.ui32[1]], xi[i.ui32[0]]); #endif } FORCE_INLINE FVec lo128() const { return _mm256_castps256_ps128(vec); } FORCE_INLINE FVec hi128() const { return _mm256_extractf128_ps(vec, 1); } //FORCE_INLINE float get0() const { return _mm256_cvtss_f32( vec ); } //FORCE_INLINE float get1() const { return _mm256_cvtss_f32( _mm256_shuffle_ps( vec, vec, 1 ) ); } //FORCE_INLINE float get2() const { return _mm256_cvtss_f32( _mm256_shuffle_ps( vec, vec, 2 ) ); } //FORCE_INLINE float get3() const { return _mm256_cvtss_f32( _mm256_shuffle_ps( vec, vec, 3 ) ); } }; FORCE_INLINE FVec operator- (const FVec& a, const FVec& b) { return _mm256_sub_ps( a, b ); } FORCE_INLINE FVec operator* (const FVec& a, const FVec& b) { return _mm256_mul_ps( a, b ); } FORCE_INLINE FVec operator/ (const FVec& a, const FVec& b) { return _mm256_div_ps( a, b ); } FORCE_INLINE IVec ftoi (const FVec& a) { return _mm256_cvttps_epi32(a); } FORCE_INLINE IVec operator<= (const FVec& a, const FVec& b) { return _mm256_castps_si256( _mm256_cmp_ps( a, b, _CMP_LE_OS) ); } FORCE_INLINE IVec operator>= (const FVec& a, const FVec& b) { return _mm256_castps_si256( _mm256_cmp_ps( a, b, _CMP_GE_OS ) ); } FORCE_INLINE IVec operator< (const FVec& a, const FVec& b) { return _mm256_castps_si256(_mm256_cmp_ps(a, b, _CMP_LT_OS )); } #ifdef USE_FMA FORCE_INLINE FVec mulSub(const FVec& a, const FVec& b, const FVec& c) { return _mm256_fmsub_ps(a, b, c); } #endif typedef VecStorage> FVec256Double; template <> struct FVec : FVec256Double { FORCE_INLINE FVec() {} FORCE_INLINE FVec( double d ) : FVec256Double( _mm256_set1_pd( d ) ) {} FORCE_INLINE FVec( const double *v ) : FVec256Double( _mm256_loadu_pd( v ) ) {} FORCE_INLINE FVec( const vec_t& v) : FVec256Double( v ) {} FORCE_INLINE FVec(double d3, double d2, double d1, double d0) : FVec256Double(_mm256_set_pd(d3, d2, d1, d0)) {} //void set0( double f ) { vec = _mm256_load_sd( &f ); } void setN( double f ) { vec = _mm256_set1_pd( f ); } FORCE_INLINE void setidx( const double *xi, const IVec& idx ) { vec = _mm256_i32gather_pd(xi, idx, 8); } FORCE_INLINE void setidx( const double *xi, const IVec& idx ) { vec = _mm256_i64gather_pd(xi, idx, 8); } // FORCE_INLINE double get0() const { return _mm256_cvtsd_f64( vec ); } // FORCE_INLINE double get1() const { return _mm256_cvtsd_f64( _mm256_shuffle_pd( vec, vec, 1 ) ); }; }; FORCE_INLINE FVec operator- (const FVec& a, const FVec& b) { return _mm256_sub_pd( a, b ); } FORCE_INLINE FVec operator* (const FVec& a, const FVec& b) { return _mm256_mul_pd( a, b ); } FORCE_INLINE FVec operator/ (const FVec& a, const FVec& b) { return _mm256_div_pd( a, b ); } FORCE_INLINE IVec ftoi (const FVec& a) { return _mm256_cvttpd_epi32(a); } FORCE_INLINE IVec operator<= (const FVec& a, const FVec& b) { return _mm256_castpd_si256(_mm256_cmp_pd( a, b, _CMP_LE_OS ) ); } FORCE_INLINE IVec operator< (const FVec& a, const FVec& b) { return _mm256_castpd_si256(_mm256_cmp_pd(a, b, _CMP_LT_OS)); } FORCE_INLINE IVec operator>= (const FVec& a, const FVec& b) { return _mm256_castpd_si256(_mm256_cmp_pd( a, b, _CMP_GE_OS ) ); } #ifdef USE_FMA FORCE_INLINE FVec mulSub(const FVec& a, const FVec& b, const FVec& c) { return _mm256_fmsub_pd(a, b, c); } #endif #endif } // namespace Details } // namespace BinSearch #endif // !defined(__aarch64__) bitsandbytes/include/Type.h000066400000000000000000000131151506752665600163110ustar00rootroot00000000000000 #pragma once #include #include #include #include "Portable.h" using std::size_t; namespace BinSearch { enum InstrSet { Scalar, SSE, AVX, Neon, VSX }; #define ALGOENUM(x, b) x, enum Algos { #include "AlgoXCodes.h" }; #undef ALGOENUM namespace Details { template struct InstrIntTraits; template struct InstrFloatTraits; // base class for algorithm supporting the method: // uint32 scalar(T z) const template struct AlgoScalarBase; // base class for algorithm supporting the following methods, constants and definitions: // static const uint32 nElem // struct Constants; // void initConstants(Constants& cst) const // void vectorial(uint32 *pr, const T *pz, const Constants& cst) const // The function vectorial processes nElem items template struct AlgoVecBase; template struct IntTraits; template <> struct IntTraits { typedef uint32 itype; }; template <> struct IntTraits { typedef uint64 itype; }; template struct Body { template FORCE_INLINE static void iteration(const Expr& e, uint32 *ri, const T* zi, const typename Expr::Constants& cst) { e.vectorial(ri, zi, cst); Body::template iteration(e, ri + D, zi + D, cst); } }; template <> struct Body<0> { template FORCE_INLINE static void iteration(const Expr& e, uint32 *ri, const T* zi, const H&) { } }; template struct Loop { typedef Algo algo_type; static const uint32 M = 4; static const uint32 D = algo_type::nElem; FORCE_INLINE static void loop(const algo_type& e, uint32 *ri, const T* zi, uint32 n) { typename algo_type::Constants cst; e.initConstants(cst); uint32 j = 0; while (j + (D*M) <= n) { Details::Body::template iteration(e, ri + j, zi + j, cst); j += (D*M); } while (j + D <= n) { e.vectorial(ri + j, zi + j, cst); j += D; } while (D > 1 && j < n) { ri[j] = e.scalar(zi[j]); j += 1; } } }; template struct _Pipeliner { template FORCE_INLINE static void go(const Expr& e, Data* d) { e.template run(d); _Pipeliner::go(e, d); } }; template struct _Pipeliner { template FORCE_INLINE static void go(const Expr& e, Data* d) { } }; template struct Pipeliner { template FORCE_INLINE static void go(const Expr& e, Data* d) { _Pipeliner::go(e, d); } }; #if 1 template char is_complete_impl(char (*)[sizeof(T)]); template long is_complete_impl(...); template struct IsComplete { static const bool value = sizeof(is_complete_impl(0)) == sizeof(char); }; #else template std::true_type is_complete_impl(T *); std::false_type is_complete_impl(...); template struct IsComplete : decltype(is_complete_impl(std::declval())) {}; #endif template struct AlgoScalarToVec : AlgoScalarBase { typedef AlgoScalarBase base_t; AlgoScalarToVec(const typename base_t::Data& d) : base_t(d) {} AlgoScalarToVec(const T* px, const uint32 n) : base_t(px, n) {} static const uint32 nElem = 1; struct Constants { }; void initConstants(Constants& cst) const { } FORCE_INLINE void vectorial(uint32 *pr, const T *pz, const Constants& cst) const { *pr = base_t::scalar(*pz); } }; template struct conditional { typedef T type; }; template struct conditional { typedef F type; }; template struct CondData { FORCE_INLINE CondData(T x) : v(x) {} FORCE_INLINE operator const T&() const { return v;} private: T v; }; template struct CondData { FORCE_INLINE CondData(T) {} FORCE_INLINE operator const T() const { return 0;} }; template struct BinAlgoBase : Details::conditional< Details::IsComplete>::value , Details::AlgoVecBase , Details::AlgoScalarToVec >::type { typedef typename Details::conditional< Details::IsComplete>::value , Details::AlgoVecBase , Details::AlgoScalarToVec >::type base_t; BinAlgoBase(const T* px, const uint32 n) : base_t(px, n) {} BinAlgoBase(const typename base_t::Data& d) : base_t(d) {} }; } // namespace Details } // namespace BinSearch bitsandbytes/install_cuda.py000066400000000000000000000071701506752665600166140ustar00rootroot00000000000000import os import subprocess import sys from urllib.request import urlretrieve cuda_versions = { "118": "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run", "120": "https://developer.download.nvidia.com/compute/cuda/12.0.1/local_installers/cuda_12.0.1_525.85.12_linux.run", "121": "https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda_12.1.1_530.30.02_linux.run", "122": "https://developer.download.nvidia.com/compute/cuda/12.2.2/local_installers/cuda_12.2.2_535.104.05_linux.run", "123": "https://developer.download.nvidia.com/compute/cuda/12.3.2/local_installers/cuda_12.3.2_545.23.08_linux.run", "124": "https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda_12.4.1_550.54.15_linux.run", "125": "https://developer.download.nvidia.com/compute/cuda/12.5.1/local_installers/cuda_12.5.1_555.42.06_linux.run", "126": "https://developer.download.nvidia.com/compute/cuda/12.6.2/local_installers/cuda_12.6.2_560.35.03_linux.run", } def install_cuda(version, base_path, download_path): formatted_version = f"{version[:-1]}.{version[-1]}" folder = f"cuda-{formatted_version}" install_path = os.path.join(base_path, folder) if os.path.exists(install_path): print(f"Removing existing CUDA version {version} at {install_path}...") subprocess.run(["rm", "-rf", install_path], check=True) url = cuda_versions[version] filename = url.split("/")[-1] filepath = os.path.join(download_path, filename) if not os.path.exists(filepath): print(f"Downloading CUDA version {version} from {url}...") urlretrieve(url, filepath) else: print(f"Installer for CUDA version {version} already downloaded.") # Make the installer executable subprocess.run(["chmod", "+x", filepath], check=True) # Install CUDA print(f"Installing CUDA version {version}...") install_command = [ "bash", filepath, "--no-drm", "--no-man-page", "--override", "--toolkitpath=" + install_path, "--toolkit", "--silent", ] print(f"Running command: {' '.join(install_command)}") try: subprocess.run(install_command, check=True) except subprocess.CalledProcessError as e: print(f"Installation failed for CUDA version {version}: {e}") return finally: # Delete the installer file os.remove(filepath) print(f"CUDA version {version} installed at {install_path}") def main(): user_base_path = os.path.expanduser("~/cuda") system_base_path = "/usr/local/cuda" base_path = user_base_path # default to user-specific installation download_path = "/tmp" # default download path if len(sys.argv) < 2: print("Usage: python install_cuda.py [user/system] [download_path]") sys.exit(1) version = sys.argv[1] if len(sys.argv) > 2: base_path = system_base_path if sys.argv[2] == "system" else user_base_path if len(sys.argv) > 3: download_path = sys.argv[3] if not os.path.exists(base_path): os.makedirs(base_path) if not os.path.exists(download_path): os.makedirs(download_path) # Install CUDA version(s) if version == "all": for ver in cuda_versions: install_cuda(ver, base_path, download_path) elif version in cuda_versions: install_cuda(version, base_path, download_path) else: print(f"Invalid CUDA version: {version}. Available versions are: {', '.join(cuda_versions.keys())}") sys.exit(1) if __name__ == "__main__": main() bitsandbytes/install_cuda.sh000066400000000000000000000043621506752665600165760ustar00rootroot00000000000000URL118=https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run URL120=https://developer.download.nvidia.com/compute/cuda/12.0.1/local_installers/cuda_12.0.1_525.85.12_linux.run URL121=https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda_12.1.1_530.30.02_linux.run URL122=https://developer.download.nvidia.com/compute/cuda/12.2.2/local_installers/cuda_12.2.2_535.104.05_linux.run URL123=https://developer.download.nvidia.com/compute/cuda/12.3.2/local_installers/cuda_12.3.2_545.23.08_linux.run URL124=https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda_12.4.1_550.54.15_linux.run URL125=https://developer.download.nvidia.com/compute/cuda/12.5.1/local_installers/cuda_12.5.1_555.42.06_linux.run URL126=https://developer.download.nvidia.com/compute/cuda/12.6.2/local_installers/cuda_12.6.2_560.35.03_linux.run CUDA_VERSION=$1 BASE_PATH=$2 EXPORT_BASHRC=$3 if [[ -n "$CUDA_VERSION" ]]; then if [[ "$CUDA_VERSION" -eq "118" ]]; then URL=$URL118 FOLDER=cuda-11.8 elif [[ "$CUDA_VERSION" -eq "120" ]]; then URL=$URL120 FOLDER=cuda-12.0 elif [[ "$CUDA_VERSION" -eq "121" ]]; then URL=$URL121 FOLDER=cuda-12.1 elif [[ "$CUDA_VERSION" -eq "122" ]]; then URL=$URL122 FOLDER=cuda-12.2 elif [[ "$CUDA_VERSION" -eq "123" ]]; then URL=$URL123 FOLDER=cuda-12.3 elif [[ "$CUDA_VERSION" -eq "124" ]]; then URL=$URL124 FOLDER=cuda-12.4 elif [[ "$CUDA_VERSION" -eq "125" ]]; then URL=$URL125 FOLDER=cuda-12.5 elif [[ "$CUDA_VERSION" -eq "126" ]]; then URL=$URL126 FOLDER=cuda-12.6 else echo "argument error: No cuda version passed as input. Choose among versions 118 to 126" fi else echo "argument error: No cuda version passed as input. Choose among versions 118 to 126" fi FILE=$(basename $URL) if [[ -n "$CUDA_VERSION" ]]; then echo $URL echo $FILE wget $URL bash $FILE --no-drm --no-man-page --override --toolkitpath=$BASE_PATH/$FOLDER/ --toolkit --silent if [ "$EXPORT_BASHRC" -eq "1" ]; then echo "export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:$BASE_PATH/$FOLDER/lib64" >> ~/.bashrc echo "export PATH=\$PATH:$BASE_PATH/$FOLDER/bin" >> ~/.bashrc source ~/.bashrc fi else echo "" fi bitsandbytes/pyproject.toml000066400000000000000000000111421506752665600165060ustar00rootroot00000000000000[build-system] requires = ["scikit-build-core", "setuptools >= 77.0.3", "trove-classifiers>=2025.8.6.13"] build-backend = "scikit_build_core.setuptools.build_meta" [project] name = "bitsandbytes" dynamic = ["version"] description = "k-bit optimizers and matrix multiplication routines." authors = [{name="Tim Dettmers", email="dettmers@cs.washington.edu"}] maintainers = [ {name="Titus von Köller", email="titus@huggingface.co"}, {name="Matthew Douglas", email="matthew.douglas@huggingface.co"} ] requires-python = ">=3.9" readme = "README.md" license = "MIT" license-files = ["LICENSE"] keywords = [ "gpu", "optimizers", "optimization", "8-bit", "quantization", "compression" ] classifiers = [ "Development Status :: 4 - Beta", "Environment :: GPU :: NVIDIA CUDA :: 11.8", "Environment :: GPU :: NVIDIA CUDA :: 12", "Environment :: GPU :: NVIDIA CUDA :: 13", "Intended Audience :: Developers", "Intended Audience :: Science/Research", "Operating System :: POSIX :: Linux", # "Operating System :: MacOS", "Operating System :: Microsoft :: Windows", "Programming Language :: C++", "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Topic :: Scientific/Engineering :: Artificial Intelligence" ] dependencies = [ "torch>=2.3,<3", "numpy>=1.17", "packaging>=20.9" ] [project.urls] homepage = "https://github.com/bitsandbytes-foundation/bitsandbytes" changelog = "https://github.com/bitsandbytes-foundation/bitsandbytes/blob/main/CHANGELOG.md" docs = "https://huggingface.co/docs/bitsandbytes/main" issues = "https://github.com/bitsandbytes-foundation/bitsandbytes/issues" [project.optional-dependencies] benchmark = ["pandas", "matplotlib"] docs = ["hf-doc-builder==0.5.0"] dev = [ "bitsandbytes[test]", "build>=1.0.0,<2", "ruff==0.11.2", "pre-commit>=3.5.0,<4", "wheel>=0.42,<1" ] test = [ "einops~=0.8.0", "lion-pytorch==0.2.3", "pytest~=8.3", "scipy>=1.11.4,<2", "transformers>=4.30.1,<5" ] [tool.setuptools] package-data = { "*" = ["libbitsandbytes*.*", "py.typed"] } [tool.setuptools.packages.find] include = ["bitsandbytes*"] [tool.setuptools.dynamic] version = {attr = "bitsandbytes.__version__"} [tool.coverage.report] exclude_also = [ # exclude backward() functions from coverage, as they are invoked from C++ 'def backward\(ctx' ] [tool.pytest.ini_options] addopts = "-rP -m 'not slow and not benchmark and not deprecated'" # ; --cov=bitsandbytes # ; # contexts: record which test ran which line; can be seen in html coverage report # ; --cov-context=test # ; --cov-report html log_cli = true log_cli_level = "INFO" log_file = "logs/pytest.log" markers = [ "benchmark: mark test as a benchmark", "deprecated: mark test as covering a deprecated feature", "slow: mark test as slow", ] [tool.ruff] src = [ "bitsandbytes", "tests", "benchmarking" ] target-version = "py39" line-length = 119 [tool.ruff.lint] select = [ "B", # bugbear: security warnings "E", # pycodestyle (error) "W", # pycodestyle (warning) "F", # pyflakes "I", # isort "ISC", # implicit string concatenation "UP", # alert you when better syntax is available in your python version "RUF", # the ruff developer's own rules ] ignore = [ "B007", # Loop control variable not used within the loop body (TODO: enable) "B028", # Warning without stacklevel (TODO: enable) "E501", # Suppress line-too-long warnings: trust yapf's judgement on this one. "E701", # Multiple statements on one line (TODO: enable) "E712", # Allow using if x == False, as it's not always equivalent to if x. "E731", # Do not use lambda "RUF012", # Mutable class attribute annotations "RUF034", # Useless if-else (TODO: enable) "ISC001", # single-line-implicit-string-concatenation incompatible with formatter ] [tool.ruff.lint.extend-per-file-ignores] "**/__init__.py" = ["F401"] # allow unused imports in __init__.py "{benchmarking,tests}/**/*.py" = [ "B007", "B011", "B023", "E701", "E731", "F841", "UP030", ] [tool.ruff.lint.isort] combine-as-imports = true detect-same-package = true force-sort-within-sections = true known-first-party = ["bitsandbytes"] [[tool.mypy.overrides]] module = "triton.*" ignore_missing_imports = true [[tool.mypy.overrides]] module = "scipy.stats" ignore_missing_imports = true bitsandbytes/scripts/000077500000000000000000000000001506752665600152625ustar00rootroot00000000000000bitsandbytes/scripts/stale.py000066400000000000000000000043251506752665600167500ustar00rootroot00000000000000# Copyright 2023 The HuggingFace Team, the AllenNLP library authors. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ Script to close stale issue. Taken in part from the AllenNLP repository. https://github.com/allenai/allennlp. """ from datetime import datetime as dt, timezone import os from github import Github # All labels that we don't want to touch LABELS_TO_EXEMPT = [ "feature-request", ] def main(): g = Github(os.environ["GITHUB_TOKEN"]) repo = g.get_repo("TimDettmers/bitsandbytes") open_issues = repo.get_issues(state="open") for issue in open_issues: comments = sorted([comment for comment in issue.get_comments()], key=lambda i: i.created_at, reverse=True) last_comment = comments[0] if len(comments) > 0 else None if ( last_comment is not None and last_comment.user.login == "github-actions[bot]" and (dt.now(timezone.utc) - issue.updated_at).days > 7 and (dt.now(timezone.utc) - issue.created_at).days >= 30 and not any(label.name.lower() in LABELS_TO_EXEMPT for label in issue.get_labels()) ): issue.edit(state="closed") elif ( (dt.now(timezone.utc) - issue.updated_at).days > 23 and (dt.now(timezone.utc) - issue.created_at).days >= 30 and not any(label.name.lower() in LABELS_TO_EXEMPT for label in issue.get_labels()) ): issue.create_comment( "This issue has been automatically marked as stale because it has not had " "recent activity. If you think this still needs to be addressed " "please comment on this thread.\n\n", ) if __name__ == "__main__": main() bitsandbytes/setup.py000066400000000000000000000022751506752665600153130ustar00rootroot00000000000000# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from distutils.errors import DistutilsModuleError from warnings import warn from setuptools import find_packages, setup from setuptools.command.build_py import build_py from setuptools.dist import Distribution # Tested with wheel v0.29.0 class BinaryDistribution(Distribution): def has_ext_modules(self): return True class ExtBuildPy(build_py): def run(self): # build_cmake needs to be called prior to build_py, as the latter # collects the files output into the package directory. try: self.run_command("build_cmake") except DistutilsModuleError: warn( "scikit-build-core not installed, CMake will not be invoked automatically. " "Please install scikit-build-core or run CMake manually to build extensions." ) super().run() setup( version="0.48.1.dev0", packages=find_packages(), distclass=BinaryDistribution, cmake_source_dir=".", cmdclass={ "build_py": ExtBuildPy, }, ) bitsandbytes/tests/000077500000000000000000000000001506752665600147355ustar00rootroot00000000000000bitsandbytes/tests/__init__.py000066400000000000000000000000001506752665600170340ustar00rootroot00000000000000bitsandbytes/tests/conftest.py000066400000000000000000000022341506752665600171350ustar00rootroot00000000000000import gc import random import numpy as np import pytest import torch def _set_seed(): torch.manual_seed(0) torch.cuda.manual_seed_all(0) torch.mps.manual_seed(0) np.random.seed(0) random.seed(0) def pytest_runtest_call(item): try: _set_seed() item.runtest() except AssertionError as ae: if str(ae) == "Torch not compiled with CUDA enabled": pytest.skip("Torch not compiled with CUDA enabled") raise except RuntimeError as re: # CUDA-enabled Torch build, but no CUDA-capable device found if "Found no NVIDIA driver on your system" in str(re): pytest.skip("No NVIDIA driver found") raise @pytest.hookimpl(trylast=True) def pytest_runtest_teardown(item, nextitem): gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() elif torch.backends.mps.is_available() and torch.backends.mps.is_built(): torch.mps.empty_cache() @pytest.fixture(scope="session") def requires_cuda() -> bool: cuda_available = torch.cuda.is_available() if not cuda_available: pytest.skip("CUDA is required") return cuda_available bitsandbytes/tests/helpers.py000066400000000000000000000064671506752665600167660ustar00rootroot00000000000000import functools from io import BytesIO from itertools import product import os import random from typing import Any import torch from bitsandbytes.cextension import HIP_ENVIRONMENT test_dims_rng = random.Random(42) TRUE_FALSE = (True, False) BOOLEAN_TRIPLES = list(product(TRUE_FALSE, repeat=3)) # all combinations of (bool, bool, bool) BOOLEAN_TUPLES = list(product(TRUE_FALSE, repeat=2)) # all combinations of (bool, bool) @functools.cache def get_available_devices(no_cpu=False): if "BNB_TEST_DEVICE" in os.environ: # If the environment variable is set, use it directly. device = os.environ["BNB_TEST_DEVICE"] return [] if no_cpu and device == "cpu" else [device] devices = [] if HIP_ENVIRONMENT else ["cpu"] if not no_cpu else [] if hasattr(torch, "accelerator"): # PyTorch 2.6+ - determine accelerator using agnostic API. if torch.accelerator.is_available(): devices += [str(torch.accelerator.current_accelerator())] else: if torch.cuda.is_available(): devices += ["cuda"] if torch.backends.mps.is_available(): devices += ["mps"] if hasattr(torch, "xpu") and torch.xpu.is_available(): devices += ["xpu"] custom_backend_name = torch._C._get_privateuse1_backend_name() custom_backend_module = getattr(torch, custom_backend_name, None) custom_backend_is_available_fn = getattr(custom_backend_module, "is_available", None) if custom_backend_is_available_fn and custom_backend_module.is_available(): devices += [custom_backend_name] return devices def torch_save_to_buffer(obj): buffer = BytesIO() torch.save(obj, buffer) buffer.seek(0) return buffer def torch_load_from_buffer(buffer): buffer.seek(0) obj = torch.load(buffer, weights_only=False) buffer.seek(0) return obj def get_test_dims(min: int, max: int, *, n: int) -> list[int]: return [test_dims_rng.randint(min, max) for _ in range(n)] def format_with_label(label: str, value: Any) -> str: if isinstance(value, bool): formatted = "T" if value else "F" elif isinstance(value, (list, tuple)) and all(isinstance(v, bool) for v in value): formatted = "".join("T" if b else "F" for b in value) elif isinstance(value, torch.dtype): formatted = describe_dtype(value) else: formatted = str(value) return f"{label}={formatted}" def id_formatter(label: str): """ Return a function that formats the value given to it with the given label. """ return lambda value: format_with_label(label, value) DTYPE_NAMES = { torch.bfloat16: "bf16", torch.bool: "bool", torch.float16: "fp16", torch.float32: "fp32", torch.float64: "fp64", torch.int32: "int32", torch.int64: "int64", torch.int8: "int8", } def describe_dtype(dtype: torch.dtype) -> str: return DTYPE_NAMES.get(dtype) or str(dtype).rpartition(".")[2] def is_supported_on_hpu( quant_type: str = "nf4", dtype: torch.dtype = torch.bfloat16, quant_storage: torch.dtype = torch.uint8 ) -> bool: """ Check if the given quant_type, dtype and quant_storage are supported on HPU. """ if quant_type == "fp4" or dtype == torch.float16 or quant_storage not in (torch.uint8, torch.bfloat16): return False return True bitsandbytes/tests/test_autograd.py000066400000000000000000000251051506752665600201570ustar00rootroot00000000000000import pytest import torch import bitsandbytes as bnb from tests.helpers import ( BOOLEAN_TRIPLES, TRUE_FALSE, describe_dtype, get_available_devices, id_formatter, is_supported_on_hpu, ) TRANSPOSE_VALS = [(False, True), (False, False)] @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dim1", [40], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [64, 0], ids=id_formatter("dim2")) @pytest.mark.parametrize("dim3", [32], ids=id_formatter("dim3")) @pytest.mark.parametrize("dim4", [48], ids=id_formatter("dim4")) @pytest.mark.parametrize("decomp", [0.0, 6.0], ids=id_formatter("decomp")) @pytest.mark.parametrize( "funcs", [(torch.matmul, bnb.matmul), (torch.matmul, bnb.research.switchback_bnb)], ids=["func=matmul", "func=switchback_bnb"], ) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype) @pytest.mark.parametrize("req_grad", BOOLEAN_TRIPLES, ids=id_formatter("req_grad")) @pytest.mark.parametrize("transpose", TRANSPOSE_VALS, ids=id_formatter("transpose")) @pytest.mark.parametrize("has_fp16_weights", TRUE_FALSE, ids=id_formatter("has_fp16_weights")) @pytest.mark.parametrize("has_bias", TRUE_FALSE, ids=id_formatter("has_bias")) def test_matmullt( device, dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose, decomp, has_fp16_weights, has_bias ): if device != "cuda": if funcs[1] == bnb.research.switchback_bnb: # TODO: Deprecate/remove? pytest.skip("switchback_bnb only works on CUDA.") if req_grad[1]: # This will be deprecated for CUDA in the future. We don't expect # this to work on any other device. pytest.skip("Deprecated feature with CUDA support only.") dimA = (dim2, dim3) if not transpose[0] else (dim3, dim2) dimB = (dim3, dim4) if not transpose[1] else (dim4, dim3) outlier_dim = torch.randint(0, dimA[1], size=(dimA[1] // 8,), device=device) if has_bias == False: req_grad = list(req_grad) req_grad[2] = False if device == "cpu" and dtype != torch.float32 and has_fp16_weights and any(req_grad): if torch.__version__ < (2, 6): pytest.xfail("mse_loss bf16/fp16 on CPU is not supported in torch < 2.6") for i in range(3): # normal multiply if funcs[0] in [torch.mm, torch.matmul]: A = torch.randn(size=dimA, device=device, requires_grad=req_grad[0], dtype=dtype) if decomp == 6.0: with torch.no_grad(): A[:, outlier_dim] = 6.0 B = torch.randn(size=dimB, device=device, requires_grad=req_grad[1], dtype=dtype) target = torch.randn( size=(dim2, dim4), device=device, requires_grad=req_grad[1], dtype=dtype, ) bias = None bias2 = None if has_bias: bias = torch.randn(dim4, device=device, dtype=dtype, requires_grad=req_grad[2]) bias2 = bias.clone() torch.nn.init.xavier_uniform_(B) B2 = B.clone() state = bnb.MatmulLtState() state.threshold = decomp state.has_fp16_weights = has_fp16_weights if not has_fp16_weights: if not transpose[0] and not transpose[1]: B2 = B2.t().contiguous() state.CB, state.SCB, _ = bnb.functional.int8_vectorwise_quant(B2.to(torch.float16)) B2 = state.CB if not transpose[0] and transpose[1]: out_torch = funcs[0](A, B.t()) out_bnb = funcs[1](A, B2, state=state, bias=bias2) elif not transpose[0] and not transpose[1]: out_torch = funcs[0](A, B) out_bnb = funcs[1](A, B2.t(), state=state, bias=bias2) if has_bias: out_torch += bias assert out_bnb.dtype == A.dtype, f"bnb matmullt received {A.dtype} but returned {out_bnb.dtype}" n = out_bnb.numel() err = torch.abs(out_bnb - out_torch).mean().item() # print(f'abs error {err:.4f}') idx = torch.isclose(out_bnb, out_torch, atol=0.01, rtol=0.1) assert (idx == 0).sum().item() <= n * (0.0175 if dtype == torch.float16 else 0.021) idx = torch.isclose(out_bnb, out_torch, atol=0.035, rtol=0.2) assert (idx == 0).sum().item() <= n * 0.001 if has_fp16_weights: if any(req_grad): out_bnb.data.copy_(out_torch) if device == "cuda": torch.cuda.synchronize() loss_bnb = torch.nn.functional.mse_loss(out_bnb, target).mean() loss_bnb.backward() gradA1 = A.grad gradB1 = B.grad A.grad = None B.grad = None if has_bias: gradBias1 = bias.grad bias.grad = None loss_torch = torch.nn.functional.mse_loss(out_torch, target).mean() loss_torch.backward() gradA2 = A.grad gradB2 = B.grad A.grad = None B.grad = None if has_bias: gradBias2 = bias.grad bias.grad = None if req_grad[0]: torch.testing.assert_close(gradA1, gradA2, atol=0.015, rtol=0.1) if req_grad[1]: n = gradB1.numel() if dim2 > 0: assert torch.abs(gradB1).sum() > 0.0 assert torch.abs(gradB2).sum() > 0.0 else: assert torch.abs(gradB1).sum() == 0.0 assert torch.abs(gradB2).sum() == 0.0 idx = torch.isclose(gradB1, gradB2, atol=0.06, rtol=0.3) assert (idx == 0).sum().item() <= n * 0.10 idx = torch.isclose(gradB1, gradB2, atol=0.10, rtol=0.3) assert (idx == 0).sum().item() <= n * 0.02 torch.testing.assert_close(gradB1, gradB2, atol=0.18, rtol=0.3) if req_grad[2]: torch.testing.assert_close(gradBias1, gradBias2) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dim1", [48], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [64, 0], ids=id_formatter("dim2")) @pytest.mark.parametrize("dim3", [64], ids=id_formatter("dim3")) @pytest.mark.parametrize("dim4", [96], ids=id_formatter("dim4")) @pytest.mark.parametrize("funcs", [(torch.matmul, bnb.matmul_4bit)], ids=["func=matmul"]) @pytest.mark.parametrize("req_grad", BOOLEAN_TRIPLES, ids=id_formatter("req_grad")) @pytest.mark.parametrize("transpose", TRANSPOSE_VALS, ids=id_formatter("transpose")) @pytest.mark.parametrize("has_bias", TRUE_FALSE, ids=id_formatter("has_bias")) @pytest.mark.parametrize("dtype", [torch.float16, torch.float32], ids=describe_dtype) @pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics")) @pytest.mark.parametrize("quant_type", ["fp4", "nf4"], ids=id_formatter("quant_type")) def test_matmul_4bit( device, dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose, has_bias, compress_statistics, quant_type, ): dimA = (dim2, dim3) if not transpose[0] else (dim3, dim2) dimB = (dim3, dim4) if not transpose[1] else (dim4, dim3) if has_bias == False: req_grad = list(req_grad) req_grad[2] = False if device == "cpu" and dtype != torch.float32 and any(req_grad) and torch.__version__ < (2, 6): pytest.xfail("mse_loss fp16 on CPU is not supported in torch < 2.6") if device == "hpu" and not is_supported_on_hpu(quant_type, dtype): pytest.skip("This configuration is not supported on HPU.") for i in range(3): # normal multiply if funcs[0] in [torch.mm, torch.matmul]: A = torch.randn(size=dimA, device=device, requires_grad=req_grad[0], dtype=dtype) B = torch.randn(size=dimB, device=device, requires_grad=req_grad[1], dtype=dtype) target = torch.randn(size=(dim2, dim4), device=device, requires_grad=req_grad[1], dtype=dtype) bias = None bias2 = None if has_bias: bias = torch.randn(dim4, device=device, dtype=dtype, requires_grad=req_grad[2]) bias2 = bias.clone() torch.nn.init.xavier_uniform_(B) B2, quant_state = bnb.functional.quantize_4bit( B, compress_statistics=compress_statistics, quant_type=quant_type, ) if not transpose[0] and transpose[1]: out_torch = funcs[0](A, B.t()) out_bnb = funcs[1](A, B2.t(), quant_state, bias=bias2) elif not transpose[0] and not transpose[1]: out_torch = funcs[0](A, B) out_bnb = funcs[1](A, B2, quant_state, bias=bias2) if has_bias: out_torch += bias assert out_bnb.dtype == A.dtype, f"bnb matmullt received {A.dtype} but returned {out_bnb.dtype}" n = out_bnb.numel() err = torch.abs(out_bnb - out_torch).float().mean().item() if n > 0: assert err < 0.115 # assert err < 0.20 if any(req_grad): out_bnb.data.copy_(out_torch) if device == "cuda": torch.cuda.synchronize() elif device == "hpu": torch.hpu.synchronize() loss_bnb = torch.nn.functional.mse_loss(out_bnb, target).mean() loss_bnb.backward() gradA1 = A.grad gradB1 = B.grad A.grad = None B.grad = None if has_bias: gradBias1 = bias.grad bias.grad = None loss_torch = torch.nn.functional.mse_loss(out_torch, target).mean() loss_torch.backward() gradA2 = A.grad gradB2 = B.grad A.grad = None B.grad = None if has_bias: gradBias2 = bias.grad bias.grad = None if req_grad[0]: torch.testing.assert_close(gradA1, gradA2, atol=0.015, rtol=0.1) if req_grad[2]: torch.testing.assert_close(gradBias1, gradBias2) bitsandbytes/tests/test_cuda_setup_evaluator.py000066400000000000000000000017421506752665600225700ustar00rootroot00000000000000import pytest from bitsandbytes.cextension import HIP_ENVIRONMENT, get_cuda_bnb_library_path from bitsandbytes.cuda_specs import CUDASpecs @pytest.fixture def cuda120_spec() -> CUDASpecs: return CUDASpecs( cuda_version_string="120", highest_compute_capability=(8, 6), cuda_version_tuple=(12, 0), ) @pytest.mark.skipif(HIP_ENVIRONMENT, reason="this test is not supported on ROCm") def test_get_cuda_bnb_library_path(monkeypatch, cuda120_spec): monkeypatch.delenv("BNB_CUDA_VERSION", raising=False) assert get_cuda_bnb_library_path(cuda120_spec).stem == "libbitsandbytes_cuda120" @pytest.mark.skipif(HIP_ENVIRONMENT, reason="this test is not supported on ROCm") def test_get_cuda_bnb_library_path_override(monkeypatch, cuda120_spec, caplog): monkeypatch.setenv("BNB_CUDA_VERSION", "110") assert get_cuda_bnb_library_path(cuda120_spec).stem == "libbitsandbytes_cuda110" assert "BNB_CUDA_VERSION" in caplog.text # did we get the warning? bitsandbytes/tests/test_deprecated.py000066400000000000000000000152471506752665600204570ustar00rootroot00000000000000import pytest import torch import bitsandbytes as bnb from bitsandbytes import functional as F from tests.helpers import BOOLEAN_TRIPLES, describe_dtype, get_test_dims, id_formatter from tests.test_autograd import TRANSPOSE_VALS @pytest.mark.deprecated def test_dynamic_quantization(): diffs = [] reldiffs = [] for i in range(100): A1 = torch.randn(1024, 1024, device="cuda") C, S = F.quantize(A1) A2 = F.dequantize(C, S) diff = torch.abs(A1 - A2) reldiff = diff / torch.abs(A1 + 1e-8) diffs.append(diff.mean().item()) reldiffs.append(reldiff.mean().item()) assert diff.mean().item() < 0.0135 print(sum(diffs) / len(diffs)) print(sum(reldiffs) / len(reldiffs)) for i in range(100): A1 = torch.rand(1024, 1024, device="cuda") C, S = F.quantize(A1) A2 = F.dequantize(C, S) diff = torch.abs(A1 - A2).mean().item() torch.testing.assert_close(A1, A2, atol=1e-2, rtol=0) assert diff < 0.004 @pytest.mark.parametrize("gtype", [torch.float32, torch.float16], ids=["float", "half"]) @pytest.mark.deprecated def test_percentile_clipping(gtype): gnorm_vec1 = torch.zeros(100, device="cuda") gnorm_vec2 = torch.zeros(100, device="cuda") n = 4 step = 0 percentile = 5 for i in range(20): step += 1 g = torch.randn(n, n, dtype=gtype, device="cuda") gnorm1, clip2, gnorm_scale = F.percentile_clipping(g, gnorm_vec2, step, percentile=percentile) assert gnorm_scale == 1.0 if gnorm1 < clip2 else clip2 / gnorm1 gnorm2 = torch.norm(g.float()) if step == 1: gnorm_vec1[:] = gnorm2 else: gnorm_vec1[step % 100] = gnorm2 vals, idx = torch.sort(gnorm_vec1) clip1 = vals[percentile] torch.testing.assert_close(gnorm_vec1, torch.sqrt(gnorm_vec2)) torch.testing.assert_close(clip1, clip2) torch.testing.assert_close(gnorm1, gnorm2) @pytest.mark.parametrize("dim1", get_test_dims(16, 64, n=1), ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [*get_test_dims(32, 96, n=1), 0], ids=id_formatter("dim2")) @pytest.mark.parametrize("dim3", get_test_dims(32, 96, n=1), ids=id_formatter("dim3")) @pytest.mark.parametrize("dim4", get_test_dims(32, 96, n=1), ids=id_formatter("dim4")) @pytest.mark.parametrize("req_grad", BOOLEAN_TRIPLES, ids=id_formatter("req_grad")) @pytest.mark.parametrize("transpose", TRANSPOSE_VALS, ids=id_formatter("transpose")) @pytest.mark.parametrize("dtype", [torch.float16, torch.float32], ids=describe_dtype) @pytest.mark.parametrize( "funcs", [(torch.matmul, bnb.research.matmul_fp8_mixed), (torch.matmul, bnb.research.matmul_fp8_global)], ids=["matmul_fp8_mixed", "matmul_fp8_global"], ) @pytest.mark.deprecated @pytest.mark.skip("Deprecated functionality, to be removed.") def test_matmul_fp8(dim1, dim2, dim3, dim4, funcs, dtype, req_grad, transpose): dimA = (dim2, dim3) if not transpose[0] else (dim3, dim2) dimB = (dim3, dim4) if not transpose[1] else (dim4, dim3) req_grad = list(req_grad) req_grad[2] = False for i in range(3): # normal multiply if funcs[0] in [torch.mm, torch.matmul]: A = torch.randn(size=dimA, device="cuda", requires_grad=req_grad[0], dtype=dtype) B = torch.randn(size=dimB, device="cuda", requires_grad=req_grad[1], dtype=dtype) target = torch.randn(size=(dim2, dim4), device="cuda", requires_grad=req_grad[1], dtype=dtype) torch.nn.init.xavier_uniform_(B) fw_code = bnb.functional.create_fp8_map(True, 4, 3, 8).to(A.device) bw_code = bnb.functional.create_fp8_map(True, 5, 2, 8).to(A.device) if not transpose[0] and transpose[1]: out_torch = funcs[0](A, B.t()) out_bnb = funcs[1](A, B.t(), fw_code, bw_code) elif not transpose[0] and not transpose[1]: out_torch = funcs[0](A, B) out_bnb = funcs[1](A, B, fw_code, bw_code) assert out_bnb.dtype == A.dtype, f"bnb matmullt received {A.dtype} but returned {out_bnb.dtype}" n = out_bnb.numel() err = torch.abs(out_bnb - out_torch).float().mean().item() if n > 0: assert err < 0.115 # assert err < 0.20 if any(req_grad): out_bnb.data.copy_(out_torch) torch.cuda.synchronize() loss_bnb = torch.nn.functional.mse_loss(out_bnb, target).mean() loss_bnb.backward() gradA1 = A.grad gradB1 = B.grad A.grad = None B.grad = None loss_torch = torch.nn.functional.mse_loss(out_torch, target).mean() loss_torch.backward() gradA2 = A.grad gradB2 = B.grad A.grad = None B.grad = None if req_grad[0]: torch.testing.assert_close(gradA1, gradA2, atol=0.015, rtol=0.1) if req_grad[1]: n = gradB1.numel() if dim2 > 0: assert torch.abs(gradB1).sum() > 0.0 assert torch.abs(gradB2).sum() > 0.0 else: assert torch.abs(gradB1).sum() == 0.0 assert torch.abs(gradB2).sum() == 0.0 idx = torch.isclose(gradB1, gradB2, atol=0.06, rtol=0.3) assert (idx == 0).sum().item() <= n * 0.1 idx = torch.isclose(gradB1, gradB2, atol=0.10, rtol=0.3) assert (idx == 0).sum().item() <= n * 0.02 grad_err = (gradB1 - gradB2).abs().mean() assert grad_err.item() < 0.003 torch.testing.assert_close(gradB1, gradB2, atol=0.18, rtol=0.3) @pytest.mark.deprecated def test_fp8linear(): b = 10 h = 1024 inp = torch.randn(b, h).cuda() fp32 = torch.nn.Linear(h, h * 2).cuda() fp8 = bnb.research.nn.LinearFP8Mixed(h, h * 2).cuda() fp32b = torch.nn.Linear(h * 2, h).cuda() fp8b = bnb.research.nn.LinearFP8Mixed(h * 2, h).cuda() fp8.weight.data.copy_(fp32.weight.data) fp8.bias.data.copy_(fp32.bias.data) fp8b.weight.data.copy_(fp32b.weight.data) fp8b.bias.data.copy_(fp32b.bias.data) a = fp32b(torch.nn.functional.gelu(fp32(inp))) b = fp8b(torch.nn.functional.gelu(fp8(inp))) err = (a - b).abs().mean() a.mean().backward() b.mean().backward() graderr = (fp8.weight.grad - fp32.weight.grad).abs().mean() bgraderr = (fp8.bias.grad - fp32.bias.grad).abs().mean() assert err < 0.05 assert graderr < 0.00002 assert bgraderr < 0.00002 bitsandbytes/tests/test_functional.py000066400000000000000000001627731506752665600205300ustar00rootroot00000000000000import math import platform import random import time import einops from packaging import version import pytest import torch import bitsandbytes as bnb from bitsandbytes import functional as F from bitsandbytes.cextension import HIP_ENVIRONMENT from tests.helpers import ( BOOLEAN_TUPLES, TRUE_FALSE, describe_dtype, get_available_devices, get_test_dims, id_formatter, is_supported_on_hpu, ) torch.set_printoptions(precision=5, sci_mode=False, linewidth=120, edgeitems=20, threshold=10000) k = 20 def assert_all_approx_close(a, b, rtol=1e-3, atol=1e-3, count=0, throw=True): idx = torch.isclose(a, b, rtol=rtol, atol=atol) sumval = (idx == 0).sum().item() if sumval > count: if throw: print(f"Too many values not close: assert {sumval} < {count}") torch.testing.assert_close(a, b, rtol=rtol, atol=atol) return sumval class FFN(torch.nn.Module): def __init__(self, input_features, hidden_size, bias=True): super().__init__() self.fc1 = torch.nn.Linear(input_features, hidden_size, bias=bias) self.fc2 = torch.nn.Linear(hidden_size, input_features, bias=bias) with torch.no_grad(): torch.nn.init.xavier_uniform_(self.fc1.weight) torch.nn.init.xavier_uniform_(self.fc2.weight) def forward(self, x): x = torch.relu(self.fc1(x)) x = self.fc2(x) return x class Timer: def __init__(self): self.starts = {} self.ends = {} self.agg = {} def tick(self, name="default"): if name not in self.starts: self.starts[name] = torch.cuda.Event(enable_timing=True) self.ends[name] = torch.cuda.Event(enable_timing=True) self.starts[name].record() else: ms = self.tock(name, evict=True, print_ms=False) def tock(self, name="default", evict=True, print_ms=True): if name in self.ends: self.ends[name].record() torch.cuda.synchronize() ms = self.starts[name].elapsed_time(self.ends[name]) if name not in self.agg: self.agg[name] = 0.0 self.agg[name] += ms if evict: self.starts.pop(name) self.ends.pop(name) if print_ms and name in self.agg: print(f"{name} took: {self.agg[name] / 1000.0:.5f}s") return self.agg[name] def reset(self): self.starts = {} self.ends = {} self.agg = {} print("Resetting benchmark data") class Test8BitBlockwiseQuantizeFunctional: @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize("nested", TRUE_FALSE, ids=id_formatter("nested")) @pytest.mark.parametrize( "blocksize", [4096, 2048, 1024, 512, 256, 128, 64] if not HIP_ENVIRONMENT else [4096, 2048, 1024, 512, 256, 128], ) @pytest.mark.parametrize("signed", TRUE_FALSE, ids=id_formatter("signed")) def test_dynamic_blockwise_quantization(self, device, dtype, nested, blocksize, signed): iters = 100 if device != "cuda": iters = 10 # This test is slow in our non-CUDA implementations, so avoid atypical use cases. if nested: pytest.skip("Not a typical use case.") if blocksize != 256: pytest.skip("Only blocksize 256 is used in CPU/MPS/XPU") if dtype != torch.float32: pytest.skip("Only float32 is used in CPU/MPS/XPU") diffs = [] reldiffs = [] for i in range(iters): A1 = torch.randn(1024, 1024, device=device, dtype=dtype) C, S = F.quantize_blockwise(A1, blocksize=blocksize, nested=nested) A2 = F.dequantize_blockwise(C, S) diff = torch.abs(A1 - A2).float() reldiff = diff / torch.abs(A1.float() + 1e-8) diffs.append(diff.mean().item()) reldiffs.append(reldiff.mean().item()) abserr = sum(diffs) / len(diffs) relerr = sum(reldiffs) / len(reldiffs) assert abserr < 0.011 assert relerr < 0.018 assert A2.dtype == dtype diffs = [] code = F.create_dynamic_map(signed=signed) for i in range(iters): A1 = torch.rand(1024, 1024, device=device, dtype=dtype) C, S = F.quantize_blockwise(A1, blocksize=blocksize, nested=nested, code=code) A2 = F.dequantize_blockwise(C, S) diff = torch.abs(A1 - A2).float() reldiff = diff / torch.abs(A1.float() + 1e-8) diffs.append(diff.mean().item()) reldiffs.append(reldiff.mean().item()) # torch.testing.assert_close(A1, A2, atol=1e-2, rtol=0) abserr = sum(diffs) / len(diffs) relerr = sum(reldiffs) / len(reldiffs) if signed: threshold_abserr = 0.0035 assert abserr < 0.0036 assert relerr < 0.015 else: assert abserr < 0.0023 assert relerr < 0.012 assert A2.dtype == dtype @pytest.mark.skipif("cpu" not in get_available_devices(), reason="CPU is required") @pytest.mark.parametrize("hidden", [128]) @pytest.mark.parametrize("blocksize", [4096, 16384]) def test_blockwise_cpu_large(self, hidden, blocksize): diffs = [] reldiffs = [] batch = 128 seq = 128 for i in range(2): A1 = torch.randn(batch, seq, hidden, device="cpu") t0 = time.time() C, S = F.quantize_blockwise(A1, blocksize=blocksize) A2 = F.dequantize_blockwise(C, S, blocksize=blocksize) print(time.time() - t0) diff = torch.abs(A1 - A2) reldiff = diff / torch.abs(A1 + 1e-8) diffs.append(diff.mean().item()) reldiffs.append(reldiff.mean().item()) assert diffs[-1] < 0.011 # print(sum(diffs)/len(diffs)) # print(sum(reldiffs)/len(reldiffs)) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("bits", range(2, 9), ids=id_formatter("bits")) @pytest.mark.parametrize("method", ["linear", "fp8", "dynamic"]) def test_few_bit_quant(self, device, bits, method): if bits != 8 and device == "cpu": pytest.skip("CPU implementation only supports 8 bits") abserrs = [] relerrs = [] code = None if method == "linear": code = F.create_linear_map(True, total_bits=bits).to(device) elif method == "fp8": ebits = math.ceil(bits / 2) pbits = bits - ebits - 1 code = F.create_fp8_map(True, ebits, pbits, bits).to(device) elif method == "dynamic": code = F.create_dynamic_map(True, bits - 0, bits).to(device) # for some data types we have no zero # for some data types we have one zero # for some data types we have two zeros assert torch.unique(code).numel() in [2**bits, 2**bits - 1], f"bits: {bits}, method: {method}" # print(method, (code==0).sum()) assert code.numel() == 256 for i in range(10): values = torch.randn(1, 32, device=device) values /= values.abs().max() # values[values.abs() < 1e-6] += 1e-5 q1 = [] v1 = [] for v in values[0]: idx = torch.abs(v - code).argmin() q1.append(idx.item()) v1.append(code[idx].item()) q1 = torch.tensor(q1, device=device) v1 = torch.tensor(v1, device=device) q2, S2 = F.quantize_blockwise(values, code=code) v2 = F.dequantize_blockwise(q2, S2) idx = torch.isclose(q1.int(), q2.int()) err2 = torch.abs(v2 - values) abserrs.append(err2.mean().item()) relerrs.append((err2 / (1e-10 + values).abs()).mean().item()) if idx.sum(): # some weird cases err1 = torch.abs(v1 - values).mean() # assert err2.mean() <= err1 else: torch.testing.assert_close(q1, q2) @pytest.mark.parametrize("device", get_available_devices()) def test_fp8_quant(self, device): # TODO if device == "cpu": pytest.skip("CPU implementation segfaults") for e_bits in range(1, 7): p_bits = 7 - e_bits code = F.create_fp8_map(True, e_bits, p_bits).to(device) abserr = [] relerr = [] for i in range(10): A1 = torch.randn(1024, 1024, device=device) C, SC = F.quantize_blockwise(A1, code=code) A2 = F.dequantize_blockwise(C, SC) diff = torch.abs(A1 - A2) reldiff = diff / torch.abs(A1 + 1e-8) abserr.append(diff.mean().item()) relerr.append(reldiff.mean().item()) # assert diff < 0.0075 # print(sum(abserr)/len(abserr)) # print(sum(relerr)/len(relerr)) abserr = [] relerr = [] for i in range(10): A1 = torch.rand(1024, 1024, device=device) C, SC = F.quantize_blockwise(A1, code=code) A2 = F.dequantize_blockwise(C, SC) diff = torch.abs(A1 - A2) reldiff = diff / torch.abs(A1 + 1e-8) abserr.append(diff.mean().item()) relerr.append(reldiff.mean().item()) # assert diff < 0.0075 # print(sum(abserr)/len(abserr)) # print(sum(relerr)/len(relerr)) abserr = [] relerr = [] for i in range(10): A1 = torch.randn(1024, 1024, device=device) C, SC = F.quantize_blockwise(A1) A2 = F.dequantize_blockwise(C, SC) diff = torch.abs(A1 - A2) reldiff = diff / torch.abs(A1 + 1e-8) abserr.append(diff.mean().item()) relerr.append(reldiff.mean().item()) # assert diff < 0.0075 # print(3, sum(abserr)/len(abserr)) # print(3, sum(relerr)/len(relerr)) @pytest.mark.benchmark def test_bench_dequantization(self): a = torch.rand(1024, 1024, device="cuda").half() code = F.create_fp8_map(True, 3, 0, 4).cuda() qa, SA = F.quantize_blockwise(a, code=code) print(qa.max()) max_theoretical_mu = 1024 * 1024 * 2 / 1024**3 / 672 * 1000 * 1000 # print(max_theoretical_mu) torch.cuda.synchronize() t0 = time.time() for i in range(100): qa, SA = F.quantize_blockwise(a) torch.cuda.synchronize() # print((time.time()-t0)/1e6) def test_stable_embedding(): layer = bnb.nn.StableEmbedding(1024, 1024) layer.reset_parameters() def quant(x): max1 = torch.abs(x).max() x = torch.round(x / max1 * 127) return max1, x.to(torch.int8) def dequant(c, maxC): return c.float() * (maxC / 127) def mm_dequant(maxA, maxB, C): return C.float() * (maxA / 127) * (maxB / 127) def quant_multi(x, dim): max1 = torch.amax(torch.abs(x), dim=dim, keepdim=True) max1[max1 == 0] = 1.0 x = torch.round(x / max1 * 127) return max1, x.to(torch.int8) def quant_multi_chunk(x, dim, chunk_size=32): if dim == 1: x_chunked = einops.rearrange(x, "(c a) b -> c a b", c=chunk_size) max1 = torch.amax(torch.abs(x_chunked), dim=dim + 1, keepdim=True) max1 = torch.tile(max1, (1, 1, x.shape[1])) max1 = max1.view(x.shape) elif dim == 0: x_chunked = einops.rearrange(x, "a (b c) -> a b c", c=chunk_size) max1 = torch.amax(torch.abs(x_chunked), dim=dim, keepdim=True) max1 = torch.tile(max1, (x.shape[0], 1, 1)) max1 = max1.view(x.shape) max1[max1 == 0] = 1.0 x = torch.round(x / max1 * 127) return max1, x.to(torch.int8) def mean(xx): return sum(xx) / float(len(xx)) methods = { "linear": ( lambda x, dim: quant(x), lambda x, dim: quant(x), dequant, dequant, mm_dequant, ), "vectorwise": (quant_multi, quant_multi, dequant, dequant, mm_dequant), } @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") class TestIGEMMFunctional: @pytest.mark.parametrize("dim1", [1024 * 2], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [1024 * 16], ids=id_formatter("dim2")) @pytest.mark.parametrize("quant_methods", methods.values(), ids=methods.keys()) @pytest.mark.parametrize("batched", TRUE_FALSE, ids=id_formatter("batched")) def test_approx_igemm(self, dim1, dim2, quant_methods, batched): dim1 = dim1 - (dim1 % 32) dim2 = dim2 - (dim2 % 32) errors = [] relerrors = [] # print("") for i in range(5): if batched: A = torch.normal(0, 0.5, size=(32, dim1, dim2 // 32), device="cuda") B = torch.normal(0, 0.5, size=(32, dim2 // 32, dim1), device="cuda") maxA, Ac = quant_methods[0](A, 2) maxB, Bc = quant_methods[1](B, 1) else: A = torch.normal(0, 0.5, size=(dim1, dim2), device="cuda") B = torch.normal(0, 0.5, size=(dim2, dim1), device="cuda") maxA, Ac = quant_methods[0](A, 1) maxB, Bc = quant_methods[1](B, 0) torch.testing.assert_close(quant_methods[2](maxA, Ac), A, atol=0.025, rtol=0.05) if batched: out2 = torch.bmm(A, B) C = torch.bmm(Ac.float(), Bc.float()) else: out2 = torch.mm(A, B) C = F.igemm(Ac, Bc) out = quant_methods[4](maxA, maxB, C) std = out2.std() out /= std out2 /= std err = torch.abs(out - out2) relerr = err / torch.abs(out2) errors.append(err.mean().item()) relerrors.append(relerr.mean().item()) # print(mean(errors)) # print(mean(relerrors)) @pytest.mark.parametrize("hidden_dim", [32, 256], ids=id_formatter("hidden_dim")) @pytest.mark.parametrize("batch_dim", [16, 256], ids=id_formatter("batch_dim")) @pytest.mark.parametrize("seq_dim", [16, 256], ids=id_formatter("seq_dim")) @pytest.mark.parametrize("transpose", BOOLEAN_TUPLES, ids=id_formatter("transpose")) def test_igemm(self, hidden_dim, batch_dim, transpose, seq_dim): hidden_dim = hidden_dim - (hidden_dim % 32) batch_dim = batch_dim - (batch_dim % 16) seq_dim = seq_dim - (seq_dim % 16) for i in range(k): shapeA = (batch_dim, hidden_dim) if not transpose[0] else (hidden_dim, batch_dim) shapeB = ( (32 * random.randint(1, 4), hidden_dim) if transpose[1] else (hidden_dim, 32 * random.randint(1, 4)) ) A = torch.randint(-128, 127, size=shapeA, device="cuda").to(torch.int8) B = torch.randint(-128, 127, size=shapeB, device="cuda").to(torch.int8) if not transpose[0] and not transpose[1]: out2 = torch.matmul(A.float(), B.float()) out = F.igemm(A, B) elif not transpose[0] and transpose[1]: out2 = torch.matmul(A.float(), B.t().float()) out = F.igemm(A, B.t()) elif transpose[0] and not transpose[1]: out2 = torch.matmul(A.t().float(), B.float()) out = F.igemm(A.t(), B) elif transpose[0] and transpose[1]: out2 = torch.matmul(A.t().float(), B.t().float()) out = F.igemm(A.t(), B.t()) torch.testing.assert_close(out.float(), out2) for i in range(k): shapeA = (batch_dim, seq_dim, hidden_dim) shapeB = ( (32 * random.randint(1, 4), hidden_dim) if transpose[1] else (hidden_dim, 32 * random.randint(1, 4)) ) A = torch.randint(-128, 127, size=shapeA, device="cuda").to(torch.int8) B = torch.randint(-128, 127, size=shapeB, device="cuda").to(torch.int8) if not transpose[0] and not transpose[1]: out2 = torch.matmul(A.float(), B.float()) out = F.igemm(A, B) elif not transpose[0] and transpose[1]: out2 = torch.matmul(A.float(), B.t().float()) out = F.igemm(A, B.t()) torch.testing.assert_close(out.float(), out2) @pytest.mark.parametrize("seq_dim", [32, 256, 512], ids=id_formatter("seq_dim")) @pytest.mark.parametrize("hidden_dim", [64, 1024, 4096], ids=id_formatter("hidden_dim")) @pytest.mark.parametrize("batch_dim", [2, 8, 16], ids=id_formatter("batch_dim")) def test_dim3_igemm(self, seq_dim, hidden_dim, batch_dim): seq_dim = seq_dim - (seq_dim % 32) hidden_dim = hidden_dim - (hidden_dim % 32) batch_dim = batch_dim - (batch_dim % 2) for i in range(25): A = torch.randint(-128, 127, size=(batch_dim, seq_dim, hidden_dim), device="cuda").to(torch.int8) B = torch.randint(-128, 127, size=(batch_dim, seq_dim, 1024), device="cuda").to(torch.int8) out2 = torch.einsum("bsi, bso->io", A.float(), B.float()) iout = torch.empty(A.shape[2], B.shape[2], dtype=torch.int32, device=A.device) out = F.igemm(A, B, out=iout) torch.testing.assert_close(out.float(), out2) @pytest.mark.parametrize("seq_dim", [32, 512], ids=id_formatter("seq_dim")) @pytest.mark.parametrize("hidden_dim", [32, 1024 * 4], ids=id_formatter("hidden_dim")) @pytest.mark.parametrize("batch_dim", [2, 16], ids=id_formatter("batch_dim")) @pytest.mark.parametrize("transpose", TRUE_FALSE, ids=id_formatter("transpose")) @pytest.mark.skipif(HIP_ENVIRONMENT, reason="this test is not supported on ROCm yet") def test_minmax_igemm(self, seq_dim, hidden_dim, batch_dim, transpose): def min_max(x): maxA = torch.amax(x, dim=2, keepdim=True) minA = torch.amin(x, dim=2, keepdim=True) scale = (maxA - minA) / 2.0 return (127 * (x - minA - scale) / scale).to(torch.int8), minA, scale seq_dim = seq_dim - (seq_dim % 16) hidden_dim = hidden_dim - (hidden_dim % 16) batch_dim = batch_dim - (batch_dim % 2) errs = [] relerrs = [] errs2 = [] relerrs2 = [] for i in range(k): A = torch.normal(0.0, 0.5, size=(batch_dim, seq_dim, hidden_dim), device="cuda") if transpose: B = torch.normal(0, 0.5, size=(256, hidden_dim), device="cuda") else: B = torch.normal(0, 0.5, size=(hidden_dim, 256), device="cuda") Ac, minA, scale = min_max(A) if transpose: maxB, Bc = quant_multi(B, dim=(1 if transpose else 0)) out = F.igemm(Ac, Bc.t()) out2 = torch.matmul(A, B.t()) offset = B.t().sum(0) * (minA + scale) out = out.float() out = (out * maxB.t() * scale / (127 * 127)) + offset maxA, Ac = quant_multi(A, dim=2) out3 = F.igemm(Ac, Bc.t()) out3 = mm_dequant(maxA, maxB.t(), out3) else: maxB, Bc = quant_multi(B, dim=0) offset = B.sum(0) * (minA + scale) out = F.igemm(Ac, Bc) out2 = torch.matmul(A, B) out = out.float() out = (out * maxB * scale / (127 * 127)) + offset maxA, Ac = quant_multi(A, dim=2) out3 = F.igemm(Ac, Bc) out3 = mm_dequant(maxA, maxB, out3) std = out2.std() out2 /= std out /= std out3 /= std err = torch.abs(out - out2) relerr = err / (torch.abs(out2) + 1e-7) err2 = torch.abs(out3 - out2) relerr2 = err2 / (torch.abs(out2) + 1e-7) errs.append(err.mean().item()) relerrs.append(relerr.mean().item()) errs2.append(err2.mean().item()) relerrs2.append(relerr2.mean().item()) # print(mean(errs)) # print(mean(relerrs)) # print(mean(errs2)) # print(mean(relerrs2)) assert mean(errs) < 0.015 # There's a higher relerr on L40S with torch 2.4+cu118. is_sm89 = torch.cuda.get_device_capability() == (8, 9) if torch.version.cuda == "11.8" and is_sm89 and torch.__version__ < (2, 5): assert mean(relerrs) < 0.41 else: assert mean(relerrs) < 0.3 @pytest.mark.parametrize("dim1", [1, 64], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [32, 128], ids=id_formatter("dim2")) @pytest.mark.parametrize("dim3", [32, 256], ids=id_formatter("dim3")) @pytest.mark.parametrize("dim4", [32, 256], ids=id_formatter("dim4")) @pytest.mark.parametrize("transpose", BOOLEAN_TUPLES, ids=id_formatter("transpose")) def test_ibmm(self, dim1, dim2, dim3, dim4, transpose): dim2 = dim2 - (dim2 % 16) dim3 = dim3 - (dim3 % 16) dim4 = dim4 - (dim4 % 16) for i in range(k): shapeA = (dim1, dim3, dim2) if transpose[0] else (dim1, dim2, dim3) shapeB = (dim1, dim4, dim3) if transpose[1] else (dim1, dim3, dim4) A = torch.randint(-128, 127, size=shapeA, device="cuda").to(torch.int8) B = torch.randint(-128, 127, size=shapeB, device="cuda").to(torch.int8) if not transpose[0] and not transpose[1]: out2 = torch.bmm(A.float(), B.float()) out = F.igemm(A, B) elif not transpose[0] and transpose[1]: out2 = torch.bmm(A.float(), B.permute([0, 2, 1]).float()) out = F.igemm(A, B.permute([0, 2, 1])) elif transpose[0] and not transpose[1]: out2 = torch.bmm(A.permute([0, 2, 1]).float(), B.float()) out = F.igemm(A.permute([0, 2, 1]), B) elif transpose[0] and transpose[1]: out2 = torch.bmm(A.permute([0, 2, 1]).float(), B.permute([0, 2, 1]).float()) out = F.igemm(A.permute([0, 2, 1]), B.permute([0, 2, 1])) torch.testing.assert_close(out.float(), out2.float()) class TestLLMInt8Functional: @staticmethod def vectorwise_mm_dequant(xq, S1, S2, dtype=torch.half): """Reference implementation for the F.int8_mm_dequant function.""" C = 127.0 x = xq.float() if len(S1.shape) == 3 and len(x.shape) == 2: S1 = S1.squeeze(0) if len(S2.shape) == 3 and len(x.shape) == 2: S2 = S2.squeeze(0) if len(S1.shape) == 2: x *= S1 / C else: x *= S1 / C x *= S2 / C return x.to(dtype) @staticmethod def vectorwise_quant(x, dim=1): """Reference implementation""" max1 = torch.amax(torch.abs(x), dim=dim, keepdim=True) xq = torch.round(x * (127.0 / max1)).to(torch.int8) return xq, max1 @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dim1", [128], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [256], ids=id_formatter("dim2")) @pytest.mark.parametrize("dim3", [499, 512], ids=id_formatter("dim3")) @pytest.mark.parametrize("dim4", [512], ids=id_formatter("dim4")) @pytest.mark.parametrize("dims", (2, 3), ids=id_formatter("dims")) @pytest.mark.parametrize("ldb", (0,), ids=id_formatter("ldb")) def test_int8_linear_matmul(self, device, dim1, dim2, dim3, dim4, dims, ldb): for i in range(k): if dims == 2: A = torch.randint(-128, 127, size=(dim1, dim3), dtype=torch.int8, device=device) elif dims == 3: A = torch.randint(-128, 127, size=(dim1, dim2, dim3), dtype=torch.int8, device=device) B = torch.randint(-128, 127, size=(dim4, dim3), dtype=torch.int8, device=device) C1 = torch.matmul(A.float(), B.t().float()) C2 = F.int8_linear_matmul(A, B) torch.testing.assert_close(C1, C2.float()) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dim1", [32], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [32], ids=id_formatter("dim2")) @pytest.mark.parametrize("dim3", [32], ids=id_formatter("dim3")) @pytest.mark.parametrize("dim4", [32], ids=id_formatter("dim4")) @pytest.mark.parametrize("dims", (2,), ids=id_formatter("dims")) def test_int8_linear_matmul_half(self, device, dim1, dim2, dim3, dim4, dims): for i in range(k): if dims == 2: A = torch.normal(0, 0.5, size=(dim1, dim3), device=device).half() elif dims == 3: A = torch.normal(0, 0.5, size=(dim1, dim2, dim3), device=device).half() B = torch.randn((dim4, dim3), device=device).half() torch.nn.init.xavier_uniform_(B) C1 = torch.matmul(A, B.t()) A = A.view(-1, A.shape[-1]) CA, statsA, _ = F.int8_vectorwise_quant(A) CB, statsB, _ = F.int8_vectorwise_quant(B) output = F.int8_mm_dequant(F.int8_linear_matmul(CA, CB), statsA, statsB) torch.testing.assert_close(C1.view(-1, C1.shape[-1]), output, atol=0.025, rtol=0.05) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dim1", (64, 256), ids=id_formatter("dim1")) @pytest.mark.parametrize("dim4", (64, 1024), ids=id_formatter("dim4")) @pytest.mark.parametrize("dims", (2,), ids=id_formatter("dims")) @pytest.mark.parametrize("has_bias", TRUE_FALSE, ids=id_formatter("has_bias")) def test_dequant_mm(self, device, dim1, dim4, dims, has_bias): inner = 128 bias = None if has_bias: bias = torch.randn(dim4, device=device, dtype=torch.float16) for i in range(1): A = torch.randn(dim1, inner, device=device) B = torch.randn(dim4, inner, device=device) C1 = torch.matmul(A.half(), B.t().half()) if has_bias: C1 += bias A1, maxA = self.vectorwise_quant(A, dim=1) B1, maxB = self.vectorwise_quant(B, dim=1) C2 = F.int8_linear_matmul(A1, B1) C4 = self.vectorwise_mm_dequant(C2.float(), maxA, maxB.t()) if has_bias: C4 += bias # TODO: is something wrong here? If so, the problem goes deeper # n = C1.numel() # p = 0.06 std = C1.std(0).view(1, -1) C1 /= std C4 /= std # assert_all_approx_close(C1, C4, atol=0.02, rtol=0.1, count=int(n*0.06)) # assert (count / n < p), f"error in more than {p} of elements: {count}/{n}={count/n}" C5 = F.int8_mm_dequant(C2, maxA, maxB, bias=bias) C5 /= std torch.testing.assert_close(C5, C4, atol=0.015, rtol=0.1) n = C5.numel() assert_all_approx_close(C1, C4, atol=0.015, rtol=0.1, count=int(0.01 * n)) @pytest.mark.parametrize("dim1", [1 * 1024], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [1 * 1024], ids=id_formatter("dim2")) @pytest.mark.parametrize("dims", (2,), ids=id_formatter("dims")) @pytest.mark.parametrize("threshold", [0.0, 3.0], ids=id_formatter("decomp")) @pytest.mark.deprecated def test_colrow_absmax(self, dim1, dim2, dims, threshold): for i in range(k): A = torch.randn(dim1, dim2, device="cuda").half() assert dims == 2 row_stats1, _ = torch.abs(A.float()).max(1) col_stats1, _ = torch.abs(A.float()).max(0) if threshold > 0.0: A_truncated = A.clone() A_truncated[torch.abs(A_truncated) >= threshold] = 0.0 row_stats1_trunc, _ = torch.abs(A_truncated.float()).max(1) col_stats1_trunc, _ = torch.abs(A_truncated.float()).max(0) row_stats2, col_stats2, nnz_block_ptr2 = F.get_colrow_absmax(A, threshold=threshold) nnz_rows1_counts = (torch.abs(A) >= threshold).sum(1).flatten() nnz_block_ptr1 = torch.zeros( nnz_rows1_counts.shape[0] + 1, dtype=nnz_rows1_counts.dtype, device=nnz_rows1_counts.device, ) nnz_block_ptr1[1:] = nnz_rows1_counts.cumsum(0) torch.testing.assert_close(col_stats1_trunc, col_stats2) torch.testing.assert_close(row_stats1_trunc, row_stats2) # torch.testing.assert_close(nnz_block_ptr1, nnz_block_ptr2) else: row_stats2, col_stats2, nnz_block_ptr2 = F.get_colrow_absmax(A, threshold=0.0) assert nnz_block_ptr2 is None torch.testing.assert_close(col_stats1, col_stats2) torch.testing.assert_close(row_stats1, row_stats2) @pytest.mark.parametrize("dim1", [2048, 4096], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [512, 1024], ids=id_formatter("dim2")) @pytest.mark.deprecated def test_int8_double_quant(self, dim1, dim2): for i in range(k): A = torch.randn(dim1, dim2, device="cuda").half() out_col1, Scol = self.vectorwise_quant(A, dim=0) out_row1, Srow = self.vectorwise_quant(A, dim=1) CA, CAt, statsA, statsAt, _ = F.int8_double_quant(A) # max difference is 1 due to rounding differences torch.testing.assert_close(CA, out_row1, atol=1, rtol=0) torch.testing.assert_close(CAt, out_col1, atol=1, rtol=0) n = CAt.numel() num_not_close_rows = (torch.isclose(CA, out_row1, atol=1) == 0).sum().item() num_not_close_cols = (torch.isclose(CAt, out_col1, atol=1) == 0).sum().item() # allow for 1:500 error due to rounding differences min_error = 1 / 500 if num_not_close_cols > (min_error * n): print( f"Min error exceeded {num_not_close_cols} elements are different. Error: {num_not_close_cols / n:.4f}" ) assert False if num_not_close_rows > (min_error * n): print( f"Min error exceeded {num_not_close_rows} elements are different. Error: {num_not_close_rows / n:.4f}" ) assert False torch.testing.assert_close(Srow.flatten().float(), statsA) torch.testing.assert_close(Scol.flatten().float(), statsAt) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize( ("dim1", "dim4", "inner"), ( pytest.param(dim1, dim4, inner, id=f"{dim1=},{dim4=},{inner=}") for (dim1, dim4, inner) in zip( (1, 8, 2048, 4096), (2, 128, 2048, 4096), (4, 256, 512, 4096), ) ), ) def test_integrated_int8_linear_matmul(self, device, dim1, dim4, inner): if device == "cpu" and inner > 2048: pytest.skip("Slow on CPU") for i in range(k): A = torch.randn(dim1, inner, device=device).half() B = torch.randn(dim4, inner, device=device).half() out1 = torch.matmul(A.half(), B.t().half()) C1a, stats1a, _ = F.int8_vectorwise_quant(A) C2a, stats2a, _ = F.int8_vectorwise_quant(B) A1, maxA = self.vectorwise_quant(A, dim=1) B1, maxB = self.vectorwise_quant(B, dim=1) torch.testing.assert_close(maxA.flatten().float(), stats1a) torch.testing.assert_close(maxB.flatten().float(), stats2a) torch.testing.assert_close(C1a, A1, rtol=0, atol=1) torch.testing.assert_close(C2a, B1, rtol=0, atol=1) out2 = F.int8_linear_matmul(A1, B1) C2 = F.int8_linear_matmul(A1, B1) out3 = self.vectorwise_mm_dequant(C2.float(), maxA, maxB.t()) err1 = torch.abs(out1 - out2).mean().item() err2 = torch.abs(out1 - out3).mean().item() assert err2 <= err1 * 1.025 @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dim1", [512, 2048], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [1024, 4096], ids=id_formatter("dim2")) def test_coo_double_quant(self, device, dim1, dim2): threshold = 2.00 for i in range(k): A = torch.randn(dim1, dim2, device=device).half() idx = torch.abs(A) >= threshold CA, statsA, outlier_cols = F.int8_vectorwise_quant(A, threshold=threshold) if outlier_cols is not None: A1 = A * idx A2 = torch.zeros_like(A) + A1 torch.testing.assert_close(A1, A2) A[:, outlier_cols] = 0 A2 = (CA.float() * statsA.unsqueeze(1) / 127).half() torch.testing.assert_close(A, A2, rtol=0.05, atol=1.5e-2) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dim1", [512, 2048], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [1024, 4096], ids=id_formatter("dim2")) def test_coo_int8_vectorwise_quant(self, device, dim1, dim2): threshold = 3.00 for i in range(k): A = torch.randn(dim1, dim2, device=device).half() idx = torch.abs(A) >= threshold CA, statsA, outlier_cols = F.int8_vectorwise_quant(A, threshold=threshold) if outlier_cols is not None: A2 = (CA.float() * statsA.unsqueeze(1) / 127).half() A[:, outlier_cols] = 0 torch.testing.assert_close(A * (idx == 0), A2, rtol=0.05, atol=1.5e-2) @pytest.mark.skipif(HIP_ENVIRONMENT, reason="this test is not supported on ROCm yet") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") class TestSpMMFunctional: @pytest.mark.parametrize("dim1", [256, 1024], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [128, 512], ids=id_formatter("dim2")) @pytest.mark.parametrize("transposed_B", TRUE_FALSE, ids=id_formatter("transposed_B")) def test_spmm_coo(self, dim1, dim2, transposed_B): threshold = 1.5 dim3 = torch.randint(32, 128, size=(1,)).item() # dim3 = 17 for i in range(k): A = torch.randn(dim1, dim2).cuda().half() if transposed_B: B = torch.randn(dim3, dim2).cuda().half() else: B = torch.randn(dim2, dim3).cuda().half() idx = torch.abs(A) >= threshold nnz = (idx == 1).sum().item() rows, cols = torch.where(idx) values = A[idx] cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values) A2 = A * idx if transposed_B: out2 = F.spmm_coo(cooA, B.t()) out1 = torch.matmul(A2, B.t()) else: out2 = F.spmm_coo(cooA, B) out1 = torch.matmul(A2, B) assert_all_approx_close(out1, out2, rtol=0.01, atol=3.0e-2, count=30) @pytest.mark.benchmark def test_spmm_bench(self): batch = 2 model = 1024 * 1 hidden = model * 4 seq = 1024 dim1 = batch * seq dim2 = model dim3 = hidden threshold = 4 A = torch.randn(dim1, dim2, device="cuda").half() B = torch.randn(dim2, dim3, device="cuda").half() for i in range(10): C1 = bnb.matmul(A, B.t()) torch.cuda.synchronize() t0 = time.time() for i in range(k): C1 = bnb.matmul(A, B.t()) torch.cuda.synchronize() t8 = time.time() - t0 idx = torch.abs(A) >= threshold nnz = (idx == 1).sum().item() print(nnz / idx.numel()) rows, cols = torch.where(idx) values = A[idx] cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values) for i in range(10): out2 = F.spmm_coo(cooA, B) torch.cuda.synchronize() t0 = time.time() for i in range(k): out2 = F.spmm_coo(cooA, B) torch.cuda.synchronize() tsp = time.time() - t0 print(tsp, t8) print(tsp / t8) @pytest.mark.parametrize("dim1", [1 * 2048], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [12288], ids=id_formatter("dim2")) @pytest.mark.parametrize("dtype", [torch.float16], ids=describe_dtype) @pytest.mark.parametrize("out_func", ["zeros", "ones"], ids=id_formatter("out_func")) def test_spmm_coo_very_sparse(self, dim1, dim2, dtype, out_func): out_func = getattr(torch, out_func) threshold = 3.3 # threshold = 2.8 # threshold = 0.0 A = torch.randn(dim1, dim2, device="cuda").half() if dtype == torch.float16: B = torch.randn(dim2, dim2 * 4, device="cuda").half() torch.nn.init.xavier_uniform_(B) else: B = torch.randn(dim2, dim2 * 4, device="cuda").half() torch.nn.init.xavier_uniform_(B) SB = torch.abs(B).max().float() B = torch.round(B / SB * 127).to(torch.int8) print("") idx = torch.abs(A) >= threshold nnz = (idx == 1).sum().item() rows, cols = torch.where(idx) values = A[idx] cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values) A2 = A * idx out1 = torch.matmul(A2.half(), B.half()) out = out_func(out1.shape, dtype=torch.float16, device=out1.device) out1 += out.clone() out2 = F.spmm_coo_very_sparse(cooA, B, out=out) # print(B) # print(out1) # print(out2) p = 200 / (2048 * 12288 * 4) n = out1.numel() count = math.ceil(p * n) std = out1.std() out1 /= std out2 /= std assert_all_approx_close(out1, out2.half(), rtol=0.01, atol=3.0e-2, count=count) # assert_all_approx_close(out1, out2.half(), rtol=0.05, atol=0.01, count=count) idx_col = torch.randint(0, A2.shape[-1], size=(15,)) # torch.testing.assert_close(out1, out2.half(), rtol=0.05, atol=0.001) # Bt = torch.randn(dim2*4, dim2, device='cuda').half() # torch.cuda.synchronize() # t0 = time.time() # print(A2.shape, B.shape) # for i in range(100): # #out3 = F.spmm_coo(cooA, Bt.t()) # #out2 = F.spmm_coo(cooA, B) # #out2 = F.spmm_coo_very_sparse(cooA, B) # #out1 = torch.matmul(A, Bt.t()) # torch.cuda.synchronize() # print(time.time() - t0) @pytest.mark.parametrize("dim1", [1 * 2048]) @pytest.mark.parametrize("dim2", [2048]) @pytest.mark.parametrize("dtype", [torch.int8]) def test_spmm_coo_dequant(self, dim1, dim2, dtype): threshold = 6.0 # threshold = 2.8 # threshold = 0.0 A = torch.randn(dim1, dim2, device="cuda").half() B = torch.empty(dim2, dim2 * 4, device="cuda", dtype=torch.float16) torch.nn.init.xavier_uniform_(B) Bt = B.t().contiguous() CB, CBt, statsB, statsBt, coo_tensor = F.int8_double_quant(B) rowidx = torch.randint(0, A.shape[-1], size=(15,)) A[:, rowidx] = 8.0 idx = torch.abs(A) >= threshold nnz = (idx == 1).sum().item() rows, cols = torch.where(idx) values = A[idx] cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values) A2 = A * idx out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt) out1 = torch.matmul(A2, B.half()) out3 = F.spmm_coo_very_sparse(cooA, CBt.half()) out3 = out3 * statsBt.half() / 127 values, counts = torch.unique(cooA.rowidx, return_counts=True) offset = counts.cumsum(0).int() max_count, max_idx = torch.sort(counts, descending=True) print(torch.median(max_count.float())) torch.testing.assert_close(out2, out3, rtol=0.05, atol=0.001) p = 200 / (2048 * 12288 * 4) n = out1.numel() count = math.ceil(p * n) assert_all_approx_close(out1, out2, rtol=0.01, atol=3.0e-2, count=count) # torch.cuda.synchronize() # t0 = time.time() # for i in range(100): # out2 = F.spmm_coo_very_sparse(cooA, B) # torch.cuda.synchronize() # print('fp16', time.time() - t0) torch.cuda.synchronize() t0 = time.time() for i in range(100): out2 = F.spmm_coo(cooA, B) torch.cuda.synchronize() print("cusparse fp16", time.time() - t0) torch.cuda.synchronize() t0 = time.time() for i in range(100): out2 = F.spmm_coo_very_sparse(cooA, CBt) torch.cuda.synchronize() print("int8", time.time() - t0) torch.cuda.synchronize() t0 = time.time() for i in range(100): out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt) torch.cuda.synchronize() print("int8+dequant", time.time() - t0) torch.cuda.synchronize() t0 = time.time() for i in range(100): out2 = torch.matmul(A, B) torch.cuda.synchronize() print("matmul", time.time() - t0) torch.cuda.synchronize() t0 = time.time() for i in range(100): out1 = bnb.matmul(A, Bt) out2 = F.spmm_coo_very_sparse(cooA, CBt, dequant_stats=statsBt) out = out1 + out2 torch.cuda.synchronize() print("sparse+ matmul", time.time() - t0) torch.cuda.synchronize() t0 = time.time() for i in range(100): out1 = bnb.matmul(A, Bt) torch.matmul(A[:, rowidx], Bt.t()[rowidx], out=out1) torch.cuda.synchronize() print("partial matmul", time.time() - t0) torch.cuda.synchronize() t0 = time.time() for i in range(100): out1 = bnb.matmul(A, Bt) torch.cuda.synchronize() print("partial matmul", time.time() - t0) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") class TestSparseTensorFunctional: def test_coo2csr(self): threshold = 1 A = torch.randn(128, 128).half().cuda() idx = torch.abs(A) >= threshold nnz = (idx == 1).sum().item() rows, cols = torch.where(idx) values = A[idx] cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values) A2 = A * idx csrA = F.coo2csr(cooA) counts = csrA.rowptr[1:] - csrA.rowptr[:-1] assert counts.numel() == A.shape[0] torch.testing.assert_close(counts.long(), (A2 != 0).sum(1)) idx = A2 != 0 torch.testing.assert_close(A2[idx], csrA.values) def test_coo2csc(self): threshold = 1 A = torch.randn(128, 128).half().cuda() idx = torch.abs(A) >= threshold nnz = (idx == 1).sum().item() rows, cols = torch.where(idx) values = A[idx] cooA = F.COOSparseTensor(A.shape[0], A.shape[1], nnz, rows.int(), cols.int(), values) A2 = A * idx cscA = F.coo2csc(cooA) counts = cscA.colptr[1:] - cscA.colptr[:-1] assert counts.numel() == A.shape[1] torch.testing.assert_close(counts.long(), (A2 != 0).sum(0)) # torch uses row-major -> use transpose to transfer to col-major idx = A2.t() != 0 torch.testing.assert_close(A2.t()[idx], cscA.values) class TestQuantize4BitFunctional: @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) @pytest.mark.parametrize( "blocksize", [64, 128, 256, 512, 1024, 2048, 4096] if not HIP_ENVIRONMENT else [128, 256, 512, 1024, 2048, 4096], ) def test_4bit_quant(self, device, dtype, quant_type, blocksize): if device == "hpu" and not is_supported_on_hpu(quant_type, dtype): pytest.skip("This configuration is not supported on HPU.") A1 = torch.randn(1024, 1024, device=device, dtype=dtype) qa, SA = F.quantize_4bit(A1, blocksize=blocksize, quant_type=quant_type) A2 = F.dequantize_4bit(qa, SA, blocksize=blocksize, quant_type=quant_type) err = (A1 - A2).abs().float() relerr = (err / (A1.abs().float() + 1e-8)).mean() err = err.mean() assert A2.dtype == dtype # With larger block sizes, we can expect this to blow up. # At blocksize>=1024, don't even bother looking at relerr. # # Actually, the above is not true anymore after fixing the integer packing bug. # The following values were taken from averaging 1k samples per test configuration after fixing the bug. error_dict = dict() error_dict["fp4"] = dict() error_dict["nf4"] = dict() error_dict["fp4"]["err"] = { 64: 0.096545, 128: 0.102947, 256: 0.108685, 512: 0.114087, 1024: 0.119312, 2048: 0.124460, 4096: 0.129573, } error_dict["fp4"]["rel_err"] = { 64: 0.260130, 128: 0.275734, 256: 0.289842, 512: 0.302852, 1024: 0.314982, 2048: 0.326402, 4096: 0.337228, } error_dict["nf4"]["err"] = { 64: 0.072792, 128: 0.076835, 256: 0.080326, 512: 0.083535, 1024: 0.086603, 2048: 0.089592, 4096: 0.092537, } error_dict["nf4"]["rel_err"] = { 64: 0.203299, 128: 0.215252, 256: 0.226044, 512: 0.236021, 1024: 0.245365, 2048: 0.254146, 4096: 0.262457, } # Allow higher tolerance for fp32 on CPU with larger block sizes reltol = 2.8e-3 if dtype == torch.float32 and blocksize >= 128 and device == "cpu" else 1e-3 errtol = 1.2e-3 if dtype == torch.float32 and blocksize >= 1024 and device == "cpu" else 1e-3 assert err < error_dict[quant_type]["err"][blocksize] + errtol assert relerr < error_dict[quant_type]["rel_err"][blocksize] + reltol @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) @pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128], ids=id_formatter("blocksize")) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=describe_dtype) def test_4bit_compressed_stats(self, device, quant_type, blocksize, dtype): if device == "hpu" and not is_supported_on_hpu(quant_type, dtype): pytest.skip("FP4 quantization is not supported on HPU.") errs1 = [] errs2 = [] for i in range(10): A1 = torch.randn(1024, 1024, device=device, dtype=dtype) q2, SA2 = F.quantize_4bit(A1, blocksize=blocksize, quant_type=quant_type) q3, SA3 = F.quantize_4bit(A1, blocksize=blocksize, compress_statistics=True, quant_type=quant_type) A2 = F.dequantize_4bit(q2, SA2, quant_type=quant_type) A3 = F.dequantize_4bit(q3, SA3, quant_type=quant_type) err = (A1 - A2).abs().float() relerr = (err / (A1.abs().float() + 1e-15)).mean() err = err.mean() errs1.append(err.item()) assert err.item() < 0.11 assert relerr.item() < 0.28 err = (A1 - A3).abs().float() relerr = (err / (A1.abs().float() + 1e-15)).mean() err = err.mean() errs2.append(err.item()) assert err.item() < 0.11 assert relerr.item() < 0.28 # @pytest.mark.parametrize("quant_type", ['fp4', 'nf4']) @pytest.mark.parametrize("quant_type", ["nf4"]) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @pytest.mark.benchmark def test_bench_4bit_dequant(self, quant_type): blocksize = 256 a = torch.rand(1024 * 12 * 4, 1024 * 12, device="cuda").half() qa, SA = F.quantize_4bit(a, blocksize=blocksize, quant_type=quant_type) input_size = a.numel() / 2 output_size = a.numel() * 2 num_bytes = input_size + output_size GB = num_bytes / 1e9 max_theoretical_s = GB / 768 # print(max_theoretical_s*1e6) b = torch.randn(128, 1024 * 12, device="cuda").half() iters = 100 torch.cuda.synchronize() t0 = time.time() for i in range(iters): F.dequantize_4bit(qa, SA, blocksize=blocksize, quant_type=quant_type) # b.copy_(a) torch.cuda.synchronize() # print((time.time()-t0)/iters*1e6) # torch.cuda.synchronize() # t0 = time.time() # for i in range(iters): # torch.matmul(b, a.t()) # torch.cuda.synchronize() # print((time.time()-t0)/iters*1e6) @pytest.mark.skipif( HIP_ENVIRONMENT, reason="gemv 4bit tests are partially enabled on MI300, others being fixed for warpsize 64" ) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("double_quant", TRUE_FALSE, ids=lambda double_quant: f"DQ_{double_quant}") @pytest.mark.parametrize("storage_type", ["nf4", "fp4"]) @pytest.mark.parametrize("kind", ["fc1", "fc2", "attn", "attn_packed"]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype) @pytest.mark.parametrize( "quant_storage", [torch.uint8, torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype, ) @pytest.mark.parametrize("dim", [128, 256, 512, 1024], ids=id_formatter("dim")) def test_gemv_4bit(self, device, dim, dtype, storage_type, quant_storage, double_quant, kind): if device == "hpu" and not is_supported_on_hpu(storage_type, dtype, quant_storage): pytest.skip("This configuration is not supported on HPU.") errs1 = [] errs2 = [] errs3 = [] relerrs1 = [] relerrs2 = [] relerrs3 = [] max_errs1 = [] max_errs2 = [] max_errs3 = [] # Large number of iterations is excessive and slow on CPU. # Keep for CUDA/XPU for now. iters = 10 if device == "cpu" else 100 for i in range(iters): if kind == "fc1": A = torch.randn(1, dim, dtype=dtype, device=device) B = torch.randn(dim * 4, dim, dtype=dtype, device=device) / math.sqrt(dim) elif kind == "fc2": A = torch.randn(1, 4 * dim, dtype=dtype, device=device) B = torch.randn(dim, 4 * dim, dtype=dtype, device=device) / math.sqrt(dim) elif kind == "attn": A = torch.randn(1, dim, dtype=dtype, device=device) B = torch.randn(dim, dim, dtype=dtype, device=device) / math.sqrt(dim) elif kind == "attn_packed": A = torch.randn(1, dim, dtype=dtype, device=device) B = torch.randn(dim * 3, dim, dtype=dtype, device=device) / math.sqrt(dim) qB, state = F.quantize_4bit( B, quant_type=storage_type, compress_statistics=double_quant, quant_storage=quant_storage, ) C3 = torch.matmul(A, B.t()) C2 = F.gemv_4bit(A, qB.t(), state=state) A.requires_grad = True C1 = bnb.matmul_4bit(A, qB.t(), state) err1 = (C1 - C2).abs().float() err2 = (C3 - C2).abs().float() err3 = (C3 - C1).abs().float() mag1 = torch.abs(C1).float() + 1e-5 mag2 = torch.abs(C3).float() + 1e-5 mag3 = torch.abs(C3).float() + 1e-5 relerr1 = err1 / mag1 relerr2 = err2 / mag2 relerr3 = err3 / mag3 max_err1 = err1.max() max_err2 = err2.max() max_err3 = err3.max() errs1.append(err1.mean().item()) errs2.append(err2.mean().item()) errs3.append(err3.mean().item()) relerrs1.append(relerr1.mean().item()) relerrs2.append(relerr2.mean().item()) relerrs3.append(relerr3.mean().item()) max_errs1.append(max_err1.item()) max_errs2.append(max_err2.item()) max_errs3.append(max_err3.item()) c = int(C1.numel() * 0.0014 * (dim / 256)) + 1 c = assert_all_approx_close(C1, C2, 1e-5, 0.01, count=0, throw=False) err1 = sum(errs1) / len(errs1) / math.sqrt(dim) err2 = sum(errs2) / len(errs2) / math.sqrt(dim) err3 = sum(errs3) / len(errs3) / math.sqrt(dim) relerr1 = sum(relerrs1) / len(relerrs1) / math.sqrt(dim) relerr2 = sum(relerrs2) / len(relerrs2) / math.sqrt(dim) relerr3 = sum(relerrs3) / len(relerrs3) / math.sqrt(dim) maxerr1 = sum(max_errs1) / len(max_errs1) / math.sqrt(dim) maxerr2 = sum(max_errs2) / len(max_errs2) / math.sqrt(dim) maxerr3 = sum(max_errs3) / len(max_errs3) / math.sqrt(dim) absratio = err2 / err3 relratio = relerr2 / relerr3 maxratio = relerr2 / relerr3 # for debugging if the tests fails # # print('='*80) # print(f'For matmul: {A.shape}, {B.shape}, {kind}, {dtype}, {storage_type}, double_quant={double_quant}:') # print(C1.flatten()[-20:]) # print(C2.flatten()[-20:]) # print(f'inference vs training abs: {err1}') # print(f'inference vs training rel: {relerr1}') # print(f'inference vs training max: {maxerr1}') # print(f'inference vs training vs torch err ratio abs: {absratio}') # print(f'inference vs training vs torch err ratio rel: {relratio}') # print(f'inference vs training vs torch err ratio max: {maxratio}') if dtype == torch.float16: if dim <= 512: assert err1 < 7e-5 # TODO(matthewdouglas): On T4, dim=128-fp16-fc2-fp4-DQ will have relerror ~ 0.00092727 if ( device == "cuda" and double_quant and storage_type == "fp4" and kind == "fc2" and torch.cuda.get_device_capability() == (7, 5) ): assert relerr1 < 0.00093 else: assert relerr1 < 0.0008 else: assert err1 < 6e-5 assert relerr1 < 2e-4 assert absratio < 1.005 and absratio > 0.995 assert relratio < 1.005 and relratio > 0.992 assert maxratio < 1.005 and maxratio > 0.992 elif dtype == torch.float32: if dim <= 512: assert err1 < 5e-8 assert relerr1 < 1e-6 assert maxerr1 < 1.05e-7 else: assert err1 < 5e-8 assert relerr1 < 8e-6 assert maxerr1 < 1e-7 assert absratio < 1.005 and absratio > 0.995 assert relratio < 1.005 and relratio > 0.995 assert maxratio < 1.005 and maxratio > 0.995 elif dtype == torch.bfloat16: if dim <= 512: relerr_thres = 0.013 if hasattr(torch, "xpu") and torch.xpu.is_available() else 0.007 assert err1 < 6e-4 assert relerr1 < relerr_thres assert maxerr1 < 0.015 else: assert err1 < 2e-4 assert relerr1 < 0.002 assert maxerr1 < 0.0012 assert absratio < 1.005 and absratio > 0.995 assert relratio < 1.05 and relratio > 0.96 assert maxratio < 1.05 and maxratio > 0.97 @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("storage_type", ["nf4", "fp4"], ids=["nf4", "fp4"]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype) @pytest.mark.skipif(HIP_ENVIRONMENT, reason="this test is not supported on ROCm yet") def test_gemv_eye_4bit(self, device, storage_type, dtype): if device == "hpu" and not is_supported_on_hpu(storage_type, dtype): pytest.skip("This configuration is not supported on HPU.") if ( device == "cpu" and platform.system() == "Windows" and version.parse(torch.__version__).release == (2, 8, 0) ): pytest.skip("Regression: CPU crash on Windows with torch 2.8.0") dims = 4 dims = get_test_dims(0, 8192, n=dims) dims = [dim + (64 - (dim % 64)) for dim in dims] # for dim in [576, 5120, 3520, 5184, 1280, 4992, 5312, 2048]: for dim in dims: A = torch.normal(0, 0.1, size=(1, 1, dim), dtype=dtype, device=device) B = torch.eye(dim, dtype=dtype, device=device) qB, state = F.quantize_4bit(B, quant_type=storage_type, compress_statistics=False) C3 = torch.matmul(A, B.t()) C2 = bnb.matmul_4bit(A, qB.t(), state) A.requires_grad = True C1 = bnb.matmul_4bit(A, qB.t(), state) torch.testing.assert_close(A, C3) torch.testing.assert_close(A, C1) torch.testing.assert_close(A, C2) # torch.testing.assert_close(A, C1, rtol=1e-5, atol=0.00001) # torch.testing.assert_close(A, C2, rtol=1e-5, atol=0.080) def test_normal_map_tree(): code = F.create_normal_map() values = code[:8].tolist() + code[-8:].tolist() num_pivots = 1 # print(values) while num_pivots < 16: idx = list(range(16 // num_pivots // 2, 16, 16 // num_pivots)) # print(idx) num_pivots *= 2 pivots = [] for i in idx: pivots.append((values[i - 1] + values[i]) / 2) # print(pivots) bitsandbytes/tests/test_generation.py000066400000000000000000000101261506752665600205010ustar00rootroot00000000000000from itertools import product import math import pytest import torch from tests.helpers import TRUE_FALSE, describe_dtype, id_formatter transformers = pytest.importorskip("transformers") def get_4bit_config(): return transformers.BitsAndBytesConfig( load_in_4bit=True, load_in_8bit=False, llm_int8_threshold=6.0, llm_int8_has_fp16_weight=False, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", ) def get_model_and_tokenizer(config): model_name_or_path, quant_type = config bnb_config = get_4bit_config() if quant_type == "16bit": bnb_config.load_in_4bit = False else: bnb_config.bnb_4bit_quant_type = quant_type model = transformers.AutoModelForCausalLM.from_pretrained( model_name_or_path, quantization_config=bnb_config, max_memory={0: "48GB"}, device_map="auto", torch_dtype=torch.bfloat16, ).eval() tokenizer = transformers.AutoTokenizer.from_pretrained(model_name_or_path) return model, tokenizer def get_prompt_for_generation_eval(text, add_roles=True): description = ( "A chat between a curious human and an artificial intelligence assistant. " "The assistant gives helpful, detailed, and polite answers to the user's questions." ) if add_roles: prompt = f"{description} ### Human: {text} ### Assistant:" else: prompt = f"{description} {text}" return prompt def generate(model, tokenizer, text, generation_config, prompt_func=get_prompt_for_generation_eval): text = prompt_func(text) inputs = tokenizer(text, return_tensors="pt").to("cuda:0") outputs = model.generate(inputs=inputs["input_ids"], generation_config=generation_config) return tokenizer.decode(outputs[0], skip_special_tokens=True) models = ["bigscience/bloom-1b7"] dtypes = ["nf4", "fp4"] @pytest.fixture(scope="session", params=product(models, dtypes)) def model_and_tokenizer(request): model, tokenizer = get_model_and_tokenizer(request.param) yield request.param, model, tokenizer del model @pytest.mark.parametrize("DQ", TRUE_FALSE, ids=id_formatter("dq")) @pytest.mark.parametrize("inference_kernel", TRUE_FALSE, ids=id_formatter("inference_kernel")) @pytest.mark.parametrize("dtype", [torch.float16], ids=describe_dtype) @pytest.mark.slow def test_pi(requires_cuda, model_and_tokenizer, inference_kernel, DQ, dtype): fixture_config, model, tokenizer = model_and_tokenizer generation_config = transformers.GenerationConfig( max_new_tokens=20, do_sample=True, top_p=0.9, temperature=0.7, ) generation_config.max_new_tokens = 20 # text = 'Please write down the first 50 digits of pi.' # text = get_prompt_for_generation_eval(text) # text += ' Sure, here the first 50 digits of pi: 3.14159' n_cases = 6 text = "3.14159" if hasattr(model.config, "quantization_config"): model.config.quantization_config.bnb_4bit_compute_dtype = dtype model.config.quantization_config.bnb_4bit_use_double_quant = DQ if not inference_kernel: text = [text] * n_cases inputs = tokenizer(text, return_tensors="pt").to("cuda:0") x = inputs["input_ids"] outputs = [] if inference_kernel: for i in range(n_cases): output = model.generate(x, generation_config=generation_config) textout = tokenizer.decode(output[0], skip_special_tokens=True) outputs.append(textout) else: outputs = model.generate(x, generation_config=generation_config) outputs = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs] assert len(outputs) == n_cases failure_count = 0 for i in range(n_cases): if outputs[i][: len(str(math.pi))] != str(math.pi): failure_count += 1 failure_max = 2 if fixture_config[0] == "huggyllama/llama-7b" else 4 if failure_count > failure_max: print(math.pi) for out in outputs: print(out) raise ValueError(f"Failure count: {failure_count}/{n_cases}") bitsandbytes/tests/test_linear4bit.py000066400000000000000000000346171506752665600204160ustar00rootroot00000000000000import copy import os import pickle import platform from tempfile import TemporaryDirectory import pytest import torch import bitsandbytes as bnb from bitsandbytes.cextension import HIP_ENVIRONMENT from tests.helpers import ( TRUE_FALSE, describe_dtype, get_available_devices, id_formatter, is_supported_on_hpu, torch_load_from_buffer, torch_save_to_buffer, ) storage = { "uint8": torch.uint8, "float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32, } @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("quant_storage", ["uint8", "float16", "bfloat16", "float32"]) @pytest.mark.parametrize("original_dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("bias", TRUE_FALSE, ids=id_formatter("bias")) @pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics")) @pytest.mark.parametrize("quant_type", ["nf4", "fp4"]) @pytest.mark.parametrize("save_before_forward", TRUE_FALSE, ids=id_formatter("save_before_forward")) def test_linear_serialization( device, quant_type, original_dtype, compress_statistics, bias, quant_storage, save_before_forward ): if device == "hpu" and not is_supported_on_hpu(quant_type, original_dtype, storage[quant_storage]): pytest.skip("This configuration is not supported on HPU.") compute_dtype = None layer_shape = (300, 400) linear = torch.nn.Linear(*layer_shape, dtype=original_dtype, device="cpu") # original layer # Quantizing original layer linear_q = bnb.nn.Linear4bit( linear.in_features, linear.out_features, bias=bias, compute_dtype=compute_dtype, compress_statistics=compress_statistics, quant_type=quant_type, device="meta", ) new_weight = bnb.nn.Params4bit(data=linear.weight, quant_type=quant_type, requires_grad=False) linear_q.weight = new_weight if bias: linear_q.bias = torch.nn.Parameter(linear.bias) linear_q = linear_q.to(device) # saving to state_dict: sd = linear_q.state_dict() # restoring from state_dict: bias_data2 = sd.pop("bias", None) weight_data2 = sd.pop("weight") weight2 = bnb.nn.Params4bit.from_prequantized(quantized_stats=sd, data=weight_data2, device=device) # creating new layer with same params: linear_q2 = bnb.nn.Linear4bit( linear.in_features, linear.out_features, bias=bias, compute_dtype=compute_dtype, compress_statistics=compress_statistics, quant_type=quant_type, device="meta", ) # loading weights from state_dict: linear_q2.weight = weight2 if bias: linear_q2.bias = torch.nn.Parameter(bias_data2) linear_q2 = linear_q2.to(device) # MATCHING a, b = linear_q.weight, linear_q2.weight # Quantizing original layer with specified quant_storage type linear_qs = bnb.nn.Linear4bit( linear.in_features, linear.out_features, bias=bias, compute_dtype=compute_dtype, compress_statistics=compress_statistics, quant_type=quant_type, quant_storage=storage[quant_storage], device="meta", ) linear_qs.weight = bnb.nn.Params4bit( data=linear.weight, requires_grad=False, quant_type=quant_type, quant_storage=storage[quant_storage], ) if bias: linear_qs.bias = torch.nn.Parameter(linear.bias) linear_qs = linear_qs.to(device) assert a.device == b.device assert a.dtype == b.dtype assert torch.equal(a, b) q0 = a.quant_state q1 = b.quant_state for attr in ("code", "dtype", "blocksize", "absmax"): c, d = getattr(q0, attr), getattr(q1, attr) if isinstance(c, torch.Tensor): assert torch.equal(c, d) else: assert c == d, f"{c} != {d}" if q0.state2 is not None: for attr in ("code", "dtype", "blocksize", "absmax"): c, d = getattr(q0.state2, attr), getattr(q1.state2, attr) if isinstance(c, torch.Tensor): assert torch.equal(c, d) else: assert c == d, f"{c} != {d}" if bias: a, b = linear_q.bias, linear_q2.bias assert a.device == b.device assert a.dtype == b.dtype assert torch.equal(a, b) if save_before_forward: bytes_4bit = torch_save_to_buffer(linear_q) # Forward test x = torch.rand(42, layer_shape[0], device=device) a = linear_q(x) b = linear_q2(x) c = linear_qs(x) assert a.device == b.device assert a.dtype == b.dtype assert a.device == c.device assert a.dtype == c.dtype assert torch.equal(a, b) assert torch.equal(a, c) if not save_before_forward: bytes_4bit = torch_save_to_buffer(linear_q) linear_q3 = torch_load_from_buffer(bytes_4bit) # Test moving to CPU and back to GPU if device != "cpu": linear_q2.to("cpu") linear_q2.to(device) d = linear_qs(x) assert c.dtype == d.dtype assert c.device == d.device assert torch.equal(c, d) d = linear_q3(x) assert c.dtype == d.dtype assert c.device == d.device assert torch.equal(c, d) # Saved size ratio test. Target set for layer_shape == (300, 400) w/ bias with TemporaryDirectory() as tmpdir: state_path_4bit = os.path.join(tmpdir, "state_4bit.pth") state_path = os.path.join(tmpdir, "state.pth") torch.save(linear.state_dict(), state_path) torch.save(linear_q.state_dict(), state_path_4bit) size_orig, size_4 = ( os.path.getsize(state_path), os.path.getsize(state_path_4bit), ) size_ratio = size_4 / size_orig target_compression = ( 0.143 if original_dtype == torch.float32 else 0.29 ) # these numbers get lower as weight shape increases ratio_error_msg = ( f"quantized_size {size_4:,} is larger on disk than {target_compression:.2%} of original size {size_orig:,}" ) assert size_ratio < target_compression, ratio_error_msg @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("quant_type", ["nf4", "fp4"]) @pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128]) @pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics")) def test_copy_param(device, quant_type, blocksize, compress_statistics): if device == "hpu" and not is_supported_on_hpu(quant_type): pytest.skip("This configuration is not supported on HPU.") tensor = torch.randn(300, 400) param = bnb.nn.Params4bit( data=tensor, quant_type=quant_type, blocksize=blocksize, compress_statistics=compress_statistics, requires_grad=False, ).to(device) shallow_copy_param = copy.copy(param) assert param.quant_state is shallow_copy_param.quant_state assert param.data.data_ptr() == shallow_copy_param.data.data_ptr() @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("quant_type", ["nf4", "fp4"]) def test_params4bit_torch_chunk_split(device, quant_type): """Test that torch.chunk and torch.split preserve Params4bit subclass for FSDP2 compatibility.""" if device == "hpu" and not is_supported_on_hpu(quant_type, torch.float16, torch.uint8): pytest.skip("This configuration is not supported on HPU.") if device == "cpu": pytest.skip("CPU quantization causes segfault, skipping CPU test") original_tensor = torch.randn(8, 4, dtype=torch.float16, device="cpu") params4bit = bnb.nn.Params4bit(data=original_tensor, quant_type=quant_type, requires_grad=False) if device != "cpu": params4bit = params4bit.to(device) chunks = torch.chunk(params4bit, 2, dim=0) assert isinstance(chunks, tuple), "torch.chunk should return tuple" for chunk in chunks: assert isinstance(chunk, bnb.nn.Params4bit), "Chunk should preserve Params4bit subclass" assert hasattr(chunk, "quant_type"), "Should preserve metadata" assert chunk.quant_type == params4bit.quant_type, "Should preserve quant_type value" splits = torch.split(params4bit, 2, dim=0) assert isinstance(splits, tuple), "torch.split should return tuple" assert len(splits) > 0, "Should have at least one split" for split in splits: assert isinstance(split, bnb.nn.Params4bit), "Split should preserve Params4bit subclass" assert hasattr(split, "quant_type"), "Should preserve metadata" assert split.quant_type == params4bit.quant_type, "Should preserve quant_type value" @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("quant_type", ["nf4", "fp4"]) @pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128]) @pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics")) def test_deepcopy_param(device, quant_type, blocksize, compress_statistics): if device == "hpu" and not is_supported_on_hpu(quant_type): pytest.skip("This configuration is not supported on HPU.") tensor = torch.randn(300, 400) param = bnb.nn.Params4bit( data=tensor, quant_type=quant_type, blocksize=blocksize, compress_statistics=compress_statistics, requires_grad=False, ).to(device) dict_keys_before = set(param.__dict__.keys()) copy_param = copy.deepcopy(param) dict_keys_after = set(param.__dict__.keys()) dict_keys_copy = set(copy_param.__dict__.keys()) assert param.quant_state is not copy_param.quant_state assert param.data.data_ptr() != copy_param.data.data_ptr() # there was a bug where deepcopy would modify the original object assert dict_keys_before == dict_keys_after assert dict_keys_before == dict_keys_copy @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("quant_type", ["nf4", "fp4"]) @pytest.mark.parametrize("blocksize", [64, 128] if not HIP_ENVIRONMENT else [128]) @pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics")) def test_params4bit_real_serialization(device, quant_type, blocksize, compress_statistics): if device == "hpu" and not is_supported_on_hpu(quant_type): pytest.skip("This configuration is not supported on HPU.") original_tensor = torch.randn(300, 400) original_param = bnb.nn.Params4bit( data=original_tensor, quant_type=quant_type, blocksize=blocksize, compress_statistics=compress_statistics, ) dict_keys_before = set(original_param.__dict__.keys()) original_param.to(device) # change device to trigger quantization serialized_param = pickle.dumps(original_param) deserialized_param = pickle.loads(serialized_param) dict_keys_after = set(original_param.__dict__.keys()) dict_keys_deserialized = set(deserialized_param.__dict__.keys()) assert torch.equal(original_param.data, deserialized_param.data) assert original_param.requires_grad == deserialized_param.requires_grad == False assert original_param.quant_type == deserialized_param.quant_type assert original_param.blocksize == deserialized_param.blocksize assert original_param.compress_statistics == deserialized_param.compress_statistics assert original_param.quant_state == deserialized_param.quant_state # there was a bug where deepcopy would modify the original object assert dict_keys_before == dict_keys_after assert dict_keys_before == dict_keys_deserialized @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("quant_type", ["nf4", "fp4"]) @pytest.mark.parametrize("compute_dtype", [torch.bfloat16, torch.float32], ids=describe_dtype) @pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics")) @pytest.mark.parametrize("bias", TRUE_FALSE, ids=id_formatter("bias")) @pytest.mark.parametrize("fullgraph", TRUE_FALSE, ids=id_formatter("fullgraph")) @pytest.mark.parametrize("mode", ["default", "reduce-overhead"], ids=id_formatter("mode")) @pytest.mark.skipif(torch.__version__ < (2, 4), reason="Not supported in torch < 2.4") def test_linear4bit_torch_compile(device, quant_type, compute_dtype, compress_statistics, bias, fullgraph, mode): if device == "hpu" and not is_supported_on_hpu(quant_type): pytest.skip("This configuration is not supported on HPU.") if fullgraph and torch.__version__ < (2, 8, 0, "dev"): pytest.skip("fullgraph mode requires torch 2.8 or higher") if device == "cuda" and platform.system() == "Windows": pytest.skip("Triton is not officially supported on Windows") # Has a strange regression on Linux aarch64 CPU in torch==2.6.0 when fullgraph=False. if ( not fullgraph and device == "cpu" and platform.machine() == "aarch64" and platform.system() == "Linux" and ((2, 7) > torch.__version__ >= (2, 6)) ): pytest.xfail("Regression in torch==2.6.0 on Linux aarch64 CPU") dim = 256 batch_size = 16 torch.compiler.reset() # Create a small network with Linear4bit layers net = torch.nn.Sequential( *[ bnb.nn.Linear4bit( dim, dim, bias=bias, compute_dtype=compute_dtype, compress_statistics=compress_statistics, quant_type=quant_type, ) for _ in range(4) ] ).to(device) # Create input tensor x = torch.randn(batch_size, dim, dtype=compute_dtype, device=device) # Get reference output before compilation with torch.no_grad(): ref_output = net(x) # Compile the model compile_backend = "hpu_backend" if device == "hpu" else "inductor" compiled_net = torch.compile(net, fullgraph=fullgraph, mode=mode, backend=compile_backend) # Get output from compiled model with torch.no_grad(): compiled_output = compiled_net(x) # Check outputs match assert compiled_output.shape == ref_output.shape assert compiled_output.device == ref_output.device assert compiled_output.dtype == ref_output.dtype torch.testing.assert_close(compiled_output, ref_output) # Test with gradients x.requires_grad_(True) y1 = net(x).sum() y1.backward() grad_ref = x.grad.clone() x.grad = None y2 = compiled_net(x).sum() y2.backward() grad_compiled = x.grad.clone() torch.testing.assert_close(grad_compiled, grad_ref) bitsandbytes/tests/test_linear8bitlt.py000066400000000000000000000300701506752665600207470ustar00rootroot00000000000000from contextlib import nullcontext import copy import os import pickle import platform from tempfile import TemporaryDirectory import pytest import torch import bitsandbytes as bnb from bitsandbytes.cextension import HIP_ENVIRONMENT from bitsandbytes.nn.modules import Linear8bitLt from tests.helpers import ( TRUE_FALSE, get_available_devices, id_formatter, torch_load_from_buffer, torch_save_to_buffer, ) # contributed by Alex Borzunov, see: # https://github.com/bigscience-workshop/petals/blob/main/tests/test_linear8bitlt.py @pytest.mark.parametrize("device", get_available_devices()) def test_linear_no_igemmlt(device): linear = torch.nn.Linear(1024, 3072) x = torch.randn(3, 1024, dtype=torch.half) linear_custom = Linear8bitLt( linear.in_features, linear.out_features, linear.bias is not None, has_fp16_weights=False, threshold=6.0, ) # TODO: Remove, this is no longer implemented linear_custom.state.force_no_igemmlt = True linear_custom.weight = bnb.nn.Int8Params( linear.weight.data.clone(), requires_grad=False, has_fp16_weights=False, ).to(linear.weight.dtype) linear_custom.bias = linear.bias linear_custom = linear_custom.to(device) linear = linear.half().to(device) x_ref = x.clone().to(device).requires_grad_(True) x_ours = x.clone().to(device).requires_grad_(True) fx_ref = linear(x_ref).float() grad_proj = torch.randn_like(fx_ref) (fx_ref * grad_proj).mean().backward() fx_ours = linear_custom(x_ours).float() (fx_ours * grad_proj).mean().backward() assert linear_custom.state.CB is not None assert not linear_custom.state.has_fp16_weights idx = torch.isclose(fx_ref, fx_ours, atol=0.02, rtol=1e-5) assert (idx == 0).sum().item() < fx_ref.numel() * 2.5e-4 torch.testing.assert_close(fx_ref, fx_ours, atol=0.03, rtol=1e-5) torch.testing.assert_close(x_ref.grad, x_ours.grad, atol=0.01, rtol=1e-5) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("has_fp16_weights", TRUE_FALSE, ids=id_formatter("has_fp16_weights")) @pytest.mark.parametrize("threshold", [0.0, 6.0], ids=id_formatter("threshold")) @pytest.mark.parametrize("serialize_before_forward", TRUE_FALSE, ids=id_formatter("serialize_before_forward")) @pytest.mark.parametrize("deserialize_before_cuda", TRUE_FALSE, ids=id_formatter("deserialize_before_cuda")) @pytest.mark.parametrize("save_before_forward", TRUE_FALSE, ids=id_formatter("save_before_forward")) @pytest.mark.parametrize("load_before_cuda", TRUE_FALSE, ids=id_formatter("load_before_cuda")) def test_linear_serialization( device, has_fp16_weights, threshold, serialize_before_forward, deserialize_before_cuda, save_before_forward, load_before_cuda, ): if device != "cuda" and has_fp16_weights: pytest.skip("has_fp16_weights is only supported on CUDA and is deprecated") linear = torch.nn.Linear(32, 96) # TODO: Fallback for bad shapes x = torch.randn(4, 32, dtype=torch.half) # x = torch.randn(3, 32, dtype=torch.half) linear_custom = Linear8bitLt( linear.in_features, linear.out_features, linear.bias is not None, has_fp16_weights=has_fp16_weights, threshold=threshold, ) linear_custom.weight = bnb.nn.Int8Params( linear.weight.data.clone(), requires_grad=has_fp16_weights, has_fp16_weights=has_fp16_weights, ) linear_custom.bias = linear.bias linear_custom = linear_custom.to(device) if serialize_before_forward: state_dict_8bit = linear_custom.state_dict() if save_before_forward: bytes_8bit = torch_save_to_buffer(linear_custom) x_first = x.clone().to(device).requires_grad_(True) fx_first = linear_custom(x_first).float() grad_proj = torch.randn_like(fx_first) (fx_first * grad_proj).mean().backward() if not serialize_before_forward: state_dict_8bit = linear_custom.state_dict() if not save_before_forward: bytes_8bit = torch_save_to_buffer(linear_custom) with TemporaryDirectory() as tmpdir: state_path_8bit = os.path.join(tmpdir, "state_8bit.pth") state_path = os.path.join(tmpdir, "state.pth") torch.save(linear.state_dict(), state_path) torch.save(state_dict_8bit, state_path_8bit) if not has_fp16_weights: assert os.path.getsize(state_path_8bit) < 0.5 * os.path.getsize(state_path) new_state_dict = torch.load(state_path_8bit, weights_only=False) new_linear_custom = Linear8bitLt( linear.in_features, linear.out_features, linear.bias is not None, has_fp16_weights=has_fp16_weights, threshold=threshold, ) if deserialize_before_cuda: with nullcontext() if has_fp16_weights else pytest.raises(RuntimeError): new_linear_custom.load_state_dict(new_state_dict, strict=True) if load_before_cuda: new_linear_custom2 = torch_load_from_buffer(bytes_8bit) new_linear_custom = new_linear_custom.to(device) if not deserialize_before_cuda: new_linear_custom.load_state_dict(new_state_dict, strict=True) if not load_before_cuda: new_linear_custom2 = torch_load_from_buffer(bytes_8bit) x_second = x.clone().to(device).requires_grad_(True) fx_second = new_linear_custom(x_second).float() (fx_second * grad_proj).mean().backward() x_third = x.clone().to(device).requires_grad_(True) fx_third = new_linear_custom2(x_third).float() (fx_third * grad_proj).mean().backward() # if 8-bit weights were loaded before .cuda, state is incorrect anyway and RuntimeError was raised if has_fp16_weights or not deserialize_before_cuda: assert torch.allclose(fx_first, fx_second, atol=1e-5) assert torch.allclose(x_first.grad, x_second.grad, atol=1e-5) assert torch.allclose(fx_first, fx_third, atol=1e-5) assert torch.allclose(x_first.grad, x_third.grad, atol=1e-5) @pytest.fixture def linear8bit(requires_cuda): linear = torch.nn.Linear(32, 96) linear_custom = Linear8bitLt( linear.in_features, linear.out_features, linear.bias is not None, has_fp16_weights=False, threshold=6.0, ) linear_custom.weight = bnb.nn.Int8Params( linear.weight.data.clone(), requires_grad=False, has_fp16_weights=False, ) linear_custom.bias = linear.bias linear_custom = linear_custom.cuda() return linear_custom def test_linear8bit_copy_param(linear8bit): shallow_copy = copy.copy(linear8bit) assert linear8bit.weight is shallow_copy.weight assert linear8bit.bias is shallow_copy.bias assert linear8bit.weight.data.data_ptr() == shallow_copy.weight.data.data_ptr() def test_linear8bit_deepcopy_param(linear8bit): deep_copy = copy.deepcopy(linear8bit) assert linear8bit.weight is not deep_copy.weight assert linear8bit.bias is not deep_copy.bias assert linear8bit.weight.data.data_ptr() != deep_copy.weight.data.data_ptr() assert torch.allclose(linear8bit.weight.data, deep_copy.weight.data) assert linear8bit.state == deep_copy.state # check for a bug where SCB and CB were not copied assert deep_copy.weight.SCB is not None assert (linear8bit.weight.SCB == deep_copy.weight.SCB).all() assert deep_copy.weight.CB is not None assert (linear8bit.weight.CB == deep_copy.weight.CB).all() def test_linear8bit_serialization(linear8bit): serialized = pickle.dumps(linear8bit) deserialized = pickle.loads(serialized) assert linear8bit.weight.data.data_ptr() != deserialized.weight.data.data_ptr() assert torch.allclose(linear8bit.weight.data, deserialized.weight.data) assert linear8bit.bias.data.data_ptr() != deserialized.bias.data.data_ptr() assert torch.allclose(linear8bit.bias.data, deserialized.bias.data) assert linear8bit.state == deserialized.state # check for a bug where SCB and CB were not copied assert (linear8bit.weight.SCB == deserialized.weight.SCB).all() assert (linear8bit.weight.CB == deserialized.weight.CB).all() @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("threshold", [0.0, 6.0], ids=id_formatter("threshold")) @pytest.mark.parametrize("bias", TRUE_FALSE, ids=id_formatter("bias")) @pytest.mark.parametrize("fullgraph", TRUE_FALSE, ids=id_formatter("fullgraph")) @pytest.mark.parametrize("mode", ["default", "reduce-overhead"], ids=id_formatter("mode")) @pytest.mark.skipif(torch.__version__ < (2, 4), reason="Not supported in torch < 2.4") @pytest.mark.skipif(HIP_ENVIRONMENT, reason="this test is not supported on ROCm yet") def test_linear8bitlt_torch_compile(device, threshold, bias, fullgraph, mode): if device == "cuda" and platform.system() == "Windows": pytest.skip("Triton is not officially supported on Windows") dim = 256 batch_size = 16 torch.compiler.reset() # Create a small network with Linear8bitLt layers net = torch.nn.Sequential( *[bnb.nn.Linear8bitLt(dim, dim, bias=bias, has_fp16_weights=False, threshold=threshold) for _ in range(4)] ).to(device) dynamic_output_shapes = fullgraph and threshold > 0 with torch._dynamo.config.patch("capture_dynamic_output_shape_ops", dynamic_output_shapes): # Create input tensor x = torch.randn(batch_size, dim, dtype=torch.float16, device=device) # Get reference output before compilation with torch.no_grad(): ref_output = net(x) # Compile the model compile_backend = "hpu_backend" if device == "hpu" else "inductor" compiled_net = torch.compile(net, fullgraph=fullgraph, mode=mode, backend=compile_backend) # Get output from compiled model with torch.no_grad(): compiled_output = compiled_net(x) # Check outputs match assert compiled_output.shape == ref_output.shape assert compiled_output.device == ref_output.device assert compiled_output.dtype == ref_output.dtype torch.testing.assert_close(compiled_output, ref_output) # Test with gradients. Currently only works with threshold=0. # Has a strange regression on Linux aarch64 CPU in torch==2.6.0. is_broken_platform = ( device == "cpu" and platform.system() == "Linux" and platform.machine() == "aarch64" and (2, 6) <= torch.__version__ < (2, 7) ) if threshold == 0 and not is_broken_platform: x.requires_grad_(True) y1 = net(x).sum() y1.backward() grad_ref = x.grad.clone() x.grad = None y2 = compiled_net(x).sum() y2.backward() grad_compiled = x.grad.clone() torch.testing.assert_close(grad_compiled, grad_ref) @pytest.mark.parametrize("device", get_available_devices(no_cpu=True)) @pytest.mark.skipif(not get_available_devices(no_cpu=True), reason="No accelerator device") def test_linear8bitlt_device_movement(device): """Test moving a Linear8bitLt layer between CPU and an accelerator device.""" # Create a Linear8bitLt layer on CPU layer = bnb.nn.Linear8bitLt(32, 128, bias=False, has_fp16_weights=False) torch.nn.init.xavier_uniform_(layer.weight) # Create a sample input. x = torch.randn(4, 32, dtype=torch.float16, device="cpu") # Move to the device. This should quantize the weights. layer = layer.to(device) assert layer.weight.data.dtype == torch.int8 # Call the layer on the accelerator device. out_accelerator = layer(x.to(device)) # Move back to CPU and call again. layer = layer.to("cpu") out_cpu = layer(x) # Move back to the accelerator device and call again. layer = layer.to(device) out_accelerator_2 = layer(x.to(device)) # Move back to the CPU and call one last time. layer = layer.to("cpu") out_cpu_2 = layer(x) # CPU outputs should match both times. torch.testing.assert_close(out_cpu_2, out_cpu, rtol=1e-8, atol=1e-8) # Accelerator outputs should match both times. torch.testing.assert_close(out_accelerator_2, out_accelerator, rtol=1e-8, atol=1e-8) bitsandbytes/tests/test_modules.py000066400000000000000000000443431506752665600200260ustar00rootroot00000000000000import inspect import pytest import torch from torch import nn import bitsandbytes as bnb from tests.helpers import get_available_devices, id_formatter, is_supported_on_hpu class MockArgs: def __init__(self, initial_data): for key in initial_data: setattr(self, key, initial_data[key]) class MLP8bit(torch.nn.Module): def __init__(self, dim1, dim2, has_fp16_weights=True, threshold=0.0): super().__init__() self.fc1 = bnb.nn.Linear8bitLt( dim1, dim2, has_fp16_weights=has_fp16_weights, threshold=threshold, ) self.fc2 = bnb.nn.Linear8bitLt( dim2, dim1, has_fp16_weights=has_fp16_weights, threshold=threshold, ) def forward(self, x): x = self.fc1(x) x = self.fc2(x) return x def get_args(): args = MockArgs([]) args.quant_type = "vector" args.use_8bit_training = "full" args.clip_freq = 9999 return args def assert_all_approx_close(a, b, atol=1e-8, rtol=1e-5, count=10): idx = torch.isclose(a, b, rtol=rtol, atol=atol) sumval = (idx == 0).sum().item() if sumval > count: print(f"Too many values not close: assert {sumval} < {count}") torch.testing.assert_close(a, b, rtol=rtol, atol=atol) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("threshold", [0.0, 3.0], ids=id_formatter("threshold")) def test_linear8bitlt_inference(device, threshold): l1 = bnb.nn.Linear8bitLt(32, 64, threshold=threshold, has_fp16_weights=False).to(device).half() assert l1.weight.device.type == device assert l1.weight.dtype == torch.int8 l1.eval() for i in range(100): b1 = torch.randn(16, 8, 32, device=device).half() o1 = l1(b1) if i == 1: assert l1.state.CB is not None # TODO: Remove support for training int8 weights @pytest.mark.parametrize("device", get_available_devices()) def test_linear8bitlt_accumulated_gradient(device): if device != "cuda": pytest.skip("Only supported on CUDA") l1 = torch.nn.Sequential(*[bnb.nn.Linear8bitLt(32, 32).to(device).half() for i in range(2)]) l2 = torch.nn.Sequential(*[torch.nn.Linear(32, 32).to(device).half() for i in range(2)]) l1[0].weight.data.copy_(l2[0].weight.data) l1[1].weight.data.copy_(l2[1].weight.data) l1[0].bias.data.copy_(l2[0].bias.data) l1[1].bias.data.copy_(l2[1].bias.data) opt1 = bnb.optim.Adam32bit(l1.parameters(), lr=0.001) opt2 = bnb.optim.Adam32bit(l2.parameters(), lr=0.001) acc_steps = 10 for i in range(15): b1 = torch.randn(16, 8, 32, device=device).half() o1 = l1(b1) o2 = l2(b1) loss1 = o1.mean() loss2 = o2.mean() loss1.backward() loss2.backward() if i == 2: assert l1[0].state.CB is not None assert l1[1].state.CB is not None if i > 0 and i % acc_steps == 0: opt1.step() opt1.zero_grad(True) opt2.step() opt2.zero_grad(True) assert_all_approx_close(l1[0].weight, l2[0].weight, rtol=1.05, atol=0.01, count=2) assert_all_approx_close(l1[1].weight, l2[1].weight, rtol=1.05, atol=0.01, count=2) # we do this copy because otherwise we have small divergences over time that add up l1[0].weight.data.copy_(l2[0].weight.data) l1[1].weight.data.copy_(l2[1].weight.data) l1[0].bias.data.copy_(l2[0].bias.data) l1[1].bias.data.copy_(l2[1].bias.data) else: assert_all_approx_close(l1[0].weight.grad, l2[0].weight.grad, rtol=1.05, atol=0.04, count=1) assert_all_approx_close(l1[1].weight.grad, l2[1].weight.grad, rtol=1.05, atol=0.04, count=1) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("threshold", [0.0, 2.0]) def test_linear8bitlt_no_fp16_weights(device, threshold): l1 = ( bnb.nn.Linear8bitLt( 32, 64, threshold=threshold, has_fp16_weights=False, ) .to(device) .half() ) assert l1.weight.dtype == torch.int8 l1.eval() for i in range(4): b1 = torch.randn(16, 8, 32, device=device, dtype=torch.float16) o1 = l1(b1) assert o1.dtype == torch.float16 mlp = MLP8bit(32, 64, threshold=threshold, has_fp16_weights=False).to(device) assert mlp.fc1.weight.dtype == torch.int8 assert mlp.fc2.weight.dtype == torch.int8 for i in range(4): b1 = torch.randn(16, 8, 32, device=device, dtype=torch.float16) o1 = mlp(b1) assert o1.dtype == torch.float16 if threshold > 0 and device not in ("cpu", "xpu"): assert mlp.fc1.state.idx is not None assert mlp.fc2.state.idx is not None mlp = MLP8bit(32, 64, threshold=threshold, has_fp16_weights=False).to(device).half() assert mlp.fc1.weight.dtype == torch.int8 assert mlp.fc2.weight.dtype == torch.int8 for i in range(4): b1 = torch.randn(16, 8, 32, device=device, dtype=torch.float16) o1 = mlp(b1) assert o1.dtype == torch.float16 if threshold > 0 and device not in ("cpu", "xpu"): assert mlp.fc1.state.idx is not None assert mlp.fc2.state.idx is not None mlp = MLP8bit(32, 64, threshold=threshold, has_fp16_weights=False).half().to(device) for i in range(4): b1 = torch.randn(16, 8, 32, device=device, dtype=torch.float16) o1 = mlp(b1) assert o1.dtype == torch.float16 if threshold > 0 and device not in ("cpu", "xpu"): assert mlp.fc1.state.idx is not None assert mlp.fc2.state.idx is not None assert mlp.fc1.weight.dtype == torch.int8 assert mlp.fc2.weight.dtype == torch.int8 mlp = ( MLP8bit( 32, 64, threshold=threshold, has_fp16_weights=False, ) .half() .to(device) ) for i in range(4): b1 = torch.randn(16, 8, 32, device=device, dtype=torch.float16) o1 = mlp(b1) assert o1.dtype == torch.float16 if threshold > 0 and device not in ("cpu", "xpu"): assert mlp.fc1.state.idx is not None assert mlp.fc2.state.idx is not None assert mlp.fc1.weight.dtype == torch.int8 assert mlp.fc2.weight.dtype == torch.int8 assert mlp.fc1.weight.device.type == device assert mlp.fc2.weight.device.type == device mlp = MLP8bit( 32, 64, threshold=threshold, has_fp16_weights=False, ) w1, w2 = mlp.fc1.weight.clone().to(device), mlp.fc2.weight.clone().to(device) # grab weights before quantization, mlp = mlp.to(device).half() # and this line triggers quantization for i in range(4): b1 = torch.randn(16, 8, 32, device=device, dtype=torch.float16) o1 = mlp(b1) assert o1.dtype == torch.float16 if threshold > 0 and device not in ("cpu", "xpu"): assert mlp.fc1.state.idx is not None assert mlp.fc2.state.idx is not None assert mlp.fc1.weight.dtype == torch.int8 assert mlp.fc2.weight.dtype == torch.int8 assert mlp.fc1.weight.device.type == device assert mlp.fc2.weight.device.type == device b1 = torch.randn(16, 8, 32, device=device, requires_grad=True, dtype=torch.half) o1 = mlp(b1) assert o1.dtype == torch.float16 assert o1.requires_grad grad_proj = torch.randn_like(o1) mlp.zero_grad() (o1 * grad_proj).sum().backward() grad_ref = grad_proj.flatten(2) @ w2.half() @ w1.half() scale = grad_ref.abs().mean() torch.testing.assert_close(b1.grad, grad_ref, rtol=0, atol=0.05 * scale) idx = torch.isclose(b1.grad, grad_ref, atol=0.01 * scale, rtol=0.1) assert (idx == 0).sum().item() <= b1.numel() * 0.005 @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize( "module", [ lambda n_in, n_out, bias=True: bnb.nn.Linear8bitLt(n_in, n_out, bias=bias, has_fp16_weights=False), bnb.nn.LinearNF4, ], ids=["Int8Lt", "NF4"], ) def test_linear_kbit_fp32_bias(device, module): # casts model to fp16 -> int8 automatically l1 = module(32, 64).to(device) assert l1.weight.dtype in [torch.int8, torch.uint8] assert l1.bias.dtype == torch.float32 for i in range(100): b1 = torch.randn(16, 8, 32, device=device, dtype=torch.float16) # casts bias to fp32 o1 = l1(b1) assert l1.bias.dtype == torch.float16 # casts model to fp16 -> int8 automatically l1 = module(32, 64, bias=False).to(device) assert l1.weight.dtype in [torch.int8, torch.uint8] assert l1.bias is None for i in range(100): b1 = torch.randn(16, 8, 32, device=device, dtype=torch.float16) o1 = l1(b1) assert l1.bias is None module_dict = { "Int8Lt": bnb.nn.Linear8bitLt, "4bit": bnb.nn.Linear4bit, "FP4": bnb.nn.LinearFP4, "NF4": bnb.nn.LinearNF4, "FP4+C": lambda d1, d2: bnb.nn.LinearFP4(d1, d2, compress_statistics=True), "NF4+C": lambda d1, d2: bnb.nn.LinearNF4(d1, d2, compress_statistics=True), "NF4+fp32": lambda d1, d2: bnb.nn.LinearNF4(d1, d2, compute_dtype=torch.float32), "NF4+fp16": lambda d1, d2: bnb.nn.LinearNF4(d1, d2, compute_dtype=torch.float16), "NF4+bf16": lambda d1, d2: bnb.nn.LinearNF4(d1, d2, compute_dtype=torch.bfloat16), } @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("module", module_dict.values(), ids=module_dict.keys()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_kbit_backprop(device, module, dtype): b = 16 dim1 = 36 dim2 = 84 # dim1 = 37 # dim2 = 83 ref = nn.Sequential(*[torch.nn.Linear(dim1, dim2), torch.nn.Linear(dim2, 128)]) torch.nn.init.kaiming_normal_(ref[0].weight) torch.nn.init.kaiming_normal_(ref[1].weight) ref[1].weight.requires_grad_(False) kbit = nn.Sequential(*[torch.nn.Linear(dim1, dim2), module(dim2, 128)]) if ( device == "hpu" and isinstance(kbit[1], bnb.nn.Linear4bit) and not is_supported_on_hpu(kbit[1].weight.quant_type, dtype) ): pytest.skip("This configuration not supported on HPU") kbit[0].weight.detach().copy_(ref[0].weight) kbit[1].weight.detach().copy_(ref[1].weight) kbit[0].bias.detach().copy_(ref[0].bias) kbit[1].bias.detach().copy_(ref[1].bias) kbit[1].weight.requires_grad_(False) ref = ref.to(device=device, dtype=dtype) kbit = kbit.to(device=device, dtype=dtype) kbit = kbit.to(device=device, dtype=dtype) errs1 = [] errs2 = [] relerrs1 = [] relerrs2 = [] for i in range(100): batch = torch.randn(b, dim1, device=device, dtype=dtype) out1 = ref(batch) out2 = kbit(batch) out1.mean().backward() out2.mean().backward() grad1 = ref[0].weight.grad grad2 = kbit[0].weight.grad bgrad1 = ref[0].bias.grad bgrad2 = kbit[0].bias.grad err1 = (out1 - out2).abs().float() err2 = (grad1 - grad2).abs().float() relerr1 = err1 / (out1.abs().float() + 1e-9) relerr2 = err2 / (grad1.abs().float() + 1e-9) errs1.append(err1.mean().item()) errs2.append(err2.mean().item()) relerrs1.append(relerr1.mean().item()) relerrs2.append(relerr2.mean().item()) if isinstance(module, bnb.nn.Linear8bitLt): assert_all_approx_close(grad1, grad2, atol=0.008, rtol=0.05, count=1) torch.testing.assert_close(bgrad1, bgrad2, atol=0.008, rtol=0.05) else: assert_all_approx_close(grad1, grad2, atol=0.015, rtol=0.05, count=1) torch.testing.assert_close(bgrad1, bgrad2, atol=0.02, rtol=0.05) ref.zero_grad() kbit.zero_grad() assert kbit[0].weight.grad is None or kbit[0].weight.grad.sum().item() == 0 assert kbit[0].weight.grad is None or kbit[0].bias.grad.sum().item() == 0 @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("embedding_dim", [64, 65]) @pytest.mark.parametrize("input_shape", [(10,), (10, 10), (10, 10, 10)], ids=str) @pytest.mark.parametrize( "embedding_class,quant_storage", [ (bnb.nn.Embedding8bit, None), (bnb.nn.EmbeddingFP4, torch.uint8), (bnb.nn.EmbeddingFP4, torch.float32), (bnb.nn.EmbeddingNF4, torch.uint8), (bnb.nn.EmbeddingNF4, torch.float32), ], ids=lambda x: x.__name__ if inspect.isclass(x) else str(x), ) def test_embedding_lossless(device, embedding_class, input_shape, embedding_dim, quant_storage): if device == "hpu": if embedding_class is bnb.nn.EmbeddingFP4: pytest.skip("FP4 is not supported on HPU") elif embedding_class is bnb.nn.EmbeddingNF4 and not is_supported_on_hpu("nf4", torch.float32, quant_storage): pytest.skip("This configuration is not supported on HPU") num_embeddings = 128 src_weight = (torch.randn((num_embeddings, embedding_dim), dtype=torch.float32) > 0).to( torch.float32 ) * 2 - 1 # Embeddings filled with {-1, 1} values. It should compress losslessly emb_base = nn.Embedding( num_embeddings=num_embeddings, embedding_dim=embedding_dim, _freeze=True, _weight=src_weight, ) if embedding_class is bnb.nn.Embedding8bit: e = embedding_class(num_embeddings=num_embeddings, embedding_dim=embedding_dim) else: e = embedding_class(num_embeddings=num_embeddings, embedding_dim=embedding_dim, quant_storage=quant_storage) e.load_state_dict(emb_base.state_dict()) emb_base.to(device) e.to(device) input_tokens = torch.randint(low=0, high=num_embeddings, size=input_shape, device=device) torch.testing.assert_close( actual=e(input_tokens), expected=emb_base(input_tokens), ) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("embedding_dim", [64, 65]) @pytest.mark.parametrize("input_shape", [(10,), (10, 10), (10, 10, 10)], ids=str) @pytest.mark.parametrize( "embedding_class,quant_storage", [ (bnb.nn.Embedding8bit, None), (bnb.nn.EmbeddingFP4, torch.uint8), (bnb.nn.EmbeddingFP4, torch.float32), (bnb.nn.EmbeddingNF4, torch.uint8), (bnb.nn.EmbeddingNF4, torch.float32), ], ids=lambda x: x.__name__ if inspect.isclass(x) else str(x), ) def test_embedding_error(device, embedding_class, input_shape, embedding_dim, quant_storage): if device == "hpu": if embedding_class is bnb.nn.EmbeddingFP4: pytest.skip("FP4 is not supported on HPU") elif embedding_class is bnb.nn.EmbeddingNF4 and not is_supported_on_hpu("nf4", torch.float32, quant_storage): pytest.skip("This configuration is not supported on HPU") is_8bit = embedding_class is bnb.nn.Embedding8bit num_embeddings = 128 src_weight = torch.rand((num_embeddings, embedding_dim), dtype=torch.float32) emb_base = nn.Embedding( num_embeddings=num_embeddings, embedding_dim=embedding_dim, _freeze=True, _weight=src_weight, ) if is_8bit: e = embedding_class(num_embeddings=num_embeddings, embedding_dim=embedding_dim) else: e = embedding_class(num_embeddings=num_embeddings, embedding_dim=embedding_dim, quant_storage=quant_storage) e.load_state_dict(emb_base.state_dict()) emb_base.to(device) e.to(device) input_tokens = torch.randint(low=0, high=num_embeddings, size=input_shape, device=device) torch.testing.assert_close( actual=e(input_tokens), expected=emb_base(input_tokens), atol=0.05 if is_8bit else 0.20, rtol=0.0, ) @pytest.mark.parametrize("device", get_available_devices()) def test_4bit_linear_warnings(device): dim1 = 64 with pytest.warns(UserWarning, match=r"inference or training"): net = nn.Sequential(*[bnb.nn.Linear4bit(dim1, dim1, quant_type="nf4") for i in range(10)]) net = net.to(device) inp = torch.rand(10, dim1, device=device, dtype=torch.float16) net(inp) with pytest.warns(UserWarning, match=r"inference."): net = nn.Sequential(*[bnb.nn.Linear4bit(dim1, dim1, quant_type="nf4") for i in range(10)]) net = net.to(device) inp = torch.rand(1, dim1, device=device, dtype=torch.float16) net(inp) with pytest.warns(UserWarning) as record: net = nn.Sequential(*[bnb.nn.Linear4bit(dim1, dim1, quant_type="nf4") for i in range(10)]) net = net.to(device) inp = torch.rand(10, dim1, device=device, dtype=torch.float16) net(inp) net = nn.Sequential(*[bnb.nn.Linear4bit(dim1, dim1, quant_type="nf4") for i in range(10)]) net = net.to(device) inp = torch.rand(1, dim1, device=device, dtype=torch.float16) net(inp) assert len(record) == 2 @pytest.mark.parametrize("device", get_available_devices()) def test_4bit_embedding_warnings(device): num_embeddings = 128 default_block_size = 64 with pytest.warns(UserWarning, match=r"inference."): net = bnb.nn.Embedding4bit( num_embeddings=num_embeddings, embedding_dim=default_block_size + 1, quant_type="nf4" ) net.to(device) inp = torch.randint(low=0, high=num_embeddings, size=(1,), device=device) net(inp) def test_4bit_embedding_weight_fsdp_fix(requires_cuda): num_embeddings = 64 embedding_dim = 32 module = bnb.nn.Embedding4bit(num_embeddings=num_embeddings, embedding_dim=embedding_dim) module.cuda() module.weight.quant_state = None input_tokens = torch.randint(low=0, high=num_embeddings, size=(1,), device="cuda") module(input_tokens) assert module.weight.quant_state is not None def test_4bit_linear_weight_fsdp_fix(requires_cuda): inp_size = 64 out_size = 32 module = bnb.nn.Linear4bit(inp_size, out_size) module.cuda() module.weight.quant_state = None input_tensor = torch.randn((1, inp_size), device="cuda") module(input_tensor) assert module.weight.quant_state is not None def test_embedding_not_implemented_error(): with pytest.raises(NotImplementedError): emb = bnb.nn.Embedding4bit(32, 32) emb.state_dict() with pytest.raises(NotImplementedError): emb = bnb.nn.Embedding8bit(32, 32) emb.state_dict() bitsandbytes/tests/test_ops.py000066400000000000000000000252411506752665600171530ustar00rootroot00000000000000from math import prod import pytest import torch import bitsandbytes from bitsandbytes.cextension import HIP_ENVIRONMENT from tests.helpers import TRUE_FALSE, get_available_devices, id_formatter, is_supported_on_hpu # torch.library.opcheck is only available in torch 2.4 and later. # When testing with older versions, we will skip it as a no-op. if torch.__version__ >= (2, 4): opcheck = torch.library.opcheck else: opcheck = lambda *args, **kwargs: None class TestLLMInt8Ops: @pytest.mark.parametrize("device", get_available_devices()) def test_int8_linear_matmul(self, device): A = torch.randint(-128, 127, (10, 20), dtype=torch.int8, device=device) B = torch.randint(-128, 127, (30, 20), dtype=torch.int8, device=device) out = torch.ops.bitsandbytes.int8_linear_matmul.default(A, B) assert out.shape == (10, 30) assert out.dtype == torch.int32 assert out.device == A.device opcheck(torch.ops.bitsandbytes.int8_linear_matmul.default, (A, B)) @pytest.mark.parametrize("device", get_available_devices()) def test_int8_linear_matmul_out(self, device): A = torch.randint(-128, 127, (10, 20), dtype=torch.int8, device=device) B = torch.randint(-128, 127, (30, 20), dtype=torch.int8, device=device) out = torch.empty((10, 30), dtype=torch.int32, device=device) torch.ops.bitsandbytes.int8_linear_matmul.out(A, B, out) assert out.shape == (10, 30) assert out.dtype == torch.int32 assert out.device == A.device opcheck(torch.ops.bitsandbytes.int8_linear_matmul.out, (A, B, out)) @pytest.mark.parametrize("threshold", [0.0, 6.0]) @pytest.mark.parametrize("device", get_available_devices()) def test_int8_vectorwise_quant(self, threshold, device): A = torch.randn(10, 20, dtype=torch.float16, device=device) A[1][0] = 1000.0 out_row, row_stats, outlier_cols = torch.ops.bitsandbytes.int8_vectorwise_quant(A, threshold=threshold) assert out_row.shape == (10, 20) assert out_row.dtype == torch.int8 assert out_row.device == A.device assert row_stats.shape == (10,) assert row_stats.dtype == torch.float32 assert row_stats.device == A.device if threshold > 0.0: assert outlier_cols is not None assert outlier_cols.dim() == 1 assert outlier_cols.shape[0] <= A.shape[1] assert outlier_cols.device == A.device else: assert outlier_cols is None opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A,)) opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A, threshold)) @pytest.mark.parametrize("device", get_available_devices()) def test_int8_mm_dequant(self, device): A = torch.randint(-128, 127, (256, 256), dtype=torch.int32, device=device) row_stats = torch.randn(256, dtype=torch.float32, device=device) col_stats = torch.randn(256, dtype=torch.float32, device=device) out = torch.ops.bitsandbytes.int8_mm_dequant(A, row_stats, col_stats) assert out.shape == A.shape assert out.dtype == torch.float16 assert out.device == A.device opcheck(torch.ops.bitsandbytes.int8_mm_dequant, (A, row_stats, col_stats)) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) @pytest.mark.parametrize("has_bias", TRUE_FALSE) def test_int8_scaled_mm(self, device, dtype, has_bias): A = torch.randint(-128, 127, (10, 20), dtype=torch.int8, device=device) B = torch.randint(-128, 127, (30, 20), dtype=torch.int8, device=device) row_stats = torch.randn(10, dtype=torch.float32, device=device) col_stats = torch.randn(30, dtype=torch.float32, device=device) bias = torch.randn(30, dtype=dtype, device=device) if has_bias else None out = torch.ops.bitsandbytes.int8_scaled_mm(A, B, row_stats, col_stats, bias=bias, dtype=dtype) assert out.shape == (10, 30) assert out.dtype == dtype assert out.device == A.device opcheck(torch.ops.bitsandbytes.int8_scaled_mm, (A, B, row_stats, col_stats, bias, dtype)) class TestInt8BlockwiseQuantOps: @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) @pytest.mark.parametrize("blocksize", [64, 128, 256, 512] if not HIP_ENVIRONMENT else [128, 256, 512]) def test_quantize_blockwise(self, device, dtype, blocksize): if device == "cpu": if dtype != torch.float32: pytest.skip("CPU implementation is only available for float32") if blocksize != 256: pytest.skip("CPU implementation is slow; only test blocksize=256") code = bitsandbytes.functional.create_dynamic_map().to(device) A = torch.randn(1024, 1024, dtype=dtype, device=device) out, absmax = torch.ops.bitsandbytes.quantize_blockwise(A, code, blocksize) assert out.shape == A.shape assert out.dtype == torch.uint8 assert out.device == A.device assert absmax.device == A.device assert absmax.dtype == torch.float32 opcheck(torch.ops.bitsandbytes.quantize_blockwise, (A, code, blocksize)) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) @pytest.mark.parametrize("blocksize", [64, 128, 256, 512] if not HIP_ENVIRONMENT else [128, 256, 512]) def test_dequantize_blockwise(self, device, dtype, blocksize): if device == "cpu" and dtype != torch.float32: pytest.skip("CPU implementation is only available for float32") A = torch.randint(0, 255, (1024, 1024), dtype=torch.uint8, device=device) code = bitsandbytes.functional.create_dynamic_map().to(device, dtype=torch.float32) n = A.numel() blocks = -(n // -blocksize) absmax = torch.randn((blocks,), device=device, dtype=torch.float32) out = torch.ops.bitsandbytes.dequantize_blockwise.default(A, absmax, code, blocksize, dtype) assert out.shape == A.shape assert out.dtype == dtype assert out.device == A.device opcheck(torch.ops.bitsandbytes.dequantize_blockwise.default, (A, absmax, code, blocksize, dtype)) class Test4bitBlockwiseQuantOps: @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) @pytest.mark.parametrize("storage_dtype", [torch.uint8, torch.bfloat16], ids=id_formatter("storage_dtype")) @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) @pytest.mark.parametrize("blocksize", [64, 128, 256, 512] if not HIP_ENVIRONMENT else [128, 256, 512]) def test_quantize_4bit(self, device, dtype, storage_dtype, quant_type, blocksize): if device == "hpu" and not is_supported_on_hpu(quant_type, dtype, storage_dtype): pytest.skip("This configuration is not supported on HPU.") A = torch.randn(1024, 1024, dtype=dtype, device=device) out, absmax = torch.ops.bitsandbytes.quantize_4bit.default(A, blocksize, quant_type, storage_dtype) assert out.device == A.device assert out.dtype == storage_dtype assert absmax.device == A.device assert absmax.dtype == torch.float32 if storage_dtype != torch.uint8: pytest.xfail("opcheck fails for storage_dtype != torch.uint8") opcheck(torch.ops.bitsandbytes.quantize_4bit.default, (A, blocksize, quant_type, storage_dtype)) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) @pytest.mark.parametrize("storage_dtype", [torch.uint8, torch.bfloat16], ids=id_formatter("storage_dtype")) @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) @pytest.mark.parametrize("blocksize", [64, 128, 256, 512] if not HIP_ENVIRONMENT else [128, 256, 512]) def test_dequantize_4bit(self, device, dtype, storage_dtype, quant_type, blocksize): if device == "hpu" and not is_supported_on_hpu(quant_type, dtype, storage_dtype): pytest.skip("This configuration is not supported on HPU.") shape = (128, 128) n = prod(shape) blocks = -(n // -blocksize) quantized_shape = ((n + 1) // (storage_dtype.itemsize * 2), 1) A = ( torch.randint(0, 255, ((n + 1) // 2,), dtype=torch.uint8, device=device) .view(storage_dtype) .reshape(quantized_shape) .contiguous() ) absmax = torch.randn((blocks,), dtype=torch.float32, device=device) out = torch.ops.bitsandbytes.dequantize_4bit.default(A, absmax, blocksize, quant_type, shape, dtype) assert out.device == A.device assert out.shape == shape opcheck( torch.ops.bitsandbytes.dequantize_4bit.default, (A, absmax, blocksize, quant_type, shape, dtype), ) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) @pytest.mark.parametrize("storage_dtype", [torch.uint8, torch.bfloat16], ids=id_formatter("storage_dtype")) @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) @pytest.mark.parametrize("blocksize", [64, 128, 256, 512] if not HIP_ENVIRONMENT else [128, 256, 512]) @pytest.mark.skipif(HIP_ENVIRONMENT, reason="this test is not supported on ROCm yet") def test_gemv_4bit(self, device, dtype, storage_dtype, quant_type, blocksize): if device == "hpu" and not is_supported_on_hpu(quant_type, dtype, storage_dtype): pytest.skip("This configuration is not supported on HPU.") out_features = 1024 in_features = 256 A = torch.randn((1, 1, in_features), dtype=dtype, device=device) B = torch.randn((out_features, in_features), dtype=dtype, device=A.device) B_q, absmax = torch.ops.bitsandbytes.quantize_4bit(B, blocksize, quant_type, storage_dtype) code = bitsandbytes.functional.get_4bit_type(quant_type, device=A.device, blocksize=blocksize) out = torch.ops.bitsandbytes.gemv_4bit.default(A, B_q, B.shape, absmax, code, blocksize) assert out.device == A.device assert out.dtype == dtype assert out.shape == (1, 1, out_features) assert out.isreal().all() opcheck(torch.ops.bitsandbytes.gemv_4bit.default, (A, B_q, B.shape, absmax, code, blocksize)) bitsandbytes/tests/test_optim.py000066400000000000000000000547101506752665600175050ustar00rootroot00000000000000import os from os.path import join import shutil import sys import time import uuid from lion_pytorch import Lion import pytest import torch import bitsandbytes as bnb import bitsandbytes.functional as F from bitsandbytes.utils import sync_gpu from tests.helpers import describe_dtype, get_available_devices, id_formatter # import apex k = 20 def assert_most_approx_close(a, b, rtol=1e-3, atol=1e-3, max_error_count=0): idx = torch.isclose(a, b, rtol=rtol, atol=atol) error_count = (idx == 0).sum().item() if error_count > max_error_count: print(f"Too many values not close: assert {error_count} < {max_error_count}") torch.testing.assert_close(a, b, rtol=rtol, atol=atol) def get_temp_dir(): path = f"/tmp/autoswap/{uuid.uuid4()}" os.makedirs(path, exist_ok=True) return path def rm_path(path): shutil.rmtree(path) str2optimizers = {} ## TODO: maybe remove these three. str2optimizers["adam_pytorch"] = (None, torch.optim.Adam, bnb.optim.Adam) str2optimizers["lion_pytorch"] = (None, Lion, bnb.optim.Lion) str2optimizers["momentum_pytorch"] = ( None, lambda pxx: torch.optim.SGD(pxx, 0.01, 0.9), bnb.optim.Adam, ) str2optimizers["adam"] = (torch.optim.Adam, bnb.optim.Adam) str2optimizers["adam8bit_blockwise"] = (torch.optim.Adam, lambda pxx: bnb.optim.Adam8bit(pxx, block_wise=True)) str2optimizers["paged_adam"] = (torch.optim.Adam, bnb.optim.PagedAdam) str2optimizers["paged_adamw"] = (torch.optim.AdamW, bnb.optim.PagedAdamW) str2optimizers["paged_adam8bit_blockwise"] = ( torch.optim.Adam, lambda pxx: bnb.optim.PagedAdam8bit(pxx, block_wise=True), ) str2optimizers["paged_adamw8bit_blockwise"] = ( torch.optim.AdamW, lambda pxx: bnb.optim.PagedAdamW8bit(pxx, block_wise=True), ) str2optimizers["ademamix"] = (bnb.optim.ademamix._ReferenceAdEMAMix, bnb.optim.AdEMAMix) str2optimizers["ademamix8bit_blockwise"] = ( bnb.optim.ademamix._ReferenceAdEMAMix, lambda pxx: bnb.optim.AdEMAMix8bit(pxx), ) str2optimizers["paged_ademamix"] = (bnb.optim.ademamix._ReferenceAdEMAMix, bnb.optim.PagedAdEMAMix) str2optimizers["paged_ademamix8bit_blockwise"] = ( bnb.optim.ademamix._ReferenceAdEMAMix, lambda pxx: bnb.optim.PagedAdEMAMix8bit(pxx), ) str2optimizers["ademamix_scheduled"] = ( lambda pxx: bnb.optim.ademamix._ReferenceAdEMAMix(pxx, t_alpha=k, t_beta3=k), lambda pxx: bnb.optim.AdEMAMix(pxx, t_alpha=k, t_beta3=k), ) str2optimizers["paged_ademamix_scheduled"] = ( lambda pxx: bnb.optim.ademamix._ReferenceAdEMAMix(pxx, t_alpha=k, t_beta3=k), lambda pxx: bnb.optim.PagedAdEMAMix(pxx, t_alpha=k, t_beta3=k), ) str2optimizers["ademamix8bit_blockwise_scheduled"] = ( lambda pxx: bnb.optim.ademamix._ReferenceAdEMAMix(pxx, t_alpha=100, t_beta3=100), lambda pxx: bnb.optim.AdEMAMix8bit(pxx, t_alpha=100, t_beta3=100), ) str2optimizers["paged_ademamix8bit_blockwise_scheduled"] = ( lambda pxx: bnb.optim.ademamix._ReferenceAdEMAMix(pxx, t_alpha=100, t_beta3=100), lambda pxx: bnb.optim.PagedAdEMAMix8bit(pxx, t_alpha=100, t_beta3=100), ) str2optimizers["lion"] = (Lion, bnb.optim.Lion) str2optimizers["paged_lion"] = (Lion, bnb.optim.PagedLion) str2optimizers["lion8bit_blockwise"] = (Lion, lambda pxx: bnb.optim.Lion8bit(pxx, block_wise=True)) str2optimizers["paged_lion8bit_blockwise"] = (Lion, lambda pxx: bnb.optim.PagedLion8bit(pxx, block_wise=True)) str2optimizers["momentum"] = ( lambda pxx: torch.optim.SGD(pxx, 0.01, 0.9), lambda pxx: bnb.optim.SGD(pxx, 0.01, 0.9, block_wise=False), ) str2optimizers["momentum8bit_blockwise"] = ( lambda pxx: torch.optim.SGD(pxx, 0.01, 0.9), lambda pxx: bnb.optim.SGD8bit(pxx, 0.01, 0.9, block_wise=True), ) str2optimizers["rmsprop"] = ( lambda pxx: torch.optim.RMSprop(pxx, 0.01, 0.9), lambda pxx: bnb.optim.RMSprop(pxx, 0.01, 0.9, block_wise=False), ) str2optimizers["rmsprop8bit_blockwise"] = ( lambda pxx: torch.optim.RMSprop(pxx, 0.01, 0.9), lambda pxx: bnb.optim.RMSprop8bit(pxx, 0.01, 0.9, block_wise=True), ) str2statenames = {} str2statenames["adam"] = [("exp_avg", "state1"), ("exp_avg_sq", "state2")] str2statenames["paged_adamw"] = [("exp_avg", "state1"), ("exp_avg_sq", "state2")] str2statenames["paged_adam"] = [("exp_avg", "state1"), ("exp_avg_sq", "state2")] str2statenames["lion"] = [("exp_avg", "state1")] str2statenames["paged_lion"] = [("exp_avg", "state1")] str2statenames["momentum"] = [("momentum_buffer", "state1")] str2statenames["lamb"] = [("exp_avg", "state1"), ("exp_avg_sq", "state2")] str2statenames["rmsprop"] = [("square_avg", "state1")] str2statenames["adam8bit_blockwise"] = [ ("exp_avg", "state1", "qmap1", "absmax1"), ("exp_avg_sq", "state2", "qmap2", "absmax2"), ] str2statenames["paged_adam8bit_blockwise"] = [ ("exp_avg", "state1", "qmap1", "absmax1"), ("exp_avg_sq", "state2", "qmap2", "absmax2"), ] str2statenames["paged_adamw8bit_blockwise"] = [ ("exp_avg", "state1", "qmap1", "absmax1"), ("exp_avg_sq", "state2", "qmap2", "absmax2"), ] str2statenames["momentum8bit_blockwise"] = [("momentum_buffer", "state1", "qmap1", "absmax1")] str2statenames["rmsprop8bit_blockwise"] = [("square_avg", "state1", "qmap1", "absmax1")] str2statenames["lion8bit_blockwise"] = [("exp_avg", "state1", "qmap1", "absmax1")] str2statenames["paged_lion8bit_blockwise"] = [("exp_avg", "state1", "qmap1", "absmax1")] str2statenames["ademamix"] = str2statenames["ademamix_scheduled"] = [("m1_m2", "state1"), ("nu", "state2")] str2statenames["paged_ademamix"] = str2statenames["paged_ademamix_scheduled"] = [("m1_m2", "state1"), ("nu", "state2")] str2statenames["ademamix8bit_blockwise"] = str2statenames["ademamix8bit_blockwise_scheduled"] = [ ("m1_m2", "state1", "qmap1", "absmax1"), ("nu", "state2", "qmap2", "absmax2"), ] str2statenames["paged_ademamix8bit_blockwise"] = [ ("m1_m2", "state1", "qmap1", "absmax1"), ("nu", "state2", "qmap2", "absmax2"), ] optimizer_names_32bit = [ "adam", "paged_adamw", "paged_adam", "momentum", "rmsprop", "lion", "paged_lion", "ademamix", "ademamix_scheduled", "paged_ademamix", "paged_ademamix_scheduled", ] @pytest.mark.parametrize("optim_name", optimizer_names_32bit, ids=id_formatter("opt")) @pytest.mark.parametrize("gtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize("dim1", [1024], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [32, 1024, 4097, 1], ids=id_formatter("dim2")) @pytest.mark.parametrize("device", get_available_devices(no_cpu=True), ids=id_formatter("device")) @pytest.mark.skipif(not get_available_devices(no_cpu=True), reason="No device") def test_optimizer32bit(dim1, dim2, gtype, optim_name, device): if device not in ["cuda", "xpu"]: pytest.skip("Optimizers are only supported on CUDA and XPU") if optim_name.startswith("paged_") and sys.platform == "win32": pytest.skip("Paged optimizers can have issues on Windows.") if optim_name.startswith("paged_") and device == "xpu": pytest.skip("Paged optimizers are not supported on XPU currently.") if gtype == torch.bfloat16 and optim_name in ["momentum", "rmsprop"]: pytest.skip() if dim1 == 1 and dim2 == 1: return p1 = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.1 p2 = p1.clone() p1 = p1.float() torch_optimizer = str2optimizers[optim_name][0]([p1]) bnb_optimizer = str2optimizers[optim_name][1]([p2]) if gtype == torch.float32: atol, rtol = 1e-6, 1e-5 elif gtype == torch.bfloat16: atol, rtol = 1e-3, 1e-2 else: atol, rtol = 1e-4, 1e-3 for i in range(k): g = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.01 p1.grad = g.clone().float() p2.grad = g.clone() bnb_optimizer.step() torch_optimizer.step() for name1, name2 in str2statenames[optim_name]: torch.testing.assert_close( torch_optimizer.state[p1][name1], bnb_optimizer.state[p2][name2].to(device), atol=atol, rtol=rtol, ) # since Lion can have pretty noisy updates where things lie at the boundary # allow up to 15 errors for Lion assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=15) if i % (k // 5) == 0 and i > 0: path = get_temp_dir() torch.save(bnb_optimizer.state_dict(), join(path, "opt.pt")) del bnb_optimizer bnb_optimizer = None bnb_optimizer = str2optimizers[optim_name][1]([p2]) bnb_optimizer.load_state_dict(torch.load(join(path, "opt.pt"))) rm_path(path) # since Lion can have pretty noisy updates where things lie at the boundary # allow up to 10 errors for Lion assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=10) for name1, name2 in str2statenames[optim_name]: # since Lion can have pretty noisy updates where things lie at the boundary # allow up to 10 errors for Lion assert_most_approx_close( torch_optimizer.state[p1][name1], bnb_optimizer.state[p2][name2], atol=atol, rtol=rtol, max_error_count=10, ) if gtype != torch.float32: # the adam buffers should also be close because they are 32-bit # but the parameters can diverge because they are 16-bit # the difference grow larger and larger with each update # --> copy the state to keep weights close p1.data = p1.data.to(p2.dtype).float() p2.copy_(p1.data) torch.testing.assert_close(p1.to(p2.dtype), p2) if optim_name in ["lars", "lamb"]: assert bnb_optimizer.state[p2]["unorm_vec"] > 0.0 @pytest.mark.parametrize("dim1", [1024], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [32, 1024, 4097], ids=id_formatter("dim2")) @pytest.mark.parametrize("gtype", [torch.float32, torch.float16], ids=describe_dtype) @pytest.mark.parametrize("device", get_available_devices(no_cpu=True)) @pytest.mark.skipif(not get_available_devices(no_cpu=True), reason="No device") def test_global_config(dim1, dim2, gtype, device): if device not in ["cuda", "xpu"]: pytest.skip("Optimizers are only supported on CUDA and XPU") if dim1 == 1 and dim2 == 1: return p1 = torch.randn(dim1, dim2, device="cpu", dtype=gtype) * 0.1 p2 = torch.randn(dim1, dim2, device="cpu", dtype=gtype) * 0.1 p3 = torch.randn(dim1, dim2, device="cpu", dtype=gtype) * 0.1 mask = torch.rand_like(p2) < 0.1 beta1 = 0.9 beta2 = 0.999 lr = 0.001 eps = 1e-8 bnb.optim.GlobalOptimManager.get_instance().initialize() bnb.optim.GlobalOptimManager.get_instance().override_config(p3, "optim_bits", 8) bnb.optim.GlobalOptimManager.get_instance().register_parameters([p1, p2, p3]) p1 = p1.to(device) p2 = p2.to(device) p3 = p3.to(device) adam2 = bnb.optim.Adam([p1, p2, p3], lr, (beta1, beta2), eps) if gtype == torch.float32: atol, rtol = 1e-6, 1e-5 else: atol, rtol = 1e-4, 1e-3 for i in range(50): g1 = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.1 + 0.001 g2 = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.1 + 0.001 g3 = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.1 + 0.001 p1.grad = g1 p2.grad = g2 p3.grad = g3 adam2.step() assert adam2.state[p3]["state1"].dtype == torch.uint8 assert adam2.state[p3]["state2"].dtype == torch.uint8 optimizer_names_8bit = [ "adam8bit_blockwise", "lion8bit_blockwise", "momentum8bit_blockwise", "rmsprop8bit_blockwise", "ademamix8bit_blockwise", "ademamix8bit_blockwise_scheduled", ] @pytest.mark.parametrize("optim_name", optimizer_names_8bit, ids=id_formatter("opt")) @pytest.mark.parametrize("gtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize("dim2", [32, 1024, 4097], ids=id_formatter("dim2")) @pytest.mark.parametrize("dim1", [1024], ids=id_formatter("dim1")) @pytest.mark.parametrize("device", get_available_devices(no_cpu=True)) @pytest.mark.skipif(not get_available_devices(no_cpu=True), reason="No device") def test_optimizer8bit(dim1, dim2, gtype, optim_name, device): if device not in ["cuda", "xpu"]: pytest.skip("8-bit optimizers are only supported on CUDA and XPU") torch.set_printoptions(precision=6) if dim1 == 1 and dim2 == 1: return p1 = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.1 p2 = p1.clone() p1 = p1.float() blocksize = 256 torch_optimizer = str2optimizers[optim_name][0]([p1]) bnb_optimizer = str2optimizers[optim_name][1]([p2]) if gtype == torch.float32: atol, rtol = 3e-3, 1e-3 patol, prtol = 1e-5, 1e-3 elif gtype == torch.bfloat16: atol, rtol = 3e-3, 1e-3 patol, prtol = 1e-4, 1e-2 else: atol, rtol = 3e-3, 1e-3 patol, prtol = 1e-5, 1e-3 errors = [] relerrors = [] for i in range(50): g = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.01 p1.grad = g.clone().float() p2.grad = g.clone() torch_optimizer.step() bnb_optimizer.step() # since Lion can have pretty noisy updates where things lie at the boundary # assert_most_approx_close(p1, p2.float(), patol, prtol, max_error_count=0) dequant_states = [] for name1, name2, qmap, max_val in str2statenames[optim_name]: ## For AdEMAMix, we need to dequantize [p2][name2][0] and [p2][name2][1] ## separately and then stack them. The qmap is shared, but absmax is also stacked. if optim_name == "ademamix8bit_blockwise" and name1 == "m1_m2": m1 = F.dequantize_blockwise( code=bnb_optimizer.state[p2][qmap], absmax=bnb_optimizer.state[p2][max_val][0], A=bnb_optimizer.state[p2][name2][0], blocksize=blocksize, ) m2 = F.dequantize_blockwise( code=bnb_optimizer.state[p2][qmap], absmax=bnb_optimizer.state[p2][max_val][1], A=bnb_optimizer.state[p2][name2][1], blocksize=blocksize, ) s1 = torch.stack((m1, m2)) else: s1 = F.dequantize_blockwise( code=bnb_optimizer.state[p2][qmap], absmax=bnb_optimizer.state[p2][max_val], A=bnb_optimizer.state[p2][name2], blocksize=blocksize, ) num_not_close = torch.isclose(torch_optimizer.state[p1][name1], s1, atol=atol, rtol=rtol) == 0 assert num_not_close.sum().item() < 20 dequant_states.append(s1.clone()) err = torch.abs(p1 - p2) relerr = err / (torch.abs(p1) + 1e-9) if g.dtype == torch.bfloat16: assert err.mean() <= 0.00017 assert relerr.mean() <= 0.0016 else: assert err.mean() < 0.00006 assert relerr.mean() < 0.0006 errors.append(err.mean().item()) relerrors.append(relerr.mean().item()) if i % 10 == 0 and i > 0: for (name1, name2, qmap, max_val), s in zip(str2statenames[optim_name], dequant_states): s1cpy = s.clone() raws1cpy = bnb_optimizer.state[p2][name2].clone() qmap1 = bnb_optimizer.state[p2][qmap].clone() path = get_temp_dir() torch.save(bnb_optimizer.state_dict(), join(path, "opt.pt")) del bnb_optimizer bnb_optimizer = None bnb_optimizer = str2optimizers[optim_name][1]([p2]) bnb_optimizer.load_state_dict(torch.load(join(path, "opt.pt"))) rm_path(path) torch.testing.assert_close(raws1cpy, bnb_optimizer.state[p2][name2]) torch.testing.assert_close(qmap1, bnb_optimizer.state[p2][qmap]) ## For AdEMAMix, we need to dequantize [p2][name2][0] and [p2][name2][1] ## separately and then stack them. The qmap is shared, but absmax is also stacked. if optim_name == "ademamix8bit_blockwise" and name1 == "m1_m2": s1 = torch.stack( ( F.dequantize_blockwise( code=bnb_optimizer.state[p2][qmap], absmax=bnb_optimizer.state[p2][max_val][0], A=bnb_optimizer.state[p2][name2][0], blocksize=blocksize, ), F.dequantize_blockwise( code=bnb_optimizer.state[p2][qmap], absmax=bnb_optimizer.state[p2][max_val][1], A=bnb_optimizer.state[p2][name2][1], blocksize=blocksize, ), ) ) else: s1 = F.dequantize_blockwise( code=bnb_optimizer.state[p2][qmap], absmax=bnb_optimizer.state[p2][max_val], A=bnb_optimizer.state[p2][name2], blocksize=blocksize, ) torch.testing.assert_close(s1cpy, s1) num_not_close = torch.isclose(torch_optimizer.state[p1][name1], s1, atol=atol, rtol=rtol) == 0 assert num_not_close.sum().item() < 20 # Lion can have pretty noisy updates where things lie at the boundary assert_most_approx_close(p1, p2.float(), patol, prtol, max_error_count=0) # the parameters diverge quickly. Here we keep them close # together so we can test against the Adam error p1.data = p1.data.to(gtype).float() p2.copy_(p1.data) torch.testing.assert_close(p1.to(gtype), p2) for (name1, name2, qmap, max_val), s in zip(str2statenames[optim_name], dequant_states): torch_optimizer.state[p1][name1].copy_(s.data) @pytest.mark.parametrize("optim_bits", [32, 8], ids=id_formatter("optim_bits")) @pytest.mark.parametrize("gtype", [torch.float32], ids=describe_dtype) @pytest.mark.parametrize("dim2", [32, 1024, 4097], ids=id_formatter("dim2")) @pytest.mark.parametrize("dim1", [1024], ids=id_formatter("dim1")) @pytest.mark.deprecated def test_adam_percentile_clipping(requires_cuda, dim1, dim2, gtype, optim_bits): if dim1 == 1 and dim2 == 1: return p1 = torch.randn(dim1, dim2, device="cpu", dtype=gtype) * 0.1 beta1 = 0.9 beta2 = 0.999 lr = 0.001 eps = 1e-8 p1 = p1.cuda() p2 = p1.clone() adam1 = bnb.optim.Adam([p1], lr, (beta1, beta2), eps, optim_bits=optim_bits) adam2 = bnb.optim.Adam( [p2], lr, (beta1, beta2), eps, optim_bits=optim_bits, percentile_clipping=5, ) gnorm_vec = torch.zeros(100).cuda() step = 0 for i in range(50): step += 1 g1 = torch.randn(dim1, dim2, device="cuda", dtype=gtype) * 0.1 + (0.01 * i) g2 = g1.clone() p2.grad = g2 current_gnorm, clip_val, gnorm_scale = F.percentile_clipping(g1, gnorm_vec, step, 5) g1 = (g1.float() * gnorm_scale).to(gtype) p1.grad = g1 adam1.step() adam2.step() # gnorm_scale is not deterministic (warp reductions), as such there can be slight differences in state if optim_bits == 32: torch.testing.assert_close(p1, p2) torch.testing.assert_close( adam1.state[p1]["state1"], adam2.state[p2]["state1"], atol=5e-5, rtol=1e-4, ) torch.testing.assert_close( adam1.state[p1]["state2"], adam2.state[p2]["state2"], atol=5e-5, rtol=1e-4, ) elif optim_bits == 8: torch.testing.assert_close(p1, p2, atol=1e-4, rtol=1e-3) torch.testing.assert_close( adam1.state[p1]["state1"], adam2.state[p2]["state1"], atol=2, rtol=1e-3, ) torch.testing.assert_close( adam1.state[p1]["state2"], adam2.state[p2]["state2"], atol=2, rtol=1e-3, ) adam1.state[p1]["state1"].copy_(adam2.state[p2]["state1"]) adam1.state[p1]["state2"].copy_(adam2.state[p2]["state2"]) if i % 10 == 0 and i > 0: path = get_temp_dir() torch.save(adam2.state_dict(), join(path, "opt.pt")) del adam2 adam2 = None adam2 = bnb.optim.Adam( [p2], lr, (beta1, beta2), eps, optim_bits=optim_bits, percentile_clipping=5, ) adam2.load_state_dict(torch.load(join(path, "opt.pt"))) optimizer_names_benchmark = [ "adam8bit_blockwise", "paged_adam8bit_blockwise", "ademamix8bit_blockwise", "paged_ademamix8bit_blockwise", "ademamix8bit_blockwise_scheduled", "paged_ademamix8bit_blockwise_scheduled", "lion8bit_blockwise", "paged_lion8bit_blockwise", "paged_ademamix8bit_blockwise", ] @pytest.mark.parametrize("dim1", [4096], ids=id_formatter("dim1")) @pytest.mark.parametrize("dim2", [4096], ids=id_formatter("dim2")) @pytest.mark.parametrize("gtype", [torch.float32, torch.bfloat16, torch.float16], ids=describe_dtype) @pytest.mark.parametrize("optim_name", optimizer_names_benchmark, ids=id_formatter("opt")) @pytest.mark.benchmark def test_benchmark_blockwise(dim1, dim2, gtype, optim_name, device): if dim1 == 1 and dim2 == 1: return p1 = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.1 bnb_optimizer = str2optimizers[optim_name][1]([p1]) g = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.01 p1.grad = g total_steps = 500 for i in range(total_steps): if i == total_steps // 5: # 100 iterations for burn-in sync_gpu(p1) t0 = time.time() bnb_optimizer.step() sync_gpu(p1) s = time.time() - t0 print("") params = (total_steps - total_steps // 5) * dim1 * dim2 print(optim_name, gtype, s, params, s / params) # assert s < 3.9 bitsandbytes/tests/test_parametrize.py000066400000000000000000000466251506752665600207060ustar00rootroot00000000000000import pytest import torch import torch.nn as nn from bitsandbytes import functional as F from bitsandbytes.cextension import HIP_ENVIRONMENT from bitsandbytes.nn.parametrize import ( Bnb4bitParametrization, replace_parameter_4bit, replace_parameter_4bit_prequantized, ) from tests.helpers import ( TRUE_FALSE, describe_dtype, get_available_devices, id_formatter, is_supported_on_hpu, ) class ParametrizeTestModule(nn.Module): """Test module with different parameter shapes for testing parametrization.""" def __init__(self, device="cpu", dtype=torch.float32): super().__init__() # 2D parameter (typical weight matrix) self.weight_2d = nn.Parameter(torch.randn(1024, 1024, device=device, dtype=dtype)) # 3D parameter (MoE expert weights - the main use case for this feature) self.expert_weights = nn.Parameter(torch.randn(8, 512, 256, device=device, dtype=dtype)) # 1D parameter (bias-like) self.bias_1d = nn.Parameter(torch.randn(1024, device=device, dtype=dtype)) # Non-parameter attribute (should not be quantizable) self.not_param = torch.randn(32, device=device, dtype=dtype) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize("quant_type", ["nf4", "fp4"]) @pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics")) @pytest.mark.parametrize( "blocksize", [64, 128, 256] if not HIP_ENVIRONMENT else [128, 256], ) def test_replace_parameter_4bit(device, dtype, quant_type, compress_statistics, blocksize): """Test basic parameter replacement with 4-bit quantization on different dtypes.""" if device == "hpu" and not is_supported_on_hpu(quant_type, dtype): pytest.skip("This configuration is not supported on HPU.") # Create module directly on target device to avoid unnecessary transfers module = ParametrizeTestModule(device=device, dtype=dtype) original_param = module.weight_2d.clone() # Apply 4-bit quantization parametrization to the weight parameter replace_parameter_4bit( module, "weight_2d", compress_statistics=compress_statistics, quant_type=quant_type, blocksize=blocksize ) # Verify that parametrization was applied correctly assert hasattr(module, "parametrizations"), "Module should have parametrizations attribute" assert "weight_2d" in module.parametrizations, "weight_2d should be parametrized" # Test that accessing the parameter returns dequantized version with correct properties reconstructed = module.weight_2d assert reconstructed.shape == original_param.shape, "Shape should be preserved" assert reconstructed.dtype == dtype, "dtype should match original" assert reconstructed.device.type == device, "Device should match target" # Verify quantization quality using same approach as functional tests err = (original_param - reconstructed.detach()).abs().float() relerr = (err / (original_param.abs().float() + 1e-8)).mean() err_mean = err.mean() # Expected error bounds from test_functional.py expected_errors = { "nf4": { 64: {"abs": 0.072792, "rel": 0.203299}, 128: {"abs": 0.076835, "rel": 0.215252}, 256: {"abs": 0.080326, "rel": 0.226044}, }, "fp4": { 64: {"abs": 0.096545, "rel": 0.260130}, 128: {"abs": 0.102947, "rel": 0.275734}, 256: {"abs": 0.108685, "rel": 0.289842}, }, } assert err_mean < expected_errors[quant_type][blocksize]["abs"] + 1e-3, f"Mean abs error {err_mean:.6f} too high" assert relerr < expected_errors[quant_type][blocksize]["rel"] + 1e-3, f"Mean rel error {relerr:.6f} too high" @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) def test_moe_parameter_shape(device, dtype): """Test parametrization with MoE-style parameter shape""" if device == "hpu" and not is_supported_on_hpu("nf4", dtype): pytest.skip("This configuration is not supported on HPU.") param_shape = (8, 64, 32) # Create module with custom parameter shape directly on target device class MoEModule(nn.Module): def __init__(self, device, dtype): super().__init__() self.param = nn.Parameter(torch.randn(*param_shape, dtype=dtype, device=device)) module = MoEModule(device=device, dtype=dtype) original_param = module.param.clone() # Apply quantization parametrization replace_parameter_4bit(module, "param", quant_type="nf4") # Verify reconstruction maintains all properties reconstructed = module.param assert reconstructed.shape == param_shape, f"Shape should be preserved: {reconstructed.shape} vs {param_shape}" assert reconstructed.dtype == dtype, "dtype should match original" assert reconstructed.device.type == device, "Device should match target" # Verify quantization quality using error calculation approach from functional tests err = (original_param - reconstructed.detach()).abs().float() relerr = (err / (original_param.abs().float() + 1e-8)).mean() err_mean = err.mean() # Use slightly looser bounds for higher dimensional tensors abs_bound = 0.085 # NF4 baseline + margin rel_bound = 0.25 # NF4 baseline + margin assert err_mean < abs_bound, f"Mean abs error {err_mean:.6f} too high for shape {param_shape}" assert relerr < rel_bound, f"Mean rel error {relerr:.6f} too high for shape {param_shape}" @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize("quant_type", ["nf4", "fp4"]) def test_prequantized_replacement(device, dtype, quant_type): """Test applying parametrization to already quantized parameters.""" if device == "hpu" and not is_supported_on_hpu(quant_type, dtype): pytest.skip("Configuration not supported on HPU.") module = ParametrizeTestModule(device=device, dtype=dtype) original_param = module.weight_2d.clone() # Manually quantize the parameter data first (simulates loading pre-quantized weights) quantized_data, quant_state = F.quantize_4bit(original_param.data, quant_type=quant_type) # Replace parameter with quantized data (what would happen during model loading) module.weight_2d = nn.Parameter(quantized_data, requires_grad=False) # Apply parametrization to handle dequantization on access replace_parameter_4bit_prequantized( module, "weight_2d", quant_state.as_dict(packed=True), device=torch.device(device) ) # Test that parameter access properly dequantizes reconstructed = module.weight_2d assert reconstructed.shape == original_param.shape, "Shape should be preserved" assert reconstructed.dtype == dtype, "dtype should match original" assert reconstructed.device.type == device, "Device should match target" @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize("quant_type", ["nf4", "fp4"]) @pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics")) @pytest.mark.skipif(torch.__version__ < (2, 5), reason="state dict hook requires torch >= 2.5.0") def test_state_dict_functionality(device, dtype, quant_type, compress_statistics): """Test that state dict saving works with quantized parameters.""" if device == "hpu" and not is_supported_on_hpu(quant_type, dtype): pytest.skip("Configuration not supported on HPU.") module = ParametrizeTestModule(device=device, dtype=dtype) # Apply parametrization to expert weights (main MoE use case) replace_parameter_4bit(module, "expert_weights", quant_type=quant_type, compress_statistics=compress_statistics) # Save state dict - should include quantization state, not parametrization internals state_dict = module.state_dict() # Verify state dict structure: quantized param + quantization metadata assert "expert_weights" in state_dict, "Quantized parameter should be in state dict" assert "expert_weights.absmax" in state_dict, "Quantization absmax should be saved" assert "expert_weights.quant_map" in state_dict, "Quantization map should be saved" assert f"expert_weights.quant_state.bitsandbytes__{quant_type}" in state_dict, "Quant state should be saved" # Verify parametrization internals are NOT saved (clean state dict) assert "parametrizations.expert_weights.original" not in state_dict, ( "Internal parametrization keys should not be saved" ) # Test that the parameter can be accessed after state dict creation reconstructed = module.expert_weights assert reconstructed.shape == (8, 512, 256), "Shape should be preserved" assert reconstructed.dtype == dtype, "dtype should match" @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) def test_moe_realistic_forward(device, dtype): """Test realistic MoE forward computation with quantized expert weights.""" if device == "hpu" and not is_supported_on_hpu("nf4", dtype): pytest.skip("Configuration not supported on HPU.") class SimpleMoE(nn.Module): def __init__(self, device, dtype): super().__init__() # Expert weights: [num_experts, input_dim, output_dim] self.expert_weights = nn.Parameter(torch.randn(4, 32, 64, dtype=dtype, device=device)) def forward(self, x, expert_idx=0): # Select and use specific expert weight matrix expert_weight = self.expert_weights[expert_idx] # Shape: [input_dim, output_dim] return torch.matmul(x, expert_weight) module = SimpleMoE(device=device, dtype=dtype) x = torch.randn(8, 32, dtype=dtype, device=device) # Get reference output before quantization with torch.no_grad(): reference_output = module(x, expert_idx=1) # Apply 4-bit quantization to expert weights replace_parameter_4bit(module, "expert_weights", quant_type="nf4") # Get output after quantization - should be very close to original with torch.no_grad(): quantized_output = module(x, expert_idx=1) # Verify outputs match within quantization tolerance assert quantized_output.shape == reference_output.shape, "Output shape should be preserved" # Calculate error like functional tests (matrix ops may amplify quantization errors) err = (reference_output - quantized_output).abs().float() relerr = (err / (reference_output.abs().float() + 1e-8)).mean() err_mean = err.mean() # Allow for error amplification through matrix multiplication assert err_mean < 0.5, f"Forward pass mean abs error {err_mean:.6f} too high" assert relerr < 2.0, f"Forward pass mean rel error {relerr:.6f} too high" def test_error_conditions(): """Test that proper errors are raised for invalid inputs.""" module = ParametrizeTestModule() # Test AttributeError for non-existent parameter with pytest.raises(AttributeError, match="Module does not have parameter 'nonexistent'"): replace_parameter_4bit(module, "nonexistent") # Test TypeError for non-Parameter attribute with pytest.raises(TypeError, match="Parameter 'not_param' is not an instance of nn.Parameter"): replace_parameter_4bit(module, "not_param") # Test same errors for prequantized version with pytest.raises(AttributeError, match="Module does not have parameter 'nonexistent'"): replace_parameter_4bit_prequantized(module, "nonexistent", {}, torch.device("cpu")) with pytest.raises(TypeError, match="Parameter 'not_param' is not an instance of nn.Parameter"): replace_parameter_4bit_prequantized(module, "not_param", {}, torch.device("cpu")) @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.skipif(torch.__version__ < (2, 5), reason="state dict hook requires torch >= 2.5.0") def test_quant_state_preservation(device, dtype): """Test that quantization state is properly preserved and accessible.""" if device == "hpu" and not is_supported_on_hpu("nf4", dtype): pytest.skip("Configuration not supported on HPU.") module = ParametrizeTestModule(device=device, dtype=dtype) blocksize = 128 if HIP_ENVIRONMENT else 64 # Apply parametrization with specific settings replace_parameter_4bit(module, "weight_2d", quant_type="nf4", compress_statistics=True, blocksize=blocksize) # Verify that quantization state is accessible through parametrization parametrization = module.parametrizations.weight_2d[0] assert isinstance(parametrization, Bnb4bitParametrization), "Should be Bnb4bitParametrization instance" # Check quantization state properties quant_state = parametrization.quant_state assert isinstance(quant_state, F.QuantState), "Should have QuantState" assert quant_state.quant_type == "nf4", "Quant type should be preserved" assert quant_state.blocksize == blocksize, "Block size should be preserved" # Verify that state dict includes all necessary quantization metadata state_dict = module.state_dict() quant_state_dict = quant_state.as_dict(packed=True) for key in quant_state_dict.keys(): full_key = f"weight_2d.{key}" assert full_key in state_dict, f"Quantization metadata '{full_key}' should be in state dict" @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.skipif(torch.__version__ < (2, 5), reason="state dict hook requires torch >= 2.5.0") def test_multiple_parameters(device, dtype): """Test applying parametrization to multiple parameters in the same module.""" if device == "hpu" and not is_supported_on_hpu("nf4", dtype): pytest.skip("Configuration not supported on HPU.") module = ParametrizeTestModule(device=device, dtype=dtype) original_2d = module.weight_2d.clone() original_3d = module.expert_weights.clone() # Apply parametrization to multiple parameters, with varying configurations replace_parameter_4bit(module, "weight_2d", quant_type="nf4", blocksize=128) replace_parameter_4bit(module, "expert_weights", quant_type="fp4", blocksize=256) # Verify both parameters are parametrized and work correctly reconstructed_2d = module.weight_2d reconstructed_3d = module.expert_weights assert reconstructed_2d.shape == original_2d.shape, "2D parameter shape should be preserved" assert reconstructed_3d.shape == original_3d.shape, "3D parameter shape should be preserved" # Check that state dict includes quantization info for both parameters state_dict = module.state_dict() assert "weight_2d" in state_dict, "2D parameter should be in state dict" assert "expert_weights" in state_dict, "3D parameter should be in state dict" assert "weight_2d.absmax" in state_dict, "2D parameter quantization metadata should be saved" assert "expert_weights.absmax" in state_dict, "3D parameter quantization metadata should be saved" @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize( "blocksize", [64, 128, 256] if not HIP_ENVIRONMENT else [128, 256], ) def test_different_blocksizes(device, dtype, blocksize): """Test parametrization with different block sizes to verify flexibility.""" if device == "hpu" and not is_supported_on_hpu("nf4", dtype): pytest.skip("Configuration not supported on HPU.") module = ParametrizeTestModule(device=device, dtype=dtype) original_param = module.expert_weights.clone() # Apply parametrization with specified block size replace_parameter_4bit(module, "expert_weights", quant_type="nf4", blocksize=blocksize) # Verify reconstruction works with different block sizes reconstructed = module.expert_weights assert reconstructed.shape == original_param.shape, "Shape should be preserved" assert reconstructed.device.type == device, "Device should match" # Verify quantization quality using error calculation approach from functional tests err = (original_param - reconstructed.detach()).abs().float() relerr = (err / (original_param.abs().float() + 1e-8)).mean() err_mean = err.mean() # Expected error bounds from functional tests (using NF4 bounds since that's what we're testing) expected_abs = {64: 0.072792, 128: 0.076835, 256: 0.080326} expected_rel = {64: 0.203299, 128: 0.215252, 256: 0.226044} assert err_mean < expected_abs[blocksize] + 0.01, ( f"Mean abs error {err_mean:.6f} too high for blocksize {blocksize}" ) assert relerr < expected_rel[blocksize] + 0.02, f"Mean rel error {relerr:.6f} too high for blocksize {blocksize}" def test_parametrization_forward_method(): """Test the Bnb4bitParametrization forward method directly.""" device = "cpu" # Create test tensor and manually quantize it original_tensor = torch.randn(64, 32, dtype=torch.float32, device=device) quantized_data, quant_state = F.quantize_4bit(original_tensor, quant_type="nf4") # Create parametrization instance parametrization = Bnb4bitParametrization(quant_state) # Test forward pass (dequantization) dequantized = parametrization.forward(quantized_data) # Verify dequantization produces correct output assert dequantized.shape == original_tensor.shape, "Shape should be preserved during dequantization" assert dequantized.dtype == torch.float32, "dtype should be preserved" assert dequantized.device == original_tensor.device, "Device should be preserved" # Check that dequantization approximates original using mean error calculation err = (original_tensor - dequantized.detach()).abs().float() relerr = (err / (original_tensor.abs().float() + 1e-8)).mean() err_mean = err.mean() # Use NF4 bounds from functional tests with small margin assert err_mean < 0.08, f"Mean abs error {err_mean:.6f} too high" assert relerr < 0.25, f"Mean rel error {relerr:.6f} too high" @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype) def test_gradient_behavior(device, dtype): """Test that quantized parameters have proper gradient behavior.""" if device == "hpu" and not is_supported_on_hpu("nf4", dtype): pytest.skip("Configuration not supported on HPU.") module = ParametrizeTestModule(device=device, dtype=dtype) # Ensure original parameter requires gradients module.weight_2d.requires_grad_(True) assert module.weight_2d.requires_grad, "Original parameter should require gradients" # Apply quantization parametrization replace_parameter_4bit(module, "weight_2d", quant_type="nf4") # Verify that quantized parameters don't require gradients (expected behavior) # The underlying quantized parameter should have requires_grad=False # The dequantized output should also not require gradients reconstructed = module.weight_2d assert not reconstructed.requires_grad, "Dequantized parameter should not require gradients" bitsandbytes/tests/test_triton.py000066400000000000000000000050611506752665600176670ustar00rootroot00000000000000import pytest import torch from bitsandbytes.nn import Linear8bitLt from bitsandbytes.nn.triton_based_modules import SwitchBackLinear from bitsandbytes.triton.triton_utils import is_triton_available from tests.helpers import TRUE_FALSE @pytest.mark.skipif( not is_triton_available() or not torch.cuda.is_available() or not torch.cuda.get_device_capability()[0] >= 8, reason="This test requires triton and a GPU with compute capability 8.0 or higher.", ) @pytest.mark.deprecated @pytest.mark.parametrize("vector_wise_quantization", TRUE_FALSE) def test_switchback(vector_wise_quantization): for dim in [83]: for batch in [13]: standard = torch.nn.Linear(dim, 4 * dim).cuda().half() switchback = ( SwitchBackLinear(dim, 4 * dim, vector_wise_quantization=vector_wise_quantization).cuda().half() ) baseline = Linear8bitLt(dim, 4 * dim).cuda().half() switchback.weight.data.copy_(standard.weight) switchback.bias.data.copy_(standard.bias) baseline.weight.data.copy_(standard.weight) baseline.bias.data.copy_(standard.bias) x1 = torch.randn(batch, dim).cuda().half().requires_grad_(True) x2 = x1.clone().detach().requires_grad_(True) x3 = x1.clone().detach().requires_grad_(True) out_standard = standard(x1) (2**10 * out_standard.abs().mean()).backward() print(x2.dtype) out_sb = switchback(x2) (2**10 * out_sb.abs().mean()).backward() out_baseline = baseline(x3) (2**10 * out_baseline.abs().mean()).backward() err_sb = (out_standard - out_sb).abs().mean() err_baseline = (out_standard - out_baseline).abs().mean() print("OUT", err_sb, err_baseline) assert err_sb < 2 * err_baseline err_sb = (standard.bias.grad - switchback.bias.grad).abs().mean() err_baseline = (standard.bias.grad - baseline.bias.grad).abs().mean() print("GW2", err_sb, err_baseline) assert err_sb < 2 * err_baseline err_sb = (standard.weight.grad - switchback.weight.grad).abs().mean() err_baseline = (standard.weight.grad - baseline.weight.grad).abs().mean() print("GW1", err_sb, err_baseline) assert err_sb < 2 * err_baseline err_sb = (x1.grad - x2.grad).abs().mean() err_baseline = (x1.grad - x3.grad).abs().mean() print("GX1", err_sb, err_baseline) assert err_sb < 2 * err_baseline