pax_global_header00006660000000000000000000000064152125723770014525gustar00rootroot0000000000000052 comment=c18a0fa8b969d789f33b7912eef57cd69c56bd9e vllm-project-compressed-tensors-c18a0fa/000077500000000000000000000000001521257237700204655ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/.buildkite/000077500000000000000000000000001521257237700225175ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/.buildkite/pipeline.yml000066400000000000000000000003671521257237700250550ustar00rootroot00000000000000steps: # Test check pipeline - label: "Compressed-Tensors Smoke, Sanity and Regression" command: | PIPELINE_FILE="$${TEST_RUNNER:-H100duo}" buildkite-agent pipeline upload .buildkite/test-check/test-check-$${PIPELINE_FILE}.yml vllm-project-compressed-tensors-c18a0fa/.buildkite/test-check/000077500000000000000000000000001521257237700245515ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/.buildkite/test-check/scripts/000077500000000000000000000000001521257237700262405ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/.buildkite/test-check/scripts/test-check.sh000066400000000000000000000014671521257237700306360ustar00rootroot00000000000000#!/usr/bin/env bash set -euo pipefail # print OS info for debugging cat /etc/issue # fetch full history and tags (setuptools_scm derives version from git tags) git fetch --tags --unshallow 2>/dev/null || git fetch --tags # install system dependencies and uv apt-get update && apt-get install -y curl g++ gcc make curl -LsSf https://astral.sh/uv/install.sh | env UV_VERSION=0.11.18 sh # set up GPU and path export LD_LIBRARY_PATH=/usr/local/nvidia/lib64 export PATH="$HOME/.local/bin:/usr/local/nvidia/bin:$PATH" nvidia-smi # create venv and install dependencies uv venv testvenv --python 3.12 source testvenv/bin/activate export UV_TORCH_BACKEND=cu130 export HF_HOME=/model-cache uv pip install .[dev] --index-strategy unsafe-best-match --extra-index-url https://download.pytorch.org/whl/cu130 # run tests make test vllm-project-compressed-tensors-c18a0fa/.buildkite/test-check/test-check-H100duo.yml000066400000000000000000000026341521257237700304510ustar00rootroot00000000000000steps: - label: "Compressed-Tensors Smoke, Sanity and Regression - H100 duo" agents: queue: RedHat-H100-WDC plugins: - kubernetes: podSpec: serviceAccountName: buildkite-anyuid securityContext: fsGroup: 0 runAsUser: 0 runAsGroup: 0 nodeSelector: vllm.ci/gpu-pool: upstream-ci-h100 containers: - name: tests image: buildkite/agent:3-ubuntu command: - chmod +x .buildkite/test-check/scripts/test-check.sh - .buildkite/test-check/scripts/test-check.sh resources: limits: nvidia.com/gpu: 2 volumeMounts: - name: ci-cache mountPath: /model-cache readOnly: false - name: devshm mountPath: /dev/shm env: - name: HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret key: token volumes: - name: ci-cache hostPath: path: /var/mnt/ci-cache type: DirectoryOrCreate - name: devshm emptyDir: medium: Memory vllm-project-compressed-tensors-c18a0fa/.buildkite/test-check/test-check-L4solo.yml000066400000000000000000000020121521257237700304730ustar00rootroot00000000000000steps: - label: "Compressed-Tensors Smoke, Sanity and Regression - L4 solo" agents: queue: RedHat-L4-GCP plugins: - kubernetes: podSpec: nodeSelector: pool: L4solo containers: - name: tests image: buildkite/agent:3-ubuntu command: - chmod +x .buildkite/test-check/scripts/test-check.sh - .buildkite/test-check/scripts/test-check.sh resources: limits: nvidia.com/gpu: 1 volumeMounts: - name: model-cache mountPath: /model-cache readOnly: false - name: dshm mountPath: /dev/shm volumes: - name: model-cache persistentVolumeClaim: claimName: buildkite-podpvc-ssd-blue - name: dshm emptyDir: medium: Memory vllm-project-compressed-tensors-c18a0fa/.claude/000077500000000000000000000000001521257237700220005ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/.claude/skills/000077500000000000000000000000001521257237700233015ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/.claude/skills/style.md000066400000000000000000000004371521257237700247670ustar00rootroot00000000000000# Code Quality Check After completing a batch of code changes: 1. Run `make style` to auto-format the code 2. Run `make quality` to validate code quality 3. Fix any issues reported by `make quality` When writing code, keep lines under 88 characters and avoid unnecessary indentation. vllm-project-compressed-tensors-c18a0fa/.claude/skills/test.md000066400000000000000000000011621521257237700246020ustar00rootroot00000000000000# Running Tests Not all tests in this project require a GPU, but most tests do require a GPU. When running tests: 1. First check if `canhazgpu` is available: `which canhazgpu` 2. If available, MUST run tests using `canhazgpu` with appropriate GPU allocation 3. Use the format: `canhazgpu --gpus 1 -- python3 -m pytest tests/...` Example: ```bash # Check if canhazgpu is available which canhazgpu # Run tests with canhazgpu (required if available) canhazgpu --gpus 1 -- python3 -m pytest tests/test_example.py ``` Note: Adjust the `--gpus` count based on test requirements (typically 1 GPU is sufficient for most tests). vllm-project-compressed-tensors-c18a0fa/.coderabbit.yaml000066400000000000000000000070231521257237700235270ustar00rootroot00000000000000# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json language: "en" early_access: false reviews: request_changes_workflow: false high_level_summary: true high_level_summary_in_walkthrough: true poem: false review_status: true collapse_walkthrough: false auto_apply_labels: true path_filters: - "!**/*.pyc" - "!**/build/**" - "!**/.pytest_cache/**" - "!**/.ruff_cache/**" - "!**/__pycache__/**" path_instructions: - path: "**/*.py" instructions: "Focus on code correctness, numerical stability, and tensor operations. Check for major PEP 8 violations and Python best practices. Pay attention to GPU memory management and model loading patterns. Ignore trivial formatting issues." - path: "tests/**/*.py" instructions: "Ensure PyTest tests are clear, comprehensive, and cover edge cases for compression and quantization scenarios. Verify proper mocking and test isolation. Don't focus on minor stylistic details." - path: "examples/**/*.py" instructions: "Review for clarity, correctness, and educational value. Ensure examples follow best practices and are runnable. Verify that comments and documentation are helpful for users learning the library." - path: "src/compressed_tensors/compressors/**/*.py" instructions: "Pay special attention to compression algorithm implementations. Verify correctness of packing, quantized storage, and MX format (MXFP4, MXFP8, NVFP4) implementations. Ensure compatibility with different model architectures and weight layouts." - path: "src/compressed_tensors/quantization/**/*.py" instructions: "Focus on quantization lifecycle correctness, scheme definitions, and observer implementations. Verify proper handling of different quantization strategies (FP8, INT4, etc.) and that calibration and activation ordering are correct." - path: "src/compressed_tensors/config/**/*.py" instructions: "Review configuration classes for correctness, serialization/deserialization safety, and backward compatibility. Ensure new fields have sensible defaults and are properly validated." - path: "src/compressed_tensors/transform/**/*.py" instructions: "Check transform factory and utility logic for correctness of weight transformations (e.g., Hadamard, smoothing). Verify numerical stability and proper application order relative to quantization." - path: "src/compressed_tensors/offload/**/*.py" instructions: "Review offloading and cache utilities for memory correctness, proper device placement, and no memory leaks. Verify that offload hooks interact correctly with model loading." - path: "docs/**/*.md" instructions: "Check for clarity, accuracy, and completeness. Ensure code examples in docs are correct and align with the current API." - path: "**/README.md" instructions: "Review for clarity, accuracy, and up-to-date information. Ensure installation instructions and examples are correct." auto_review: enabled: true ignore_title_keywords: - "WIP" - "DO NOT MERGE" - "DRAFT" drafts: false auto_incremental_review: false base_branches: - "main" knowledge_base: linked_repositories: - repository: "vllm-project/llm-compressor" instructions: "Upstream library that uses compressed-tensors for model compression workflows. Useful for understanding how compression configs, quantization schemes, and compressed checkpoints are consumed by end-to-end compression pipelines." chat: auto_reply: true issue_enrichment: auto_enrich: enabled: true vllm-project-compressed-tensors-c18a0fa/.github/000077500000000000000000000000001521257237700220255ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/.github/.gitkeep000066400000000000000000000000001521257237700234440ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/.github/actions/000077500000000000000000000000001521257237700234655ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/.github/actions/test/000077500000000000000000000000001521257237700244445ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/.github/actions/test/action.yml000066400000000000000000000050441521257237700264470ustar00rootroot00000000000000name: test compressed-tensors description: 'test compressed-tensors' inputs: venv: description: "path of virtualenv" required: true suitename: description: "test suite name" required: true code_coverage: description: whether to collect code coverage metrics during test run default: 'false' outputs: status: description: "final status from test" value: ${{ steps.test.outputs.status }} runs: using: composite steps: - name: install wheel uses: neuralmagic/nm-actions/actions/install-whl@v1.2.0 with: venv: ${{ inputs.venv }} name: compressed extra: "[dev]" - name: clean up run: | echo "cleaning up disk space..." find . -type f -name '*.whl' -exec rm -rf {} \; python -m pip cache purge sudo rm -rf /usr/local/.ghcup sudo rm -rf /opt/hostedtoolcache/CodeQL sudo rm -rf /usr/local/lib/android/sdk/ndk sudo rm -rf /usr/share/dotnet sudo rm -rf /opt/ghc sudo rm -rf /usr/local/share/boost if [[ "$(cat /etc/issue)" =~ Ubuntu ]]; then sudo apt-get clean fi df -h shell: bash - name: test id: test run: | source ${{ inputs.venv }}/bin/activate rm -rf src if [[ "${ENABLE_COVERAGE}" == "true" ]]; then echo "::group::Installing code coverage requirements via pip" pip install https://github.com/neuralmagic/pytest-nm-releng/archive/v0.4.0.tar.gz pip install coverage pytest-cov # Adding Code coverage to the tests nmre-generate-coverage-flags --package "compressed_tensors" --output-file ".coverage_flags.sh" source .coverage_flags.sh echo "::endgroup::" fi echo "::group::running tests" echo "PYTEST_ADDOPTS set to: ${PYTEST_ADDOPTS}" SUCCESS=0 pytest tests --junitxml=test-results/report.xml -o junit_suite_name="${{ inputs.suitename }}" || SUCCESS=$? echo "status=${SUCCESS}" >> "$GITHUB_OUTPUT" echo "::endgroup::" if [[ "${ENABLE_COVERAGE}" == "true" ]]; then echo "::group::check coverage reports" if [ ! -d coverage-html ]; then echo "ERROR: coverage-html folder not found" exit 1 fi echo "::endgroup::" fi deactivate exit ${SUCCESS} shell: bash env: ENABLE_COVERAGE: ${{ inputs.code_coverage || false }} vllm-project-compressed-tensors-c18a0fa/.github/mergify.yml000066400000000000000000000032211521257237700242100ustar00rootroot00000000000000pull_request_rules: - name: label-documentation description: Automatically apply documentation label conditions: - label != stale - -closed - or: - files~=^[^/]+\.md$ - files~=^docs/ - files~=^examples/ actions: label: add: - documentation - name: ping author on conflicts and add 'needs-rebase' label conditions: - label != stale - conflict - -closed actions: label: add: - needs-rebase comment: message: | This pull request has merge conflicts that must be resolved before it can be merged. Please rebase the PR, @{{author}}. https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork - name: remove 'needs-rebase' label when conflict is resolved conditions: - -conflict - -closed actions: label: remove: - needs-rebase - name: add quality-failed label conditions: - label != stale - check-failure = quality-check - -closed actions: label: add: - quality-failed comment: message: | The quality checks have failed. Please run `make style` and `make quality` under the root directory to adddress the lint failures. You will need to install the dev optional install to get the required linting packages. - name: remove quality-failed label conditions: - label != stale - -check-failure = quality-check - -closed actions: label: remove: - quality-failedvllm-project-compressed-tensors-c18a0fa/.github/scripts/000077500000000000000000000000001521257237700235145ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/.github/scripts/step-status000077500000000000000000000004011521257237700257310ustar00rootroot00000000000000#!/bin/bash -e # echo "green encased checkmark" if "${1} == 0" # echo "red X" if "${1} != 0" STEP_STATUS=${1} if [ "$STEP_STATUS" -eq 0 ]; then # green check echo -e "\xE2\x9C\x85" else # red x echo -e "\xE2\x9D\x8C" fi vllm-project-compressed-tensors-c18a0fa/.github/workflows/000077500000000000000000000000001521257237700240625ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/.github/workflows/quality-check.yaml000066400000000000000000000012421521257237700275100ustar00rootroot00000000000000name: Quality Checks on: push: branches: - main - 'release/*' pull_request: branches: - main - 'release/*' jobs: quality-check: runs-on: ubuntu-24.04 steps: - uses: actions/setup-python@v5 with: python-version: '3.10' - uses: actions/checkout@v4 with: fetch-depth: 0 fetch-tags: true - name: Set Env run: | pip3 install --upgrade pip && pip3 install --upgrade setuptools - name: "⚙️ Install dependencies" run: pip3 install .[dev] - name: "🧹 Running quality checks" run: make qualityvllm-project-compressed-tensors-c18a0fa/.github/workflows/stale.yml000066400000000000000000000032421521257237700257160ustar00rootroot00000000000000name: 'Close inactive PRs' on: schedule: - cron: '0 17 * * *' jobs: close-pull-requests: if: github.repository == 'vllm-project/compressed-tensors' permissions: issues: write pull-requests: write actions: write runs-on: ubuntu-latest steps: - uses: actions/stale@997185467fa4f803885201cee163a9f38240193d with: operations-per-run: 1000 exempt-draft-pr: true exempt-issue-labels: 'keep-open' exempt-pr-labels: 'keep-open' days-before-issue-stale: 90 days-before-issue-close: 30 stale-issue-label: 'stale' stale-issue-message: > This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you! close-issue-message: > This issue has been automatically closed due to inactivity. Please feel free to reopen if you feel it is still relevant. Thank you! days-before-pr-stale: 90 days-before-pr-close: 30 stale-pr-label: 'stale' stale-pr-message: > This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. close-pr-message: > This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it.vllm-project-compressed-tensors-c18a0fa/.gitignore000066400000000000000000000061031521257237700224550ustar00rootroot00000000000000# Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class # C extensions *.so # Distribution / packaging .Python build/ develop-eggs/ dist/ downloads/ eggs/ .eggs/ lib/ lib64/ parts/ sdist/ var/ wheels/ share/python-wheels/ *.egg-info/ .installed.cfg *.egg MANIFEST src/compressed_tensors/version.py # 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/ cover/ # 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 .pybuilder/ target/ # Jupyter Notebook .ipynb_checkpoints # IPython profile_default/ ipython_config.py # pyenv # For a library or package, you might want to ignore these files since the code is # intended to run in multiple environments; otherwise, check them in: # .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 # poetry # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. # This is especially recommended for binary packages to ensure reproducibility, and is more # commonly ignored for libraries. # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control #poetry.lock # pdm # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. #pdm.lock # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it # in version control. # https://pdm.fming.dev/#use-with-ide .pdm.toml # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm __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/ # pytype static type analyzer .pytype/ # Cython debug symbols cython_debug/ # PyCharm # JetBrains specific template is maintained in a separate JetBrains.gitignore that can # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ examples/**/*.safetensors vllm-project-compressed-tensors-c18a0fa/LICENSE000066400000000000000000000261351521257237700215010ustar00rootroot00000000000000 Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright [yyyy] [name of copyright owner] 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. vllm-project-compressed-tensors-c18a0fa/Makefile000066400000000000000000000022431521257237700221260ustar00rootroot00000000000000.PHONY: build docs test BUILD_TYPE ?= dev # set nightly to build nightly release PYCHECKDIRS := src tests PYCHECKGLOBS := 'src/**/*.py' 'tests/**/*.py' 'utils/**/*.py' 'examples/**/*.py' setup.py # run checks on all files for the repo quality: @echo "Running copyright checks"; python utils/copyright.py quality $(PYCHECKGLOBS) @echo "Running python quality checks"; black --target-version py310 --check $(PYCHECKDIRS); isort --check-only $(PYCHECKDIRS); flake8 $(PYCHECKDIRS); # style the code according to accepted standards for the repo style: @echo "Running copyright checks"; python utils/copyright.py style $(PYCHECKGLOBS) @echo "Running python styling"; black --target-version py310 $(PYCHECKDIRS); isort $(PYCHECKDIRS); # run tests for the repo test: @echo "Running python tests"; pytest -ra tests; @echo "Running emulated XPU tests"; pytest -ra -c pytest-xpu.ini --emulate-xpu; # creates wheel file build: @echo "Building the wheel for the repository"; BUILD_TYPE=$(BUILD_TYPE) python3 setup.py sdist bdist_wheel; # clean package clean: @echo "Cleaning up"; rm -rf .pytest_cache; find $(PYCHECKDIRS) | grep -E "(__pycache__|\.pyc|\.pyo)" | xargs rm -rf; vllm-project-compressed-tensors-c18a0fa/README.md000066400000000000000000000070641521257237700217530ustar00rootroot00000000000000# compressed-tensors The `compressed-tensors` library extends the [safetensors](https://github.com/huggingface/safetensors) format, providing a versatile and efficient way to store and manage compressed tensor data. This library supports various quantization and sparsity schemes, making it a unified format for handling different model optimizations like GPTQ, AWQ, SmoothQuant, INT8, FP8, SparseGPT, and more. ## Why `compressed-tensors`? As model compression becomes increasingly important for efficient deployment of LLMs, the landscape of quantization and compression techniques has become increasingly fragmented. Each method often comes with its own storage format and loading procedures, making it challenging to work with multiple techniques or switch between them. `compressed-tensors` addresses this by providing a single, extensible format that can represent a wide variety of compression schemes. * **Unified Checkpoint Format**: Supports various compression schemes in a single, consistent format. * **Wide Compatibility**: Works with popular quantization methods like GPTQ, SmoothQuant, and FP8. See [llm-compressor](https://github.com/vllm-project/llm-compressor) * **Flexible Quantization Support**: * Weight-only quantization (e.g., W4A16, W8A16, WnA16) * Activation quantization (e.g., W8A8) * KV cache quantization * Non-uniform schemes (different layers can be quantized in different ways!) * **Sparsity Support**: Handles both unstructured and semi-structured (e.g., 2:4) sparsity patterns. * **Open-Source Integration**: Designed to work seamlessly with Hugging Face models and PyTorch. This allows developers and researchers to easily experiment with composing different quantization methods, simplify model deployment pipelines, and reduce the overhead of supporting multiple compression formats in inference engines. ## Installation ### From [PyPI](https://pypi.org/project/compressed-tensors) Stable release: ```bash pip install compressed-tensors ``` Nightly release: ```bash pip install --pre compressed-tensors ``` ### From Source ```bash git clone https://github.com/vllm-project/compressed-tensors cd compressed-tensors pip install -e . ``` ## Getting started ## Saving a Compressed Model with PTQ We can use compressed-tensors to run basic post training quantization (PTQ) and save the quantized model compressed on disk ```python model_name = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T" model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cuda:0", torch_dtype="auto") config = QuantizationConfig.parse_file("./examples/bit_packing/int4_config.json") config.quantization_status = QuantizationStatus.CALIBRATION apply_quantization_config(model, config) dataset = load_dataset("ptb_text_only")["train"] tokenizer = AutoTokenizer.from_pretrained(model_name) def tokenize_function(examples): return tokenizer(examples["sentence"], padding=False, truncation=True, max_length=1024) tokenized_dataset = dataset.map(tokenize_function, batched=True) data_loader = DataLoader(tokenized_dataset, batch_size=1, collate_fn=DefaultDataCollator()) with torch.no_grad(): for idx, sample in tqdm(enumerate(data_loader), desc="Running calibration"): sample = {key: value.to(device) for key,value in sample.items()} _ = model(**sample) if idx >= 512: break model.apply(freeze_module_quantization) model.apply(compress_quantized_weights) output_dir = "./ex_llama1.1b_w4a16_packed_quantize" compressor = ModelCompressor.from_pretrained_model(model) compressor.compress_model(model) model.save_pretrained(output_dir) ``` vllm-project-compressed-tensors-c18a0fa/examples/000077500000000000000000000000001521257237700223035ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/examples/bit_packing/000077500000000000000000000000001521257237700245555ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/examples/bit_packing/ex_quantize_and_pack.py000066400000000000000000000047461521257237700313160ustar00rootroot00000000000000#### # # The following example shows how to run QDQ inside `compressed-tensors` # QDQ (quantize & de-quantize) is a way to evaluate quantized model # accuracy but will not lead to a runtime speedup. # See `../llama_1.1b/ex_config_quantization.py` to go beyond QDQ # and quantize models that will run more performantly. # #### # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from pathlib import Path import torch from compressed_tensors.compressors import ModelCompressor from compressed_tensors.quantization import ( QuantizationConfig, QuantizationStatus, apply_quantization_config, ) from datasets import load_dataset from torch.utils.data import DataLoader from tqdm import tqdm from transformers import AutoModelForCausalLM, AutoTokenizer config_file = Path(__file__).parent / "int4_config.json" model_name = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T" dataset_name = "garage-bAInd/Open-Platypus" split = "train" num_calibration_samples = 128 max_seq_length = 512 pad_to_max_length = False output_dir = "./llama1.1b_new_quant_out_test_packing" device = "cuda:0" if torch.cuda.is_available() else "cpu" model = AutoModelForCausalLM.from_pretrained( model_name, device_map=device, torch_dtype="auto" ) model.eval() # no grad or updates needed for base model config = QuantizationConfig.model_validate_json(config_file.read_text()) # set status to calibration config.quantization_status = QuantizationStatus.CALIBRATION # initialize quantization apply_quantization_config(model, config) # create dataset dataset = load_dataset(dataset_name, split=f"train[:{num_calibration_samples}]") tokenizer = AutoTokenizer.from_pretrained(model_name) def tokenize_function(examples): return tokenizer( examples["output"], padding=False, truncation=True, max_length=1024 ) tokenized_dataset = dataset.map(tokenize_function, batched=True) data_loader = DataLoader( tokenized_dataset, batch_size=1, ) with torch.no_grad(): for idx, sample in tqdm(enumerate(data_loader), desc="Running calibration"): sample = {k: v.to(model.device) for k, v in sample.items()} _ = model(**sample) if idx >= num_calibration_samples: break # convert model to QDQ model compressor = ModelCompressor(quantization_config=config) compressed_state_dict = compressor.compress(model) # save QDQ model model.save_pretrained(output_dir, state_dict=compressed_state_dict) compressor.update_config(output_dir) vllm-project-compressed-tensors-c18a0fa/examples/bit_packing/int4_config.json000066400000000000000000000005641521257237700276600ustar00rootroot00000000000000{ "quant_method": "compressed-tensors", "format": "pack-quantized", "global_compression_ratio": null, "config_groups": { "group_1": { "weights": { "num_bits": 4, "type": "int", "symmetric": false, "strategy": "tensor" }, "targets": ["Linear"] } } }vllm-project-compressed-tensors-c18a0fa/examples/convert_checkpoint/000077500000000000000000000000001521257237700261725ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/examples/convert_checkpoint/autoawq_example.py000066400000000000000000000007771521257237700317530ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from compressed_tensors.entrypoints.convert import AutoAWQConverter, convert_checkpoint MODEL_ID = "Qwen/Qwen2.5-14B-Instruct-AWQ" SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] # Convert AutoAWQ GEMM tensors into compressed-tensors W4A16 format. convert_checkpoint( model_stub=MODEL_ID, save_directory=SAVE_DIR, converter=AutoAWQConverter.from_pretrained(MODEL_ID), max_workers=8, ) vllm-project-compressed-tensors-c18a0fa/examples/convert_checkpoint/deepseek32_fpblock_example.py000066400000000000000000000020311521257237700337050ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from compressed_tensors.entrypoints.convert import ( convert_checkpoint, FP8BlockDequantizer, ) # deepseek-ai/DeepSeek-V3.2 checkpoint has layers that are quantized in the FP8 # quant method's FP8_BLOCK scheme. This script will upconvert to bfloat16 so that # the model can be compressed in another configuration. MODEL_ID = "deepseek-ai/DeepSeek-V3.2" SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-bf16" # Convert DeepSeek-V3.2 back to dense bfloat16 format convert_checkpoint( model_stub=MODEL_ID, save_directory=SAVE_DIR, converter=FP8BlockDequantizer( # `deepseek-ai/DeepSeek-V3.2` fp8-block-quantized layers, found by inspection targets=[ r"re:.*mlp.*\.(gate_up|gate|up|down)_proj$", r"re:.*self_attn.*\.(kv_b|o|q_a|q_b)_proj$", r"re:.*self_attn.kv_a_proj_with_mqa$", r"re:.*self_attn.indexer.(wk|wq_b)$", ], ), max_workers=4, ) vllm-project-compressed-tensors-c18a0fa/examples/convert_checkpoint/kimi_k26_example.py000066400000000000000000000016001521257237700316670ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from compressed_tensors.entrypoints.convert import ( convert_checkpoint, CompressedTensorsDequantizer, ) # moonshotai/Kimi-K2.6 checkpoint is published in compressed-tensors format. # This script will upconvert to bfloat16 so that the model can be compressed # in another configuration. MODEL_ID = "moonshotai/Kimi-K2.6" SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-bf16" convert_checkpoint( model_stub=MODEL_ID, save_directory=SAVE_DIR, converter=CompressedTensorsDequantizer( MODEL_ID, ignore=[ "re:.*mlp.gate$", "re:.*lm_head", "re:.*embed_tokens$", # ignore anything not in language_model "re:.*mm_projector.*", "re:.*vision.*", ], ), max_workers=1, ) vllm-project-compressed-tensors-c18a0fa/examples/convert_checkpoint/qwen3_fpblock_example.py000066400000000000000000000014011521257237700330100ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from compressed_tensors.entrypoints.convert import ( convert_checkpoint, FP8BlockDequantizer, ) MODEL_ID = "qwen-community/Qwen3-4B-FP8" SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1].rstrip("-FP8") # Convert Qwen3-4B-FP8 back to dense bfloat16 format convert_checkpoint( model_stub=MODEL_ID, save_directory=SAVE_DIR, converter=FP8BlockDequantizer( # qwen-community/Qwen3-4B-FP8's fp8-block-quantized layers, found by inspection targets=[ r"re:.*mlp.*\.(gate_up|gate|up|down)_proj$", r"re:.*self_attn.*\.(q|k|v|o)_proj$", ], weight_block_size=[128, 128], ), max_workers=8, ) vllm-project-compressed-tensors-c18a0fa/examples/convert_checkpoint/qwen3_nvfp4_example.py000066400000000000000000000017321521257237700324340ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from compressed_tensors.entrypoints.convert import ( convert_checkpoint, ModelOptNvfp4Converter, ) from compressed_tensors.quantization import QuantizationArgs, QuantizationType MODEL_ID = "nvidia/Qwen3-32B-NVFP4" SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] # Convert modelopt nvfp4 example to compressed-tensors format convert_checkpoint( model_stub=MODEL_ID, save_directory=SAVE_DIR, converter=ModelOptNvfp4Converter( # nvidia/Qwen3-32B-NVFP4's nvfp4-quantized layers, found by inspection targets=[ "re:.*mlp.*\.(gate_up|gate|up|down)_proj$", "re:.*self_attn.*\.(q|k|v|o)_proj$", ], # nvidia/Qwen3-32B-NVFP4's kv_cache_scheme, found by inspection kv_cache_scheme=QuantizationArgs( num_bits=8, dynamic=False, type=QuantizationType.FLOAT ), ), max_workers=8, ) vllm-project-compressed-tensors-c18a0fa/examples/llama_1.1b/000077500000000000000000000000001521257237700241125ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/examples/llama_1.1b/ex_config_quantization.py000066400000000000000000000101301521257237700312260ustar00rootroot00000000000000#### # # The following example shows how a model can be calibrated and # compressed entirely with primitives within `compressed-tensors` # using PyTorch hooks. # The resulting model's .safetensors file should be 1.2GB, # whereas the original model's .safetensors file is 4.1GB. # See `./ex_llmcompressor_quantization.py` for how this can be # simplified using the vllm's `llm-compressor` package # #### # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from pathlib import Path import torch from compressed_tensors.compressors import ModelCompressor from compressed_tensors.quantization import ( QuantizationConfig, QuantizationStatus, apply_quantization_config, ) from datasets import load_dataset from torch.utils.data import DataLoader, RandomSampler from tqdm import tqdm from transformers import AutoModelForCausalLM, AutoTokenizer, DefaultDataCollator config_file = Path(__file__).parent / "example_quant_config.json" model_name = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T" dataset_name = "garage-bAInd/Open-Platypus" split = "train" num_calibration_samples = 512 max_seq_length = 1024 pad_to_max_length = False output_dir = "./llama1.1b_new_quant_out" device = "cuda:0" if torch.cuda.is_available() else "cpu" model = AutoModelForCausalLM.from_pretrained( model_name, device_map=device, torch_dtype="auto" ) model.eval() # no grad or updates needed for base model config = QuantizationConfig.model_validate_json(config_file.read_text()) # set status to calibration config.quantization_status = QuantizationStatus.CALIBRATION # initialize quantization apply_quantization_config(model, config) # create hook to keep track of scales and zero points on each module with a quantization_scheme def update_scale_zp_hook( module: torch.nn.Module, input: torch.Tensor, _output: torch.Tensor ): from compressed_tensors.quantization.utils import calculate_qparams from compressed_tensors.offload import update_offload_parameter quantization_scheme = getattr(module, "quantization_scheme", None) if not quantization_scheme: # no quantization scheme nothing to do return # update weight scale / zero-point quantization_args = getattr(quantization_scheme, "weights", None) min_val, max_val = torch.aminmax(module.weight.data) scale, zp = calculate_qparams(min_val, max_val, quantization_args) update_offload_parameter(module, "weight_scale", scale) update_offload_parameter(module, "weight_zero_point", zp) # update input_activations scale / zero-point quantization_args = getattr(quantization_scheme, "input_activations", None) min_val, max_val = torch.aminmax(input[0]) scale, zp = calculate_qparams(min_val, max_val, quantization_args) update_offload_parameter(module, "input_scale", scale) update_offload_parameter(module, "input_zero_point", zp) return # register hook on each submodule in model (recursively) model.apply(lambda module: module.register_forward_hook(update_scale_zp_hook)) # create dataset dataset = load_dataset(dataset_name, split=f"train[:{num_calibration_samples}]") tokenizer = AutoTokenizer.from_pretrained(model_name) def tokenize_function(examples): return tokenizer( examples["output"], padding=False, truncation=True, max_length=1024 ) tokenized_dataset = dataset.map(tokenize_function, batched=True) data_loader = DataLoader( tokenized_dataset, batch_size=1, collate_fn=DefaultDataCollator(), sampler=RandomSampler(tokenized_dataset), ) # run calibration, hook will update scales and zero points where applicable with torch.no_grad(): for idx, sample in tqdm(enumerate(data_loader), desc="Running calibration"): sample = {k: v.to(model.device) for k, v in sample.items()} _ = model(**sample) if idx >= num_calibration_samples: break # apply compression compressor = ModelCompressor(quantization_config=config) compressed_state_dict = compressor.compress(model) # save quantized model model.save_pretrained(output_dir, state_dict=compressed_state_dict) compressor.update_config(output_dir) vllm-project-compressed-tensors-c18a0fa/examples/llama_1.1b/ex_llmcompressor_quantization.py000066400000000000000000000021011521257237700326610ustar00rootroot00000000000000#### # # The following example shows how the example in `ex_config_quantization.py` # can be done within vllm's llm-compressor project # Be sure to `pip install llmcompressor` before running # See https://github.com/vllm-project/llm-compressor for more information # #### # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from pathlib import Path import torch from llmcompressor import oneshot recipe = str(Path(__file__).parent / "example_quant_recipe.yaml") model_name = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T" dataset_name = "open_platypus" split = "train" num_calibration_samples = 512 max_seq_length = 1024 pad_to_max_length = False output_dir = "./llama1.1b_llmcompressor_quant_out" device = "cuda:0" if torch.cuda.is_available() else "cpu" oneshot( model=model_name, dataset=dataset_name, output_dir=output_dir, overwrite_output_dir=True, max_seq_length=max_seq_length, num_calibration_samples=num_calibration_samples, recipe=recipe, pad_to_max_length=pad_to_max_length, ) vllm-project-compressed-tensors-c18a0fa/examples/llama_1.1b/example_quant_config.json000066400000000000000000000010541521257237700311750ustar00rootroot00000000000000{ "quant_method": "compressed-tensors", "format": "naive-quantized", "global_compression_ratio": null, "config_groups": { "group_1": { "weights": { "num_bits": 8, "type": "int", "symmetric": true, "strategy": "tensor" }, "input_activations": { "num_bits": 8, "type": "int", "symmetric": true, "strategy": "tensor" }, "targets": ["Linear"] } } }vllm-project-compressed-tensors-c18a0fa/examples/llama_1.1b/example_quant_recipe.yaml000066400000000000000000000021261521257237700311710ustar00rootroot00000000000000test_stage: quant_modifiers: QuantizationModifier: ignore: - model.layers.0.mlp.down_proj - LlamaRotaryEmbedding - LlamaRMSNorm - SiLU - MatMulLeftInput_QK - MatMulRightInput_QK - MatMulOutput_QK - MatMulLeftInput_PV - MatMulRightInput_PV - MatMulOutput_PV scheme_overrides: Linear: weights: num_bits: 8 symmetric: true strategy: "tensor" input_activations: num_bits: 8 symmetric: false strategy: "tensor" output_activations: null Embedding: weights: num_bits: 8 symmetric: true strategy: "tensor" input_activations: null output_activations: nullvllm-project-compressed-tensors-c18a0fa/pyproject.toml000066400000000000000000000011121521257237700233740ustar00rootroot00000000000000[build-system] requires = ["setuptools", "wheel", "setuptools_scm==8.2.0"] build-backend = "setuptools.build_meta" [tool.black] line-length = 88 target-version = ['py36'] [tool.pytest.ini_options] markers = [ "unit: tests to ensure code correctness and regression test functionality", "smoke: quick tests to check basic functionality", "sanity: tests to ensure that new changes do not break existing functionality", "regression: detailed tests to ensure major functions work correctly", "integration: tests which integrate with a third party service such as HF", ]vllm-project-compressed-tensors-c18a0fa/pytest-xpu.ini000066400000000000000000000010131521257237700233230ustar00rootroot00000000000000[pytest] addopts = -ra -m "not skip_xpu" testpaths = tests/test_offload/cache/test_cpu.py tests/test_offload/cache/test_device.py tests/test_offload/cache/test_disk.py tests/test_offload/cache/test_dist_cpu.py tests/test_offload/cache/test_dist_device.py tests/test_offload/cache/test_dist_disk.py tests/test_offload/test_dispatch.py tests/test_offload/test_interface.py tests/test_offload/test_module.py markers = skip_xpu: mark test to be skipped on xpu unit: unit test marker vllm-project-compressed-tensors-c18a0fa/setup.cfg000066400000000000000000000006071521257237700223110ustar00rootroot00000000000000[isort] profile = black default_section = FIRSTPARTY ensure_newline_before_comments = True force_grid_wrap = 0 include_trailing_comma = True sections = FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER skip = src/compressed_tensors/version.py line_length = 88 lines_after_imports = 2 multi_line_output = 3 use_parentheses = True [flake8] ignore = E203, E251, E701, W503 max-line-length = 88 vllm-project-compressed-tensors-c18a0fa/setup.py000066400000000000000000000072731521257237700222100ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os from setuptools import setup, find_packages from typing import List, Dict, Tuple # Set the build type using an environment variable to give us # different package names based on the reason for the build. VALID_BUILD_TYPES = {"release", "nightly", "dev"} BUILD_TYPE = os.environ.get("BUILD_TYPE", "dev") if BUILD_TYPE not in VALID_BUILD_TYPES: raise ValueError( f"Unsupported build type {BUILD_TYPE!r}, must be one of {VALID_BUILD_TYPES}" ) from setuptools_scm import ScmVersion def version_func(version: ScmVersion) -> str: from setuptools_scm.version import guess_next_version if BUILD_TYPE == "nightly": # Nightly builds use alpha versions to ensure they are marked # as pre-releases on pypi.org. return version.format_next_version( guess_next=guess_next_version, fmt="{guessed}.a{node_date:%Y%m%d}", ) if ( BUILD_TYPE == "release" and not version.dirty and (version.exact or version.node is None) ): # When we have a tagged version, use that without modification. return version.format_with("{tag}") # In development mode or when the local repository is dirty, treat # it is as local development version. return version.format_next_version( guess_next=guess_next_version, fmt="{guessed}.dev{distance}", ) def localversion_func(version: ScmVersion) -> str: from setuptools_scm.version import get_local_node_and_date # When we are building nightly versions, we guess the next release # and add the date as an alpha version. We cannot publish packages # with local versions, so we do not add one. if BUILD_TYPE == "nightly": return "" # When we have an exact tag, with no local changes, do not append # anything to the local version field. if ( BUILD_TYPE == "release" and not version.dirty and (version.exact or version.node is None) ): return "" # In development mode or when the local repository is dirty, # return a string that includes the git SHA (node) and a date, # formatted as a local version tag. return get_local_node_and_date(version) def _setup_long_description() -> Tuple[str, str]: return open("README.md", "r", encoding="utf-8").read(), "text/markdown" def _setup_packages() -> List: return find_packages( "src", include=["compressed_tensors", "compressed_tensors.*"], exclude=["*.__pycache__.*"] ) def _setup_install_requires() -> List: return ["torch>=2.10.0", "transformers>=4.45.0", "pydantic>=2.0", "loguru"] def _setup_extras() -> Dict: return { "dev": ["black==22.12.0", "isort==5.8.0", "wheel>=0.36.2", "flake8>=3.8.3", "pytest>=6.0.0", "nbconvert>=7.16.3", "accelerate"], "accelerate": ["accelerate"] } setup( name="compressed-tensors", use_scm_version={ "version_scheme": version_func, "local_scheme": localversion_func, "version_file": "src/compressed_tensors/version.py", }, author="The vLLM Project", author_email="vllm-questions@lists.berkeley.edu", license="Apache 2.0", description="Library for utilization of compressed safetensors of neural network models", long_description=_setup_long_description()[0], long_description_content_type=_setup_long_description()[1], url="https://github.com/vllm-project/compressed-tensors", extras_require=_setup_extras(), install_requires=_setup_install_requires(), package_dir={"": "src"}, package_data={"": ["transform/utils/hadamards.safetensors"]}, packages=_setup_packages(), ) vllm-project-compressed-tensors-c18a0fa/src/000077500000000000000000000000001521257237700212545ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/__init__.py000066400000000000000000000001531521257237700233640ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/000077500000000000000000000000001521257237700251755ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/__init__.py000066400000000000000000000012741521257237700273120ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa # isort: off from .logger import LoggerConfig, configure_logger, logger from .base import * from .compressors import * from .config import * from .quantization import QuantizationConfig, QuantizationStatus # avoid resolving compressed_tensors.offload as compressed_tensors.utils.offload from .utils.offload import * from .utils.helpers import * from .utils.internal import * from .utils.match import * from .utils.permutations_24 import * from .utils.safetensors_load import * from .utils.semi_structured_conversions import * from .utils.type import * from .version import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/base.py000066400000000000000000000006101521257237700264560ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # configs QUANTIZATION_CONFIG_NAME = "quantization_config" SPARSITY_CONFIG_NAME = "sparsity_config" TRANSFORM_CONFIG_NAME = "transform_config" # required fields COMPRESSION_VERSION_NAME = "version" QUANTIZATION_METHOD_NAME = "quant_method" QUANTIZATION_METHOD = "compressed-tensors" vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/000077500000000000000000000000001521257237700275545ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/__init__.py000066400000000000000000000005401521257237700316640ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .base import * # New per-format directories from .dense import * from .model_compressors import * from .mxfp4 import * from .mxfp8 import * from .naive_quantized import * from .nvfp4 import * from .pack_quantized import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/base.py000066400000000000000000000207151521257237700310450ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC from typing import Optional import torch from compressed_tensors.compressors.format import infer_module_format from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import QuantizationScheme, QuantizationStatus from compressed_tensors.registry import RegistryMixin from compressed_tensors.utils import ( TensorStateDict, get_direct_state_dict, replace_direct_state_dict, ) __all__ = [ "BaseCompressor", "compress_module", "decompress_module", "COMPRESSIBLE_MODULE_TYPES", ] # Module types whose weights can be compressed/decompressed. Compression operates # on a per-module weight state dict and is agnostic to the module's forward, so any # module that stores a quantizable 2D `weight` works. Quantization of these types is # initialized in `initialize_module_for_quantization`. COMPRESSIBLE_MODULE_TYPES = (torch.nn.Linear, torch.nn.Embedding) class BaseCompressor(RegistryMixin, ABC): """ Base class representing a model compression algorithm. New quantization compressors (dense, naive_quantized, pack_quantized, nvfp4, mxfp4) use the classmethod interface — they are never instantiated. Look up via BaseCompressor.get_value_from_registry(format) and call compress/decompress directly on the returned class. Legacy sparse compressors (sparse_bitmask, sparse_24_bitmask, marlin_24) still use the instance-based interface and are instantiated via load_from_registry. """ @classmethod def compression_param_names(cls, scheme: QuantizationScheme) -> tuple[str]: """ Returns a tuple of compression parameter names introduced by the compressor during compression. This is necessary so that the state dict can be recreated to pass into decompress in model-free pathways. See example usage in compressed_tensors.entrypoints.convert.converters.CompressedTensorsDequantizer """ raise NotImplementedError( f"{cls.__name__} does not implement the classmethod " "compression_param_names interface" ) @classmethod def compress( cls, state_dict: TensorStateDict, scheme: QuantizationScheme ) -> TensorStateDict: """ Compress a per-module state dict. Does not modify the input. Keys are *local* names (``weight``, ``weight_scale``, …), not prefixed with the module path. :param state_dict: per-module state dict with local parameter names :param scheme: quantization scheme containing quantization parameters :return: compressed per-module state dict """ raise NotImplementedError( f"{cls.__name__} does not implement the classmethod compress interface" ) @classmethod def decompress( cls, state_dict: TensorStateDict, scheme: QuantizationScheme ) -> TensorStateDict: """ Decompress a per-module state dict. Does not modify the input. Keys are *local* names (``weight_packed``, ``weight_scale``, …). :param state_dict: compressed per-module state dict with local parameter names :param scheme: quantization scheme containing quantization parameters :return: decompressed per-module state dict """ raise NotImplementedError( f"{cls.__name__} does not implement the classmethod decompress interface" ) @classmethod def compress_module(cls, module: torch.nn.Module) -> None: """ Compress a module in-place by compressing its state dict. Extracts the module's parameters and buffers, compresses them using the compress classmethod, and replaces the module's state with the compressed version. :param module: the module to compress in-place """ scheme = getattr(module, "quantization_scheme") state_dict = get_direct_state_dict(module) compressed_state_dict = cls.compress(state_dict, scheme) replace_direct_state_dict(module, compressed_state_dict) module.quantization_status = QuantizationStatus.COMPRESSED @classmethod def decompress_module(cls, module: torch.nn.Module) -> None: """ Decompress a module in-place by decompressing its state dict. Extracts the module's parameters and buffers, decompresses them using the decompress classmethod, and replaces the module's state with the decompressed version. :param module: the module to decompress in-place """ scheme = getattr(module, "quantization_scheme") state_dict = get_direct_state_dict(module) decompressed_state_dict = cls.decompress(state_dict, scheme) replace_direct_state_dict(module, decompressed_state_dict) module.quantization_status = QuantizationStatus.DECOMPRESSED @classmethod def can_compress(cls, module_type: type, scheme: QuantizationScheme) -> bool: """ Determine if this compressor is applicable for the given module type and scheme. Examines the module type and quantization scheme and determines whether this compressor can handle the module's compression requirements. :param module_type: the type of the module to check for compatibility :param scheme: the quantization scheme to check for compatibility :return: True if this compressor can handle the module, False otherwise """ raise NotImplementedError(f"{cls.__name__} does not implement match") @classmethod def _remove_symmetric_zp( cls, state_dict: TensorStateDict, scheme: QuantizationScheme ) -> TensorStateDict: """ Remove zero points where quantization arguments do not specify asymmetric quant. This is required during compression because vLLM does not support loading zero point parameters for symmetric schemes. TODO: remove zero points from initialization :param state_dict: state dict containing extra zero point parameters :return: state dict with extra zero point parameters removed """ if scheme.input_activations and scheme.input_activations.symmetric: state_dict.pop("input_zero_point", None) if scheme.weights and scheme.weights.symmetric: state_dict.pop("weight_zero_point", None) if scheme.output_activations and scheme.output_activations.symmetric: state_dict.pop("output_zero_point", None) return state_dict def compress_module( module: torch.nn.Module, format: Optional[CompressionFormat] = None ): """ Compress a module which has had quantization applied to it. Sets the module's quantization format attribute to whichever format was used to compress. The format used to compress will be found from one of the following locations: 1. the `format` argument passed to this function 2. the attached `quantization_scheme.format` attribute 3. format inferred from `infer_module_format` :param module: module to compress inplace :param format: force override for compression format """ scheme = getattr(module, "quantization_scheme", None) if not isinstance(scheme, QuantizationScheme): return scheme.format = CompressionFormat( format or scheme.format or infer_module_format(type(module), scheme) ) compressor = BaseCompressor.get_value_from_registry(scheme.format.value) compressor.compress_module(module) def decompress_module( module: torch.nn.Module, format: Optional[CompressionFormat] = None ): """ Decompress a module which has had quantization applied to it. Sets the module's quantization format attribute to whichever format was used to decompress. The format used to decompress will be found from one of the following locations: 1. the `format` argument passed to this function 2. the attached `quantization_scheme.format` attribute 3. format inferred from `infer_module_format` :param module: module to decompress inplace :param format: force override for decompression format """ scheme = getattr(module, "quantization_scheme", None) if not isinstance(scheme, QuantizationScheme): return scheme.format = CompressionFormat( format or scheme.format or infer_module_format(type(module), scheme) ) compressor = BaseCompressor.get_value_from_registry(scheme.format.value) compressor.decompress_module(module) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/dense/000077500000000000000000000000001521257237700306525ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/dense/__init__.py000066400000000000000000000002201521257237700327550ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .base import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/dense/base.py000066400000000000000000000034611521257237700321420ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from compressed_tensors.compressors.base import BaseCompressor from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import QuantizationScheme from compressed_tensors.utils import TensorStateDict __all__ = ["DenseCompressor"] @BaseCompressor.register(name=CompressionFormat.dense.value) class DenseCompressor(BaseCompressor): """ Identity compressor for dense models — both compress and decompress return the state dict unchanged. """ @classmethod def compression_param_names(cls, scheme: QuantizationScheme) -> tuple[str]: """ DenseCompressor adds no new qparams """ return ("weight",) @classmethod def compress( cls, state_dict: TensorStateDict, scheme: QuantizationScheme ) -> TensorStateDict: """ Compress a per-module state dict. For dense models, this is a no-op. :param state_dict: local-name state dict (weight, ...) :param scheme: quantization scheme (unused) :return: unmodified state dict """ return state_dict @classmethod def decompress( cls, state_dict: TensorStateDict, scheme: QuantizationScheme ) -> TensorStateDict: """ Decompress a per-module state dict. For dense models, this is a no-op. :param state_dict: local-name state dict (weight, ...) :param scheme: quantization scheme (unused) :return: unmodified state dict """ return state_dict @classmethod def can_compress(cls, module_type: type, scheme: QuantizationScheme) -> bool: """Dense compressor matches when there's no weight quantization.""" return True vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/format.py000066400000000000000000000074131521257237700314230ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import List, Optional import torch from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import QuantizationScheme from compressed_tensors.quantization.utils import is_module_quantized from loguru import logger __all__ = ["infer_model_format", "infer_module_format"] # Priority order for compression format matching # More specific formats should come before more general ones COMPRESSION_FORMAT_PRIORITY: List[CompressionFormat] = [ CompressionFormat.mxfp4_pack_quantized, CompressionFormat.mxfp8_quantized, CompressionFormat.nvfp4_pack_quantized, CompressionFormat.pack_quantized, CompressionFormat.int_quantized, CompressionFormat.float_quantized, CompressionFormat.naive_quantized, CompressionFormat.dense, ] def infer_model_format( model: torch.nn.Module, force_compression_format: Optional[str] = None, ) -> CompressionFormat: """ Infers the quantization format for a model based on its modules For a summary of the formats, see `docs/guides/compression_formats.md`. :param model: model to check for quantization :param quantization_format: optional global format to override the per module formats :return: list of formats applied to modules (excluding dense format) """ formats = set() for _, module in model.named_modules(remove_duplicate=True): if not is_module_quantized(module): continue # infer format using priority list scheme: QuantizationScheme = module.quantization_scheme format = infer_module_format(type(module), scheme) # user provides a global override format if force_compression_format is not None: if force_compression_format != format.value: logger.warning( f"The provided format {force_compression_format} does not match " f"the inferred format {format.value}. Compression may fail", log_once=True, ) format = force_compression_format # user provides a format via QuantizationScheme.format elif scheme.format is not None: format = scheme.format scheme.format = CompressionFormat(format) if format != CompressionFormat.dense: formats.add(format) return _flatten_formats(formats) def infer_module_format( module_type: type, scheme: QuantizationScheme ) -> CompressionFormat: """ Infer the module's compression format using the module's type and quant scheme :param module_type: module type, typically linear :param scheme: quantization applied to module :return: format that should be used to compress the module """ # avoid circular imports from compressed_tensors.compressors import BaseCompressor return next( ( format for format in COMPRESSION_FORMAT_PRIORITY if BaseCompressor.get_value_from_registry(format.value).can_compress( module_type, scheme ) ) ) def _flatten_formats(formats: set[CompressionFormat]) -> CompressionFormat: """ Reduce a list of compression formats to a single summary format. Returns dense if the list is empty, the single format if there's only one, or mixed_precision if there are multiple different formats. :param formats: list of compression formats found in the model :return: single compression format representing the overall model """ if len(formats) <= 0: return CompressionFormat.dense if len(formats) == 1: return list(formats)[0] if len(formats) >= 2: return CompressionFormat.mixed_precision vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/model_compressors/000077500000000000000000000000001521257237700333135ustar00rootroot00000000000000__init__.py000066400000000000000000000002351521257237700353450ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/model_compressors# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .model_compressor import * model_compressor.py000066400000000000000000000241751521257237700371730ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/model_compressors# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import os from functools import partial from typing import Optional import compressed_tensors import torch from compressed_tensors.base import ( COMPRESSION_VERSION_NAME, QUANTIZATION_CONFIG_NAME, QUANTIZATION_METHOD, QUANTIZATION_METHOD_NAME, SPARSITY_CONFIG_NAME, TRANSFORM_CONFIG_NAME, ) from compressed_tensors.compressors.base import compress_module, decompress_module from compressed_tensors.compressors.format import infer_model_format from compressed_tensors.config import CompressionFormat from compressed_tensors.distributed import replace_module_parallel from compressed_tensors.offload import is_distributed from compressed_tensors.quantization import QuantizationConfig, QuantizationStatus from compressed_tensors.quantization.utils.helpers import is_module_quantized from compressed_tensors.transform import TransformConfig from loguru import logger from tqdm import tqdm from transformers import CompressedTensorsConfig from transformers.file_utils import CONFIG_NAME __all__ = ["ModelCompressor"] class ModelCompressor: """ Orchestrates compression and decompression of a quantized model. Compression Lifecycle - model.save_pretrained_wrapper(quantized_path) - compressor = ModelCompressor.from_pretrained_model(model) - compressor.compress_model(model) - model.save_pretrained(quantized_path) - compressor.update_config(quantized_path) Decompression Lifecycle - model = AutoModelForCausalLM.from_pretrained(quantized_path) - CompressedTensorsHfQuantizer.__init__ - compressor = ModelCompressor.from_compression_config(ct_config) - CompressedTensorsHfQuantizer._process_model_before_weight_loading - apply_quantization_config(model, ct_config.quantization_config) - compressor.compress_model(model) - CompressedTensorsHfQuantizer._process_model_after_weight_loading - if run_compressed == False: compressor.decompress_model(model) """ # these attributes are used by `CompressedTensorsHfQuantizer` to apply configs # during `_process_model_before_weight_loading` quantization_config: QuantizationConfig | None transform_config: TransformConfig | None force_compression_format: CompressionFormat | None @classmethod def from_compression_config(cls, compression_config: CompressedTensorsConfig): """ HFQuantizer uses this entrypoint to load quantized models by passing an instance of `CompressedTensorsConfig` :param compression_config: instance of `CompressedTensorsConfig` containing quantized model information """ # vLLM Cutlass24 implementation is no longer supported if not isinstance(compression_config, CompressedTensorsConfig): raise ValueError( f"Support for compression config of type {type(compression_config)} " "is no longer supported. If you are attempting to use a Sparse24 " "model, note that the Sparse24 format is not longer supported by " "as of `compressed-tensors>0.14.0`" ) # transform_config is added by transformers#42887 q_config = compression_config.quantization_config t_config = getattr(compression_config, "transform_config", None) return cls(quantization_config=q_config, transform_config=t_config) @classmethod def from_pretrained_model( cls, model: torch.nn.Module, sparsity_config_or_format: Optional[object] = None, quantization_format: Optional[str] = None, ): """ LLM Compressor uses this entrypoint to compress models before saving Given a path to a model config, extract the sparsity and/or quantization configs and load a ModelCompressor. Note: passing `sparsity_config_or_format` is no longer supported :param model: model with quantization config applied :param quantization_format: string corresponding to a quantization format that should be applied to the entire model, overrides inferred formats for all quantized modules """ if sparsity_config_or_format is not None: logger.warning("Passing sparsity config or format is no longer supported") # reconstruct qconfig from qschemes that are attached to the model quantization_config = QuantizationConfig.from_pretrained(model) transform_config = getattr(model, TRANSFORM_CONFIG_NAME, None) # update quantization config format for better UX when reading config.json if quantization_config is not None: quantization_config.format = infer_model_format(model, quantization_format) return cls( quantization_config=quantization_config, transform_config=transform_config, force_compression_format=quantization_format, ) def __init__( self, quantization_config: Optional[QuantizationConfig] = None, transform_config: Optional[TransformConfig] = None, force_compression_format: Optional[str] = None, ): self.quantization_config = quantization_config self.transform_config = transform_config self.force_compression_format = ( CompressionFormat(force_compression_format) if force_compression_format is not None else None ) def compress_model(self, model: torch.nn.Module) -> None: """ Compress the model's parameters in memory :param model: model whose parameters should be compressed in place """ # Collect all quantized modules desc = "Compressing model" modules = [ module for _, module in model.named_modules(remove_duplicate=True) if is_module_quantized(module) ] # Compress modules using distributed or sequential if not is_distributed(): for module in tqdm(modules, desc=desc): compress_module(module, self.force_compression_format) else: compress_fn = partial(compress_module, format=self.force_compression_format) replace_module_parallel(modules, compress_fn, desc=desc) # update config status to reflect compression if self.quantization_config is not None: self.quantization_config.quantization_status = QuantizationStatus.COMPRESSED # attempting to perform forward passes with a compressed model # will cause to the model to decompress. This allows for generation # without requiring CT to handle quantized kernels self.add_decompress_hook(model) def decompress_model(self, model: torch.nn.Module) -> None: """ Decompress the model's parameters in memory :param model: model whose parameters should be decompressed in place """ desc = "Decompressing model" modules = [ module for _, module in model.named_modules(remove_duplicate=True) if is_module_quantized(module) ] # TODO: support distributed decompression for module in tqdm(modules, desc=desc): decompress_module(module, self.force_compression_format) # update config status to reflect decompression if self.quantization_config is not None: self.quantization_config.quantization_status = ( QuantizationStatus.DECOMPRESSED ) # decompression hook is no longer necessary self.remove_decompression_hook(model) def update_config(self, save_directory: str) -> None: """ Update the model config located at save_directory with compression configs :param save_directory: path to a folder containing a HF model config """ if not any((self.quantization_config, self.transform_config)): return config_file_path = os.path.join(save_directory, CONFIG_NAME) if os.path.exists(config_file_path): with open(config_file_path, "r") as file: config_data = json.load(file) else: config_data = {} qconfig_data = ( self.quantization_config.model_dump(exclude=["quant_method"]) if self.quantization_config is not None else {} ) tconfig_data = ( self.transform_config.model_dump() if self.transform_config is not None else {} ) config_data[QUANTIZATION_CONFIG_NAME] = { COMPRESSION_VERSION_NAME: compressed_tensors.__version__, QUANTIZATION_METHOD_NAME: QUANTIZATION_METHOD, SPARSITY_CONFIG_NAME: {}, # sparsity is removed for now TRANSFORM_CONFIG_NAME: tconfig_data, **qconfig_data, } with open(config_file_path, "w") as config_file: json.dump(config_data, config_file, indent=2, sort_keys=True) def add_decompress_hook(self, model: torch.nn.Module): """ Register a forward pre-hook that decompresses the model on first forward pass. The hook automatically removes itself after decompression, allowing the model to run in decompressed state for inference. :param model: model to attach the decompression hook to """ def ct_decompress_hook(model, args): self.decompress_model(model) # decompress_model already removes the hook via remove_decompression_hook model.ct_decompress_hook = model.register_forward_pre_hook(ct_decompress_hook) def remove_decompression_hook(self, model: torch.nn.Module): """ Remove the decompression hook from the model if it exists. Called after manual decompression to clean up the hook that would otherwise trigger on the next forward pass. :param model: model to remove the decompression hook from """ if hasattr(model, "ct_decompress_hook"): model.ct_decompress_hook.remove() delattr(model, "ct_decompress_hook") vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/mx_utils.py000066400000000000000000000026431521257237700317770ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ Shared utilities for MX-format (MXFP4/MXFP8) scale compression and decompression. MX scales are stored as E8M0 exponents (uint8), where each value represents a biased power-of-2 exponent with bias 127: ``scale_float = 2 ** (exp - 127)``. """ import torch __all__ = ["compress_mx_scale", "decompress_mx_scale"] def compress_mx_scale(scale: torch.Tensor, scale_dtype: torch.dtype) -> torch.Tensor: """ Convert a float scale tensor to E8M0 exponent format for MX storage. Extracts the power-of-2 exponent from each scale value and adds the E8M0 bias (127). :param scale: float scale tensor (e.g. from observer output) :param scale_dtype: target dtype for the compressed scale (typically uint8) :return: biased exponent tensor in the requested dtype """ scale_exp = 127 + torch.floor(torch.log2(scale)).to(torch.int32) return scale_exp.to(scale_dtype) def decompress_mx_scale(scale: torch.Tensor) -> torch.Tensor: """ Convert E8M0 exponent scale back to float (bfloat16). Reverses the E8M0 encoding by subtracting the bias and raising 2 to the resulting power. :param scale: uint8 tensor of biased E8M0 exponents :return: bfloat16 tensor of float scale values """ scale_exp = scale.to(torch.int32) - 127 return 2.0 ** scale_exp.to(torch.bfloat16) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/mxfp4/000077500000000000000000000000001521257237700306125ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/mxfp4/__init__.py000066400000000000000000000002201521257237700327150ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .base import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/mxfp4/base.py000066400000000000000000000032551521257237700321030ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.compressors.base import ( COMPRESSIBLE_MODULE_TYPES, BaseCompressor, ) from compressed_tensors.compressors.mx_utils import ( compress_mx_scale, decompress_mx_scale, ) from compressed_tensors.compressors.nvfp4.base import NVFP4PackedCompressor from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import ( QuantizationArgs, QuantizationScheme, QuantizationType, ) __all__ = ["MXFP4PackedCompressor"] @BaseCompressor.register(name=CompressionFormat.mxfp4_pack_quantized.value) class MXFP4PackedCompressor(NVFP4PackedCompressor): """ Compressor for MXFP4 quantized models. Overrides scale compression to use log2 encoding (bias-127 exponent). """ @classmethod def _compress_scale( cls, scale: torch.Tensor, weights: QuantizationArgs ) -> torch.Tensor: scale_dtype = weights.scale_dtype or torch.uint8 return compress_mx_scale(scale, scale_dtype) @classmethod def _decompress_scale(cls, scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: return decompress_mx_scale(scale).to(dtype) @classmethod def can_compress(cls, module_type: type, scheme: QuantizationScheme) -> bool: """MXFP4 matches FP4 with group_size=32.""" return ( module_type in COMPRESSIBLE_MODULE_TYPES and scheme.weights is not None and scheme.weights.num_bits == 4 and scheme.weights.type == QuantizationType.FLOAT.value and scheme.weights.group_size == 32 ) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/mxfp8/000077500000000000000000000000001521257237700306165ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/mxfp8/__init__.py000066400000000000000000000002201521257237700327210ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .base import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/mxfp8/base.py000066400000000000000000000071541521257237700321110ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.compressors.base import ( COMPRESSIBLE_MODULE_TYPES, BaseCompressor, ) from compressed_tensors.compressors.mx_utils import ( compress_mx_scale, decompress_mx_scale, ) from compressed_tensors.compressors.naive_quantized.base import ( NaiveQuantizationCompressor, ) from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import ( QuantizationArgs, QuantizationScheme, QuantizationType, ) from compressed_tensors.utils import TensorStateDict __all__ = ["MXFP8QuantizationCompressor"] @BaseCompressor.register(name=CompressionFormat.mxfp8_quantized.value) class MXFP8QuantizationCompressor(NaiveQuantizationCompressor): """ Compressor for MXFP8 quantized models. Weights are stored as float8_e4m3fn with E8M0 (power-of-2) scales stored as uint8 exponents. Extends NaiveQuantizationCompressor by converting float scales to/from E8M0 exponent format during compression/decompression. """ @classmethod def _compress_scale( cls, scale: torch.Tensor, weights: QuantizationArgs ) -> torch.Tensor: scale_dtype = weights.scale_dtype or torch.uint8 return compress_mx_scale(scale, scale_dtype) @classmethod def _decompress_scale(cls, scale: torch.Tensor) -> torch.Tensor: return decompress_mx_scale(scale) @classmethod def compress( cls, state_dict: TensorStateDict, scheme: QuantizationScheme ) -> TensorStateDict: """ Compress a per-module state dict for MXFP8 format. Quantizes the weight via the parent class, then converts the float scale to E8M0 exponent format (uint8) for storage. :param state_dict: local-name state dict (weight, weight_scale, ...) :param scheme: quantization scheme for the weight :return: compressed state dict with E8M0 scales """ state_dict = NaiveQuantizationCompressor.compress(state_dict, scheme) # Convert float scale to E8M0 exponent format (uint8) for storage scale = state_dict["weight_scale"] state_dict["weight_scale"] = cls._compress_scale(scale, scheme.weights) return state_dict @classmethod def decompress( cls, state_dict: TensorStateDict, scheme: QuantizationScheme ) -> TensorStateDict: """ Decompress a per-module state dict for MXFP8 format. Converts E8M0 exponent scales back to float, then dequantizes the weight via the parent class. :param state_dict: local-name state dict (weight, weight_scale, ...) :param scheme: quantization scheme for the weight :return: decompressed state dict with weight in float dtype """ state_dict = state_dict.copy() # Convert E8M0 scale back to bfloat16 for consistency with model dtype scale = state_dict["weight_scale"] state_dict["weight_scale"] = cls._decompress_scale(scale) return NaiveQuantizationCompressor.decompress(state_dict, scheme) @classmethod def can_compress(cls, module_type: type, scheme: QuantizationScheme) -> bool: """MXFP8 matches FP8 with group_size=32 and uint8 scale_dtype.""" return ( module_type in COMPRESSIBLE_MODULE_TYPES and scheme.weights is not None and scheme.weights.num_bits == 8 and scheme.weights.type == QuantizationType.FLOAT.value and scheme.weights.group_size == 32 and scheme.weights.scale_dtype == torch.uint8 ) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/naive_quantized/000077500000000000000000000000001521257237700327425ustar00rootroot00000000000000__init__.py000066400000000000000000000002201521257237700347660ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/naive_quantized# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .base import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/naive_quantized/base.py000066400000000000000000000133011521257237700342240ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from compressed_tensors.compressors.base import ( COMPRESSIBLE_MODULE_TYPES, BaseCompressor, ) from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import ( ActivationOrdering, QuantizationScheme, QuantizationStrategy, QuantizationType, ) from compressed_tensors.quantization.lifecycle.forward import dequantize, quantize from compressed_tensors.quantization.utils import maybe_pad_tensor_for_block_quant from compressed_tensors.utils import TensorStateDict, getattr_chain __all__ = [ "NaiveQuantizationCompressor", "IntQuantizationCompressor", "FloatQuantizationCompressor", ] @BaseCompressor.register(name=CompressionFormat.naive_quantized.value) class NaiveQuantizationCompressor(BaseCompressor): """ Naive quantization compressor. Each quantized layer's weight is converted from its original float dtype to the closest PyTorch dtype for the bit-width specified by QuantizationArgs. """ @classmethod def compression_param_names(cls, scheme: QuantizationScheme) -> tuple[str]: param_names = ( "weight", "weight_scale", ) if not getattr_chain(scheme, "weights.symmetric", True): param_names += ("weight_zero_point",) if getattr_chain(scheme, "weights.actorder", None) == ActivationOrdering.GROUP: param_names += ("weight_g_idx",) return param_names @classmethod def compress( cls, state_dict: TensorStateDict, scheme: QuantizationScheme ) -> TensorStateDict: """ Compress a per-module state dict. Quantizes the weight to the dtype specified by the scheme's QuantizationArgs. Handles block quantization padding if needed. :param state_dict: local-name state dict (weight, weight_scale, …) :param scheme: quantization scheme for the weight :return: compressed state dict """ state_dict = state_dict.copy() weight = state_dict.pop("weight") scale = state_dict.get("weight_scale") zero_point = state_dict.get("weight_zero_point", None) g_idx = state_dict.get("weight_g_idx", None) weights = scheme.weights original_weight_shape = weight.shape # For block quantization, pad weight to divisible dimensions if ( weights.strategy == QuantizationStrategy.BLOCK and weights.block_structure is not None ): block_structure = tuple(weights.block_structure) weight = maybe_pad_tensor_for_block_quant(weight, block_structure) quantized_weight = quantize( x=weight, scale=scale, zero_point=zero_point, g_idx=g_idx, args=weights, dtype=weights.pytorch_dtype(), ) # Truncate back to original shape if padding was added if quantized_weight.shape != original_weight_shape: quantized_weight = quantized_weight[ tuple([slice(v) for v in original_weight_shape]) ] state_dict["weight"] = quantized_weight state_dict = cls._remove_symmetric_zp(state_dict, scheme) return state_dict @classmethod def decompress( cls, state_dict: TensorStateDict, scheme: QuantizationScheme ) -> TensorStateDict: """ Decompress a per-module state dict. Dequantizes the weight back to float dtype using the scale and zero-point from the state dict. :param state_dict: local-name state dict (weight, weight_scale, …) :param scheme: quantization scheme for the weight :return: decompressed state dict with weight in float dtype """ state_dict = state_dict.copy() weight = state_dict.pop("weight") scale = state_dict.get("weight_scale") zero_point = state_dict.get("weight_zero_point", None) g_idx = state_dict.get("weight_g_idx", None) state_dict["weight"] = dequantize( x_q=weight, scale=scale, zero_point=zero_point, g_idx=g_idx, ) return state_dict @classmethod def can_compress(cls, module_type: type, scheme: QuantizationScheme) -> bool: """ Naive quantization is the fallback compressor - it matches any quantized scheme that doesn't match a more specific compressor. """ return module_type in COMPRESSIBLE_MODULE_TYPES and scheme.weights is not None @BaseCompressor.register(name=CompressionFormat.int_quantized.value) class IntQuantizationCompressor(NaiveQuantizationCompressor): """Alias for integer quantized models.""" @classmethod def can_compress(cls, module_type: type, scheme: QuantizationScheme) -> bool: """Int quantized matches w8a8 int quantization.""" return ( module_type in COMPRESSIBLE_MODULE_TYPES and scheme.input_activations is not None and scheme.weights is not None and scheme.weights.type == QuantizationType.INT.value ) @BaseCompressor.register(name=CompressionFormat.float_quantized.value) class FloatQuantizationCompressor(NaiveQuantizationCompressor): """Alias for fp quantized models.""" @classmethod def can_compress(cls, module_type: type, scheme: QuantizationScheme) -> bool: """Float quantized matches w8a8 float quantization.""" return ( module_type in COMPRESSIBLE_MODULE_TYPES and scheme.input_activations is not None and scheme.weights is not None and scheme.weights.type == QuantizationType.FLOAT.value ) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/nvfp4/000077500000000000000000000000001521257237700306115ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/nvfp4/__init__.py000066400000000000000000000002471521257237700327250ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .base import * from .helpers import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/nvfp4/base.py000066400000000000000000000114421521257237700320770ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.compressors.base import ( COMPRESSIBLE_MODULE_TYPES, BaseCompressor, ) from compressed_tensors.compressors.nvfp4.helpers import ( pack_fp4_to_uint8, unpack_fp4_from_uint8, ) from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import ( QuantizationArgs, QuantizationScheme, QuantizationStrategy, QuantizationType, ) from compressed_tensors.quantization.lifecycle.forward import dequantize, quantize from compressed_tensors.utils import TensorStateDict, getattr_chain __all__ = ["NVFP4PackedCompressor"] @BaseCompressor.register(name=CompressionFormat.nvfp4_pack_quantized.value) class NVFP4PackedCompressor(BaseCompressor): """ Compressor for FP4 quantized models. Weights of each quantized layer are packed into uint8. Only supports symmetric weight compression. """ @classmethod def compression_param_names(cls, scheme: QuantizationScheme) -> tuple[str]: param_names = ( "weight_packed", "weight_scale", "weight_global_scale", ) if not getattr_chain(scheme, "weights.symmetric", True): param_names += ("weight_zero_point",) if ( getattr_chain(scheme, "input_activations.strategy", None) == QuantizationStrategy.TENSOR_GROUP ): param_names += ("input_global_scale",) return param_names @classmethod def _compress_scale( cls, scale: torch.Tensor, weights: QuantizationArgs ) -> torch.Tensor: scale_dtype = weights.scale_dtype or torch.float8_e4m3fn return scale.to(scale_dtype) @classmethod def _decompress_scale(cls, scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: return scale.to(dtype) @classmethod def compress( cls, state_dict: TensorStateDict, scheme: QuantizationScheme ) -> TensorStateDict: """ Compress a per-module state dict. Quantizes the weight and packs into uint8 as ``weight_packed``. Compresses the scale according to ``scheme.weights.scale_dtype``. Removes the raw ``weight``. :param state_dict: local-name state dict (weight, weight_scale, …) :param scheme: quantization scheme for the weight :return: compressed state dict """ state_dict = state_dict.copy() weight = state_dict.pop("weight") scale = state_dict.pop("weight_scale") global_scale = state_dict.get("weight_global_scale", None) zero_point = state_dict.get("weight_zero_point", None) weights = scheme.weights quantized_weight = quantize( x=weight, scale=scale, global_scale=global_scale, zero_point=zero_point, args=weights, ) state_dict["weight_packed"] = pack_fp4_to_uint8(quantized_weight) state_dict["weight_scale"] = cls._compress_scale(scale, weights) state_dict = cls._remove_symmetric_zp(state_dict, scheme) return state_dict @classmethod def decompress( cls, state_dict: TensorStateDict, scheme: QuantizationScheme ) -> TensorStateDict: """ Decompress a per-module state dict. Unpacks ``weight_packed`` back to fp4 values and dequantizes. Converts ``weight_scale`` back to float for dequantization. :param state_dict: local-name state dict (weight_packed, weight_scale, …) :param scheme: quantization scheme for the weight :return: decompressed state dict with weight in float dtype """ state_dict = state_dict.copy() packed = state_dict.pop("weight_packed") scale = state_dict.get("weight_scale") global_scale = state_dict.get("weight_global_scale", None) m, n = packed.shape unpacked = unpack_fp4_from_uint8(packed, m, n * 2) scale_float = cls._decompress_scale(scale, unpacked.dtype) state_dict["weight"] = dequantize( x_q=unpacked, scale=scale_float, global_scale=global_scale, dtype=unpacked.dtype, ) state_dict["weight_scale"] = torch.nn.Parameter( scale_float, requires_grad=False ) return state_dict @classmethod def can_compress(cls, module_type: type, scheme: QuantizationScheme) -> bool: """NVFP4 matches FP4 with group_size != 32 (or None).""" return ( module_type in COMPRESSIBLE_MODULE_TYPES and scheme.weights is not None and scheme.weights.num_bits == 4 and scheme.weights.type == QuantizationType.FLOAT.value and scheme.weights.group_size == 16 ) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/nvfp4/helpers.py000066400000000000000000000066441521257237700326370ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ Helper functions for packing and unpacking FP4 (E2M1) quantized weights. FP4 E2M1 format uses 1 sign bit, 2 exponent bits, and 1 mantissa bit, supporting 8 positive and 8 negative values. This module provides efficient packing of two FP4 values into a single uint8 for storage. """ import torch __all__ = ["pack_fp4_to_uint8", "unpack_fp4_from_uint8"] FLOAT_TO_E2M1 = [ 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, ] kE2M1ToFloat = torch.tensor( [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32 ) def pack_fp4_to_uint8(x: torch.Tensor) -> torch.Tensor: """ Packs a tensor with values in the fp4 range into uint8. As there are 16 valid fp4 values, two fp4 values can be packed into one uint8. Each fp4 value is mapped to its particular index (e.g. 0.5 is mapped to index 1, 6.0 is mapped to index 7) which is then represented using 4 bits. Consecutive pairs of 4 bits are then packed into an uint8. :param x: tensor to pack :returns: a packed tensor in uint8 """ m, n = x.shape device = x.device if n % 2 != 0: raise ValueError( "tensor must have an even number of columns for nvfp4 compression" ) # Create lookup table for FP4 values to indices kE2M1 = torch.tensor(FLOAT_TO_E2M1, device=device, dtype=x.dtype) # Find closest valid FP4 value index for each element abs_x = torch.abs(x) abs_diff_x = torch.abs(abs_x.unsqueeze(-1) - kE2M1) # [m, n, 8] abs_indices = torch.argmin(abs_diff_x, dim=-1) # [m, n] # Apply sign bit (bit 3) to get final 4-bit representation indices = abs_indices + (torch.signbit(x).to(torch.long) << 3) # Reshape to prepare for packing pairs of values indices = indices.reshape(-1) # Reshape to pair consecutive elements indices = indices.reshape(-1, 2) # Pack pairs of 4-bit values into 8-bit values packed = (indices[:, 0] | (indices[:, 1] << 4)).to(torch.uint8) return packed.reshape(m, n // 2) # reference: https://github.com/vllm-project/vllm/pull/16362 def unpack_fp4_from_uint8( a: torch.Tensor, m: int, n: int, dtype: torch.dtype | None = torch.bfloat16 ) -> torch.Tensor: """ Unpacks uint8 values into fp4. Each uint8 consists of two fp4 values (i.e. first four bits correspond to one fp4 value, last four correspond to a consecutive fp4 value). The bits represent an index, which are mapped to an fp4 value. :param a: tensor to unpack :param m: original dim 0 size of the unpacked tensor :param n: original dim 1 size of the unpacked tensor :param dtype: dense dtype to cast the unpacked tensor to """ assert a.dtype == torch.uint8 # Vectorized nibble processing a_flat = a.flatten() high = (a_flat & 0xF0) >> 4 # Upper nibbles low = a_flat & 0x0F # Lower nibbles # Combine nibbles for batch processing combined = torch.stack((low, high), dim=1).flatten() # Vectorized sign and magnitude extraction signs = (combined & 0x08).to(torch.bool) # Sign bits abs_vals = (combined & 0x07).to(torch.long) # Magnitude indices # Device-aware lookup and sign application kE2M1 = kE2M1ToFloat.to(device=a.device) values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0) # Reshape to final form return values.reshape(m, n).to(dtype=dtype) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/pack_quantized/000077500000000000000000000000001521257237700325565ustar00rootroot00000000000000__init__.py000066400000000000000000000002471521257237700346130ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/pack_quantized# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .base import * from .helpers import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/pack_quantized/base.py000066400000000000000000000132561521257237700340510ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.compressors.base import ( COMPRESSIBLE_MODULE_TYPES, BaseCompressor, ) from compressed_tensors.compressors.pack_quantized.helpers import ( pack_to_int32, unpack_from_int32, ) from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import ( ActivationOrdering, QuantizationScheme, QuantizationStrategy, QuantizationType, ) from compressed_tensors.quantization.lifecycle.forward import dequantize, quantize from compressed_tensors.utils import TensorStateDict, getattr_chain __all__ = ["PackedQuantizationCompressor"] PACK_ZP_STRATS = [ QuantizationStrategy.GROUP.value, QuantizationStrategy.CHANNEL.value, ] @BaseCompressor.register(name=CompressionFormat.pack_quantized.value) class PackedQuantizationCompressor(BaseCompressor): """ Compresses a quantized weight by packing multiple sub-8-bit INT values into an int32. Supports num_bits in [1, 8]; each int32 holds ``pack_factor = 32 // num_bits`` values, with any unused high bits left as zero. """ @classmethod def compression_param_names(cls, scheme: QuantizationScheme) -> tuple[str]: param_names = ( "weight_packed", "weight_scale", "weight_shape", ) if not getattr_chain(scheme, "weights.symmetric", True): param_names += ("weight_zero_point",) if getattr_chain(scheme, "weights.actorder", None) == ActivationOrdering.GROUP: param_names += ("weight_g_idx",) if ( getattr_chain(scheme, "input_activations.strategy", None) == QuantizationStrategy.TENSOR_GROUP ): param_names += ("input_global_scale",) return param_names @classmethod def compress( cls, state_dict: TensorStateDict, scheme: QuantizationScheme ) -> TensorStateDict: """ Compress a per-module state dict. Quantizes the weight, packs it into int32 as ``weight_packed``, stores the original shape as ``weight_shape``, and removes ``weight``. If the quantization is asymmetric (GROUP or CHANNEL strategy) the zero-point is also packed; otherwise it is dropped. :param state_dict: local-name state dict (weight, weight_scale, …) :param quantization_args: quantization parameters for the weight :param device: device to move compressed tensors to :return: compressed state dict """ state_dict = state_dict.copy() weight = state_dict.pop("weight") scale = state_dict.get("weight_scale") zero_point = state_dict.get("weight_zero_point", None) g_idx = state_dict.get("weight_g_idx", None) weights = scheme.weights quantized_weight = quantize( x=weight, scale=scale, zero_point=zero_point, g_idx=g_idx, args=scheme.weights, dtype=torch.int8, ) state_dict["weight_packed"] = pack_to_int32(quantized_weight, weights.num_bits) state_dict["weight_shape"] = torch.tensor(weight.shape) if not weights.symmetric and weights.strategy in PACK_ZP_STRATS: assert zero_point is not None, "Asymmetric quant requires zero-point values" packed_zp = pack_to_int32(zero_point, weights.num_bits, packed_dim=0) state_dict["weight_zero_point"] = packed_zp.contiguous() state_dict = cls._remove_symmetric_zp(state_dict, scheme) return state_dict @classmethod def decompress( cls, state_dict: TensorStateDict, scheme: QuantizationScheme ) -> TensorStateDict: """ Decompress a per-module state dict. Unpacks ``weight_packed`` back to the original weight, removes ``weight_packed``, and unpacks the zero-point if present. :param state_dict: local-name state dict (weight_packed, weight_scale, …) :param quantization_args: quantization parameters for the weight :return: decompressed state dict with weight in float dtype """ state_dict = state_dict.copy() packed = state_dict.pop("weight_packed") scale = state_dict.get("weight_scale") zero_point = state_dict.get("weight_zero_point", None) g_idx = state_dict.get("weight_g_idx", None) original_shape = state_dict.get("weight_shape") weights = scheme.weights # Unpack zero_point before dequantization if needed if not weights.symmetric and weights.strategy in PACK_ZP_STRATS: assert zero_point is not None, "Asymmetric quant requires zero-point values" original_zp_shape = (*original_shape[:-1], scale.shape[-1]) zero_point = unpack_from_int32( zero_point, weights.num_bits, original_zp_shape, packed_dim=0 ) state_dict["weight_zero_point"] = zero_point unpacked = unpack_from_int32(packed, weights.num_bits, original_shape) state_dict["weight"] = dequantize( x_q=unpacked, scale=scale, zero_point=zero_point, g_idx=g_idx, ) return state_dict @classmethod def can_compress(cls, module_type: type, scheme: QuantizationScheme) -> bool: """Pack quantized matches weight-only INT quantization with 1..8 bits.""" return ( module_type in COMPRESSIBLE_MODULE_TYPES and scheme.weights is not None and scheme.input_activations is None and 1 <= scheme.weights.num_bits <= 8 and scheme.weights.type == QuantizationType.INT.value ) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/compressors/pack_quantized/helpers.py000066400000000000000000000113721521257237700345760ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ Helper functions for packing and unpacking quantized weights into int32 format. These functions enable efficient storage of sub-8-bit quantized weights by packing multiple values into 32-bit integers. """ import math from typing import Literal import torch __all__ = ["pack_to_int32", "unpack_from_int32"] def pack_to_int32( value: torch.Tensor, num_bits: int, packed_dim: Literal[0, 1] = 1, ) -> torch.Tensor: """ Packs a tensor of quantized weights stored in int8 into int32s with padding Pseudocode: 1. Shift wrt num_bits to convert to unsigned. num_bits=8 [1,2] -> [129, 130] 2. Pad to fill in 32 bits [129, 130] -> [129, 130, 0, 0] 3. convert to binary align in order [129, 130, 0, 0] -> 00000000 00000000 10000010 10000001 4. convert aligned binary to number 00000000000000001000001010000001 -> 33409 5. covert back to uint32 33409 -> 33409 :param value: tensor to pack :param num_bits: number of bits used to store underlying data, must be at least 1 :returns: packed int32 tensor """ if value.dtype is not torch.int8: raise ValueError("Tensor must be quantized to torch.int8 before packing") if not 1 <= num_bits <= 8: raise ValueError( f"Packing is only supported for num_bits in [1, 8], got {num_bits}" ) # Handle N-dimensional tensors (e.g. MoE 3D weights) by packing each 2D slice if value.ndim > 2: return torch.stack( [ pack_to_int32(value[i], num_bits, packed_dim) for i in range(value.shape[0]) ] ) # Convert to unsigned range for packing, matching quantization offset offset = 1 << (num_bits - 1) value = (value + offset).to(torch.uint8) device = value.device pack_factor = 32 // num_bits if packed_dim == 0: value = value.transpose(0, 1) rows, cols = value.shape padded_cols = math.ceil(cols / pack_factor) * pack_factor pad_len = padded_cols - cols if pad_len > 0: value = torch.nn.functional.pad(value, (0, pad_len)) num_groups = padded_cols // pack_factor # Use int32 here reshaped = value.view(rows, num_groups, pack_factor).to(torch.int32) bit_shifts = torch.arange(pack_factor, device=device, dtype=torch.int32) * num_bits packed = (reshaped << bit_shifts).sum(dim=2, dtype=torch.int32) if packed_dim == 0: packed = packed.transpose(0, 1) return packed def unpack_from_int32( value: torch.Tensor, num_bits: int, shape: torch.Size, packed_dim: Literal[0, 1] = 1, ) -> torch.Tensor: """ Unpacks a tensor of packed int32 weights into individual int8s, maintaining the original bit range. Return tensors in int8 :param value: tensor to upack :param num_bits: number of bits to unpack each data point into :param shape: shape to unpack into, used to remove padding :returns: unpacked int8 tensor """ if value.dtype is not torch.int32: raise ValueError( f"Expected {torch.int32} but got {value.dtype}, Aborting unpack." ) if not 1 <= num_bits <= 8: raise ValueError( f"Unpacking is only supported for num_bits in [1, 8], got {num_bits}" ) # Handle N-dimensional tensors (e.g. MoE 3D weights) by unpacking each 2D slice if value.ndim > 2: return torch.stack( [ unpack_from_int32(value[i], num_bits, shape[1:], packed_dim) for i in range(value.shape[0]) ] ) pack_factor = 32 // num_bits # unpack mask = (1 << num_bits) - 1 if packed_dim == 1: unpacked = torch.zeros( (value.shape[0], value.shape[1] * pack_factor), device=value.device, dtype=torch.int32, ) for i in range(pack_factor): unpacked[:, i::pack_factor] = (value >> (num_bits * i)) & mask # remove padding original_row_size = int(shape[1]) unpacked = unpacked[:, :original_row_size] else: unpacked = torch.zeros( (value.shape[0] * pack_factor, value.shape[1]), device=value.device, dtype=torch.int32, ) for i in range(pack_factor): unpacked[i::pack_factor, :] = (value >> (num_bits * i)) & mask # remove padding original_row_size = int(shape[0]) unpacked = unpacked[:original_row_size, :] # bits are packed in unsigned format, reformat to signed # update the value range from unsigned to signed offset = pow(2, num_bits) // 2 unpacked = (unpacked - offset).to(torch.int8) return unpacked vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/config/000077500000000000000000000000001521257237700264425ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/config/__init__.py000066400000000000000000000003431521257237700305530ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .base import * from .dense import * from .sparse_24_bitmask import * from .sparse_bitmask import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/config/base.py000066400000000000000000000062051521257237700277310ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum, unique from typing import List, Optional from compressed_tensors.registry import RegistryMixin from pydantic import BaseModel __all__ = ["SparsityCompressionConfig", "CompressionFormat", "SparsityStructure"] @unique class CompressionFormat(str, Enum): dense = "dense" sparse_bitmask = "sparse-bitmask" # legacy format sparse_24_bitmask = "sparse-24-bitmask" # legacy format int_quantized = "int-quantized" float_quantized = "float-quantized" naive_quantized = "naive-quantized" pack_quantized = "pack-quantized" marlin_24 = "marlin-24" # legacy format mixed_precision = "mixed-precision" nvfp4_pack_quantized = "nvfp4-pack-quantized" mxfp4_pack_quantized = "mxfp4-pack-quantized" mxfp8_quantized = "mxfp8-quantized" @unique class SparsityStructure(Enum): """ An enumeration to represent different sparsity structures. Attributes ---------- TWO_FOUR : str Represents a 2:4 sparsity structure. ZERO_ZERO : str Represents a 0:0 sparsity structure. UNSTRUCTURED : str Represents an unstructured sparsity structure. Examples -------- >>> SparsityStructure('2:4') >>> SparsityStructure('unstructured') >>> SparsityStructure('2:4') == SparsityStructure.TWO_FOUR True >>> SparsityStructure('UNSTRUCTURED') == SparsityStructure.UNSTRUCTURED True >>> SparsityStructure(None) == SparsityStructure.UNSTRUCTURED True >>> SparsityStructure('invalid') Traceback (most recent call last): ... ValueError: invalid is not a valid SparsityStructure """ TWO_FOUR = "2:4" UNSTRUCTURED = "unstructured" ZERO_ZERO = "0:0" def __new__(cls, value): obj = object.__new__(cls) obj._value_ = value.lower() if value is not None else value return obj @classmethod def _missing_(cls, value): # Handle None and case-insensitive values if value is None: return cls.UNSTRUCTURED for member in cls: if member.value == value.lower(): return member raise ValueError(f"{value} is not a valid {cls.__name__}") class SparsityCompressionConfig(RegistryMixin, BaseModel): """ Base data class for storing sparsity compression parameters :param format: name of compression format :param targets: List of layer names or layer types that aren't sparse and should be ignored during compression. By default, assume all layers are targeted :param ignore: List of layer names (unique) to ignore from targets. Defaults to None :param global_sparsity: average sparsity of the entire model :param sparsity_structure: structure of the sparsity, such as "unstructured", "2:4", "8:16" etc """ format: str targets: Optional[List[str]] = None ignore: Optional[List[str]] = None global_sparsity: Optional[float] = 0.0 sparsity_structure: Optional[str] = "unstructured" vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/config/dense.py000066400000000000000000000014471521257237700301200ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Optional from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig __all__ = ["DenseSparsityConfig"] @SparsityCompressionConfig.register(name=CompressionFormat.dense.value) class DenseSparsityConfig(SparsityCompressionConfig): """ Identity configuration for storing a sparse model in an uncompressed dense format :param global_sparsity: average sparsity of the entire model :param sparsity_structure: structure of the sparsity, such as "unstructured", "2:4", "8:16" etc """ format: str = CompressionFormat.dense.value global_sparsity: Optional[float] = 0.0 sparsity_structure: Optional[str] = "unstructured" vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/config/sparse_24_bitmask.py000066400000000000000000000015731521257237700323360ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Optional from compressed_tensors.config import ( CompressionFormat, SparsityCompressionConfig, SparsityStructure, ) __all__ = ["Sparse24BitMaskConfig"] @SparsityCompressionConfig.register(name=CompressionFormat.sparse_24_bitmask.value) class Sparse24BitMaskConfig(SparsityCompressionConfig): """ Configuration for storing a 24 sparse model using bytemask compression :param global_sparsity: average sparsity of the entire model :param sparsity_structure: structure of the sparsity, should always be "2:4" for this compression format """ format: str = CompressionFormat.sparse_24_bitmask.value global_sparsity: Optional[float] = 0.0 sparsity_structure: Optional[str] = SparsityStructure.TWO_FOUR.value vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/config/sparse_bitmask.py000066400000000000000000000014361521257237700320270ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Optional from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig __all__ = ["BitmaskConfig"] @SparsityCompressionConfig.register(name=CompressionFormat.sparse_bitmask.value) class BitmaskConfig(SparsityCompressionConfig): """ Configuration for storing a sparse model using bitmask compression :param global_sparsity: average sparsity of the entire model :param sparsity_structure: structure of the sparsity, such as "unstructured", "2:4", "8:16" etc """ format: str = CompressionFormat.sparse_bitmask.value global_sparsity: Optional[float] = 0.0 sparsity_structure: Optional[str] = "unstructured" vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/distributed/000077500000000000000000000000001521257237700275175ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/distributed/README.md000066400000000000000000000000551521257237700307760ustar00rootroot00000000000000Collection of distributed algorithm templatesvllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/distributed/__init__.py000066400000000000000000000011301521257237700316230ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from .assign import greedy_bin_packing from .module_parallel import replace_module_parallel from .utils import ( as_broadcastable, get_source_rank, init_dist, is_distributed, is_source_process, set_source_process, wait_for_comms, ) __all__ = [ "greedy_bin_packing", "replace_module_parallel", "as_broadcastable", "get_source_rank", "init_dist", "is_distributed", "is_source_process", "wait_for_comms", "set_source_process", ] vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/distributed/assign.py000066400000000000000000000032521521257237700313570ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Callable, Hashable, TypeVar __all__ = ["greedy_bin_packing"] T = TypeVar("T", bound=Hashable) def greedy_bin_packing( items: list[T], num_bins: int, item_weight_fn: Callable[[T], int | float] = lambda x: 1, ) -> tuple[list[T], list[list[T]], dict[T, int]]: """Distribute items across bins using a greedy bin-packing heuristic. Items are sorted by weight in descending order, then each item is assigned to the bin with the smallest current total weight. This approximates an even distribution of weight across bins. :param items: items to distribute. Sorted in-place by descending weight. :param num_bins: number of bins to distribute items across. :param item_weight_fn: callable that returns the weight of an item. Defaults to uniform weight of 1. :return: a 3-tuple of: - items: the input list, now sorted by descending weight. - bin_to_items: list of length ``num_bins`` where each element is the list of items assigned to that bin. - item_to_bin: mapping from each item to its assigned bin index. """ items.sort(key=item_weight_fn, reverse=True) bin_to_items: list[list[T]] = [[] for _ in range(num_bins)] item_to_bin: dict[T, int] = dict() bin_weights: list[float] = [0 for _ in range(num_bins)] for item in items: target_bin = bin_weights.index(min(bin_weights)) bin_to_items[target_bin].append(item) item_to_bin[item] = target_bin bin_weights[target_bin] += item_weight_fn(item) return items, bin_to_items, item_to_bin vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/distributed/module_parallel.py000066400000000000000000000066031521257237700332370ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Callable, Optional, TypeVar import torch import torch.distributed as dist from compressed_tensors.distributed.assign import greedy_bin_packing from compressed_tensors.distributed.utils import set_source_process from compressed_tensors.offload.utils import as_single_threaded, module_size from compressed_tensors.utils.module import ( get_direct_state_dict, replace_direct_state_dict, ) from tqdm import tqdm __all__ = ["replace_module_parallel"] T = TypeVar("T", bound=torch.nn.Module) def replace_module_parallel( modules: list[T], apply_fn: Callable[[T], None], weight_fn: Callable[[T], int | float] = module_size, desc: Optional[str] = None, ): """Apply a function to modules in parallel across distributed ranks. Distributes modules across ranks using greedy bin packing, then applies the function to each module on its assigned rank. Non-processing ranks temporarily move their modules to meta device to avoid increasing peak memory usage during compression. This implements the 4-step algorithm: 1. Decouple: Move non-processing rank modules to meta device 2. Compress On Meta: Apply function on meta device (prepare for step 4) 3. Compress On Device: Processing rank applies function without sync 4. Recouple: Broadcast offload pointer information across ranks :param modules: list of modules to process :param apply_fn: function to apply to each module :param weight_fn: function that returns the weight/size of a module for load balancing across ranks :param desc: optional description for the progress bar (shown on each rank) """ from compressed_tensors.offload import OffloadCache, disable_onloading, to_meta _, _, assigned_rank = greedy_bin_packing(modules, dist.get_world_size(), weight_fn) # Count modules assigned to this rank and create progress bar rank = dist.get_rank() num_assigned = sum(int(a == rank) for a in assigned_rank.values()) progress = tqdm( total=num_assigned, desc=desc, position=rank, disable=(desc is None) ) # Step 1 & 2: Decouple and compress on meta for non-processing ranks with disable_onloading(): for module in modules: if assigned_rank[module] != dist.get_rank(): to_meta(module) # 1. remove non-processing rank pointers apply_fn(module) # 2. compress on meta to match state dict for step 4 # Step 3: Apply on device for processing rank with as_single_threaded(): for module in modules: if assigned_rank[module] == dist.get_rank(): apply_fn(module) # 3. compress without triggering sync progress.update(1) # Step 4: Recouple - broadcast source offload across ranks for module in modules: with disable_onloading(): state_dict = get_direct_state_dict(module) # If module is not offloaded, manually broadcast tensors via object list if not isinstance(module._parameters, OffloadCache): broadcast_obj = [state_dict] dist.broadcast_object_list(broadcast_obj, src=assigned_rank[module]) state_dict = broadcast_obj[0] with set_source_process(assigned_rank[module]): replace_direct_state_dict(module, state_dict) # 4. broadcast vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/distributed/utils.py000066400000000000000000000063421521257237700312360ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib import os from typing import Iterable import torch import torch.distributed as dist __all__ = [ "is_source_process", "is_distributed", "init_dist", "as_broadcastable", "wait_for_comms", "get_source_rank", ] SRC_RANK = 0 def get_source_rank() -> int: """Get the current source rank for distributed operations.""" return SRC_RANK def is_source_process() -> bool: return not is_distributed() or dist.get_rank() == SRC_RANK @contextlib.contextmanager def set_source_process(src_rank: int): """ Context manager to temporarily designate a different rank as the source process. Similar to "is_rank0", this function chooses which rank serves as the source when broadcasting. This is useful when work is partioned across ranks, and then results need to be broadcasted across ranks from different sources. :param src_rank: the rank to designate as the main process within the context """ global SRC_RANK restore_rank, SRC_RANK = SRC_RANK, src_rank try: yield finally: SRC_RANK = restore_rank def is_distributed() -> bool: return dist.is_available() and dist.is_initialized() def init_dist(): if "TORCHELASTIC_RUN_ID" not in os.environ: raise ValueError( "Cannot find distributed environment. " "Please make sure you are using `torchrun --nproc-per-node ...`." ) rank = int(os.environ["RANK"]) local_rank = int(os.environ["LOCAL_RANK"]) world_size = int(os.environ["WORLD_SIZE"]) device = torch.device(f"cuda:{local_rank}") torch.cuda.set_device(device) dist.init_process_group( backend="nccl", init_method="env://", rank=rank, world_size=world_size, device_id=device, ) dist.barrier() _FP8_DTYPES = ( torch.float8_e4m3fn, torch.float8_e5m2, torch.float8_e4m3fnuz, torch.float8_e5m2fnuz, ) def as_broadcastable(tensor: torch.Tensor) -> torch.Tensor: """Return a view of the tensor that is compatible with ``dist.broadcast``. NCCL does not support broadcasting FP8 dtypes on hardware without sm_90 (Hopper or later). This function works around the limitation by viewing FP8 tensors as ``uint8``, which NCCL can broadcast on any hardware. Non-FP8 tensors are returned unchanged. :param tensor: the tensor to prepare for broadcasting :return: the original tensor, or a ``uint8`` view if the dtype is FP8 """ if tensor.dtype in _FP8_DTYPES: return tensor.data.view(torch.uint8) else: return tensor def wait_for_comms(pending_comms: Iterable[dist.Work]) -> None: """Block until all pending async distributed operations complete. Calls `wait()` on each work handle, then clears the list in-place so it can be reused for the next batch of operations. :param pending_comms: mutable list of async communication handles (returned by `dist.reduce`, `dist.broadcast`, etc. with `async_op=True`). The list is cleared after all operations have completed. """ for comm in pending_comms: comm.wait() pending_comms.clear() vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/entrypoints/000077500000000000000000000000001521257237700275735ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/entrypoints/__init__.py000066400000000000000000000001531521257237700317030ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/entrypoints/convert/000077500000000000000000000000001521257237700312535ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/entrypoints/convert/__init__.py000066400000000000000000000003051521257237700333620ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa # isort: off from .converters import * from .convert_checkpoint import * convert_checkpoint.py000066400000000000000000000114241521257237700354370ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/entrypoints/convert# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os import shutil from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path import tqdm from compressed_tensors.entrypoints.convert.convert_file import ( convert_file, validate_file, write_checkpoint_quantization_config, ) from compressed_tensors.entrypoints.convert.converters import ( Converter, build_inverse_weight_maps, ) from compressed_tensors.utils.safetensors_load import ( get_checkpoint_files, get_weight_map, is_weights_file, update_safetensors_index, ) from loguru import logger __all__ = ["convert_checkpoint", "exec_jobs"] def convert_checkpoint( model_stub: str | os.PathLike, save_directory: str | os.PathLike, converter: Converter, max_workers: int = 1, ): """ Convert a model checkpoint to either: - its equivalent quantized format in compressed-tensors - the unquantized format without loading it up in memory, instead operating directly on the model safetensors files. This entrypoint operates on a model stub or folder containing weights saved in safetensors files, and updates the corresponding quantization_config field in the config.json. All additional files will be copied to new checkpoint. :param model_stub: huggingface model hub or path to local weights files :param save_directory: new checkpoint will be saved in this directory. :param max_workers: number of worker threads to process files with :param device: gpu device to accelerate quantization with :param converters: converter we wish to apply to the checkpoint, e.g. conversion of some layers from some format to compressed-tensors """ # get all model_files for checkpoint model_files = get_checkpoint_files(model_stub) weight_map = get_weight_map(model_files) # Build inverse_weight_maps, so that each job knows how to load up every necessary # weight and its dependencies inverse_weight_maps = build_inverse_weight_maps( weight_map=weight_map, model_files=model_files, converters=[converter], ) # Build validation/conversion jobs, copy over any other file validate_jobs = [] convert_jobs = [] for shard_name, resolved_path in model_files.items(): save_path = Path(save_directory) / shard_name if shard_name.endswith("safetensors"): if shard_name not in inverse_weight_maps: raise ValueError( f"Could not find inverse_weight_map for shard {shard_name}" ) validate_jobs.append( (validate_file, inverse_weight_maps[shard_name], converter) ) convert_jobs.append( (convert_file, inverse_weight_maps[shard_name], save_path, converter) ) else: if is_weights_file(shard_name): logger.warning(f"Skip processing for weights file {shard_name}") if str(resolved_path) != str(save_path): save_path.parent.mkdir(parents=True, exist_ok=True) logger.info(f"Copying {shard_name} {save_path}") shutil.copyfile(resolved_path, save_path) # Validate before long-running procssing job exec_jobs(validate_jobs, max_workers, desc="Validating") # Process weights, accumulating total bytes used and the new weight_map total_size = 0 weight_map = dict() convert_results = exec_jobs(convert_jobs, max_workers, desc="Converting") for _total_size, _weight_map in convert_results: total_size += _total_size weight_map.update(_weight_map) # Update config and safetensors index write_checkpoint_quantization_config(save_directory, converter) update_safetensors_index(save_directory, total_size, weight_map) def exec_jobs( jobs: list[tuple[Callable, ...]], max_workers: int = 1, desc: str = "Executing Jobs" ) -> list: """ Execute jobs in parallel, using ThreadPoolExecutor :param jobs: list of tuples, the first entry of which is the callable, and the remaining elements are the inputs args to the callable :param max_workers: number of workers to use :param desc: tqdm description """ results = [] # For easier debugging, don't run single-threaded jobs via ThreadPoolExecutor if max_workers == 1: for job in tqdm.tqdm(jobs, desc=desc): results.append(job[0](*job[1:])) return results with ThreadPoolExecutor(max_workers) as executor: futures = [executor.submit(*job) for job in jobs] for future in tqdm.tqdm(as_completed(futures), total=len(futures), desc=desc): results.append(future.result()) return results vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/entrypoints/convert/convert_file.py000066400000000000000000000112051521257237700343030ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import os from compressed_tensors import __version__ as ct_version from compressed_tensors.base import COMPRESSION_VERSION_NAME, QUANTIZATION_CONFIG_NAME from compressed_tensors.entrypoints.convert import Converter from compressed_tensors.utils.safetensors_load import ( InverseWeightMap, find_config_path, load_tensors_from_inverse_weight_map, ) from loguru import logger from safetensors.torch import save_file __all__ = [ "validate_file", "convert_file", "write_checkpoint_quantization_config", ] def write_checkpoint_quantization_config( save_directory: str | os.PathLike, converter: Converter, ): """ Write the quantization config produced by `converter` into the model config file (config.json or params.json) in save_directory. This is called after the convert checkpoint pathway completes to record which quantization was applied. The quantization_config section is replaced entirely with the new config. :param save_directory: directory containing the model config file :param converter: Converter instance whose create_config() produces the updated quantization config """ quant_config_data = None if (quant_config := converter.create_config()) is not None: quant_config_data = quant_config.model_dump() quant_config_data[COMPRESSION_VERSION_NAME] = ct_version config_file_path = find_config_path(save_directory) if config_file_path is not None: with open(config_file_path, "r") as file: config_data = json.load(file) if quant_config_data is None: # if no new quant config, make sure checkpoint quant config is empty if QUANTIZATION_CONFIG_NAME in config_data: del config_data[QUANTIZATION_CONFIG_NAME] # some checkpoints have quant config nested inside text_config elif ( "text_config" in config_data and QUANTIZATION_CONFIG_NAME in config_data["text_config"] ): del config_data["text_config"][QUANTIZATION_CONFIG_NAME] else: # if new quant config, overwrite checkpoint quant config config_data[QUANTIZATION_CONFIG_NAME] = quant_config_data with open(config_file_path, "w") as file: json.dump(config_data, file, indent=2, sort_keys=True) else: logger.warning( f"Could not find config file in {save_directory}. Please add to config " f"{json.dumps(quant_config_data, indent=2, sort_keys=True)}" ) def validate_file( inverse_weight_map: InverseWeightMap, converter: Converter, ): """ Validate that each quantizable tensor in a safetensors file can be quantized. :param inverse_weight_map: mapping of resolved source file path -> list of tensor names to load from that file. Precomputed by build_inverse_weight_map() in the job-building phase. Example: {"/path/shard0.safetensors": ["q_proj.weight"], "/path/shard1.safetensors": ["k_proj.weight", "v_proj.weight"]} :param converter: converter we wish to apply to the checkpoint, e.g. conversion of some layers from some format to compressed-tensors """ tensors = load_tensors_from_inverse_weight_map(inverse_weight_map) converter.validate(tensors) def convert_file( inverse_weight_map: InverseWeightMap, save_path: str | os.PathLike, converter: Converter, ) -> tuple[int, dict[str, str]]: """ Convert tensors in a given safetensors file :param inverse_weight_map: mapping of resolved source file path -> list of tensor names to load from that file. Precomputed by build_inverse_weight_map() in the job-building phase. Example: {"/path/shard0.safetensors": ["q_proj.weight"], "/path/shard1.safetensors": ["k_proj.weight", "v_proj.weight"]} :param save_path: save path of file with quantized weights :param converter: converter we wish to apply to the checkpoint, e.g. conversion of some layers from some format to compressed-tensors :returns: tuple of (total_size, weight_map), respectively the total size in bytes of the saved file and dictionary of weight name -> save path """ tensors = load_tensors_from_inverse_weight_map(inverse_weight_map) tensors = converter.process(tensors) save_file(tensors, save_path) total_size = sum(tensor.nbytes for tensor in tensors.values()) weight_map = {key: os.path.basename(save_path) for key in tensors.keys()} return total_size, weight_map vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/entrypoints/convert/converters/000077500000000000000000000000001521257237700334455ustar00rootroot00000000000000__init__.py000066400000000000000000000004321521257237700354760ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/entrypoints/convert/converters# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa # isort: skip_file from .base import * from .ct_dequantizer import * from .autoawq import * from .modelopt_nvfp4 import * from .fp8block_dequantizer import * autoawq.py000066400000000000000000000221331521257237700354220ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/entrypoints/convert/converters# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import re from collections.abc import Iterable from typing import Any, cast import torch from compressed_tensors.compressors.pack_quantized.helpers import pack_to_int32 from compressed_tensors.config import CompressionFormat from compressed_tensors.entrypoints.convert.converters import Converter from compressed_tensors.quantization import ( QuantizationArgs, QuantizationConfig, QuantizationScheme, QuantizationStatus, QuantizationStrategy, QuantizationType, ) from compressed_tensors.utils.match import match_name from transformers import AutoConfig __all__ = ["AutoAWQConverter"] class AutoAWQConverter(Converter): """ Convert AutoAWQ GEMM checkpoint tensors to compressed-tensors WNA16 tensors. AutoAWQ GEMM stores quantized linear layers as qweight/qzeros/scales. This converter unpacks qweight into compressed-tensors' signed integer convention and preserves the per-group quantization parameters as weight_scale and weight_zero_point. """ AWQ_REVERSE_ORDER = [0, 4, 1, 5, 2, 6, 3, 7] def __init__( self, bits: int = 4, group_size: int = 128, zero_point: bool = True, version: str = "gemm", ignore: Iterable[str] = ("lm_head",), targets: Iterable[str] = ("Linear",), ): if bits != 4: raise ValueError("AutoAWQConverter currently supports only 4-bit weights") if version != "gemm": raise ValueError(f"Unsupported AutoAWQ version: {version}") self.bits = bits self.group_size = group_size self.zero_point = zero_point self.version = version self.ignore = list(ignore) self.targets = list(targets) @classmethod def from_pretrained( cls, model_name_or_path: str, targets: Iterable[str] = ("Linear",), trust_remote_code: bool = False, ) -> "AutoAWQConverter": config = AutoConfig.from_pretrained( model_name_or_path, trust_remote_code=trust_remote_code ) autoawq_config = getattr(config, "quantization_config", None) if autoawq_config is None: raise ValueError("Model config does not contain quantization_config") autoawq_config = cast(dict[str, Any], autoawq_config) if autoawq_config.get("quant_method") != "awq": raise ValueError("Model config is not an AutoAWQ config") return cls.from_autoawq_config( autoawq_config, targets=targets, ) @classmethod def from_autoawq_config( cls, autoawq_config: dict[str, Any], targets: Iterable[str] = ("Linear",), ) -> "AutoAWQConverter": ignore = ["lm_head"] for module in autoawq_config.get("modules_to_not_convert") or []: ignore.append(f"re:.*{re.escape(module)}.*") return cls( bits=autoawq_config.get("bits", 4), group_size=autoawq_config.get("group_size", 128), zero_point=autoawq_config.get("zero_point", True), version=autoawq_config.get("version", "gemm"), ignore=ignore, targets=targets, ) def process(self, tensors: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: for name in list(tensors): if not name.endswith(".qweight"): continue module_name = name.removesuffix(".qweight") if not self._is_targeted(module_name): continue qweight = tensors.pop(f"{module_name}.qweight") qzeros = tensors.pop(f"{module_name}.qzeros", None) scales = tensors.pop(f"{module_name}.scales") weight, weight_scale, weight_zero_point = self._convert_gemm_module( qweight, scales, qzeros ) tensors[f"{module_name}.weight_scale"] = weight_scale tensors[f"{module_name}.weight_packed"] = pack_to_int32(weight, self.bits) tensors[f"{module_name}.weight_shape"] = torch.tensor(weight.shape) if weight_zero_point is not None: weight_zero_point = pack_to_int32( weight_zero_point, self.bits, packed_dim=0 ).contiguous() tensors[f"{module_name}.weight_zero_point"] = weight_zero_point return tensors def validate(self, tensors: dict[str, torch.Tensor]): for name in tensors: module_name, _, param_name = name.rpartition(".") if param_name in {"qweight", "qzeros", "scales"}: if not self._is_targeted(module_name): raise ValueError(f"Found unexpected non-targeted tensor {name}") if param_name != "qweight" or not self._is_targeted(module_name): continue for dependency in self.get_dependencies(name): if dependency not in tensors: raise ValueError( f"Found qweight without corresponding {dependency}" ) def create_config(self) -> QuantizationConfig: weights = QuantizationArgs( num_bits=self.bits, type=QuantizationType.INT, symmetric=not self.zero_point, group_size=self.group_size, strategy=QuantizationStrategy.GROUP, ) return QuantizationConfig( config_groups={ "config_group_0": QuantizationScheme( targets=self.targets, weights=weights, format=CompressionFormat.pack_quantized.value, ) }, ignore=self.ignore, format=CompressionFormat.pack_quantized.value, quantization_status=QuantizationStatus.COMPRESSED.value, ) def get_dependencies(self, weight_name: str) -> set[str]: module_name, _, suffix = weight_name.rpartition(".") if suffix == "qweight" and self._is_targeted(module_name): dependencies = {f"{module_name}.scales"} if self.zero_point: dependencies.add(f"{module_name}.qzeros") return dependencies return set() def _convert_gemm_module( self, qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: if self.zero_point and qzeros is None: raise ValueError("Found qweight without corresponding qzeros") iweight, izeros = self.unpack_awq(qweight, qzeros, self.bits) iweight, izeros = self.reverse_awq_order(iweight, izeros, self.bits) iweight = torch.bitwise_and(iweight, (2**self.bits) - 1) quantized_weight = iweight - 2 ** (self.bits - 1) weight_zero_point = None if self.zero_point: assert izeros is not None weight_zero_point = torch.bitwise_and(izeros, (2**self.bits) - 1) weight_zero_point = weight_zero_point - 2 ** (self.bits - 1) weight_zero_point = weight_zero_point.T.contiguous() return ( quantized_weight.T.contiguous(), scales.T.contiguous(), weight_zero_point, ) def _is_targeted(self, module_name: str) -> bool: if any(match_name(module_name, ignore) for ignore in self.ignore): return False if len(self.targets) == 0 or "Linear" in self.targets: return True return any(match_name(module_name, target) for target in self.targets) @staticmethod def unpack_awq( qweight: torch.Tensor, qzeros: torch.Tensor | None, bits: int ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Unpack AutoAWQ GEMM int32-packed weights and zero-points into int8 values. """ shifts = torch.arange(0, 32, bits, device=qweight.device) iweights = torch.bitwise_right_shift( qweight[:, :, None], shifts[None, None, :] ).to(torch.int8) iweights = iweights.view(iweights.shape[0], -1) if qzeros is None: return iweights, None izeros = torch.bitwise_right_shift( qzeros[:, :, None], shifts[None, None, :] ).to(torch.int8) izeros = izeros.view(izeros.shape[0], -1) return iweights, izeros @staticmethod def reverse_awq_order( iweights: torch.Tensor, izeros: torch.Tensor | None, bits: int ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Undo AutoAWQ's special intra-int32 packing order. """ reverse_order_tensor = torch.arange( iweights.shape[-1], dtype=torch.int32, device=iweights.device, ) reverse_order_tensor = reverse_order_tensor.view(-1, 32 // bits) reverse_order_tensor = reverse_order_tensor[ :, AutoAWQConverter.AWQ_REVERSE_ORDER ] reverse_order_tensor = reverse_order_tensor.view(-1) iweights = iweights[:, reverse_order_tensor] if izeros is not None: izeros = izeros[:, reverse_order_tensor] return iweights, izeros base.py000066400000000000000000000131021521257237700346470ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/entrypoints/convert/converters# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from __future__ import annotations from collections import defaultdict from typing import TYPE_CHECKING, Protocol import torch from compressed_tensors.utils.safetensors_load import InverseWeightMap __all__ = ["Converter", "build_inverse_weight_maps"] if TYPE_CHECKING: from compressed_tensors.quantization import QuantizationConfig class Converter(Protocol): """ Converter interface, to modify safetensors files based on tensor name and pointer to torch.Tensor, and create the QuantizationConfig """ def process(self, tensors: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """ Process tensors, returning converted set to be saved in a safetensors file. Converted tensors are dequantized (i.e. no quantization config) or in a compressed-tensors compatible format. Examples: - rename tensor or invert weights to match compressed-tensors convention. - dequantize to full-precision :param tensors: dictionary of tensor name to tensor, as loaded from safetensors file. Tensor name is a concatenation of module name and parameter name, e.g. - `model.layers.0.self_attn.q_proj.weight` - `model.layers.0.mlp.up_proj.weight_packed` :returns: dictionary of converted tensor name to tensor, to be saved in a safetensors file. Same format as input param tensors. """ raise NotImplementedError() def validate(self, tensors: dict[str, torch.Tensor]): """ Validation layer to quickly log warnings or raise an error if the safetensors file is not compatible with Converter. :param tensors: dictionary of tensor name to tensor, as loaded from safetensors file. """ raise NotImplementedError() def create_config(self) -> QuantizationConfig | None: """ Create compressed-tensors QuantizationConfig so that it can be set in the new model checkpoint's config.json. If the converter is moving checkpoint to full-precision, have this function return None, and quantization_config will be removed from config.json """ raise NotImplementedError() def get_dependencies(self, weight_name: str) -> set[str]: """ Given a weight name, return a set of all dependency weight names, so that weights can be processed correctly and in a parallelized fashion. If there are no dependencies, an empty dict should be returned. :returns: set[str] of dependency weight names """ raise NotImplementedError() def build_inverse_weight_maps( weight_map: dict[str, str], model_files: dict[str, str], converters: list[Converter], ) -> dict[str, InverseWeightMap]: """ For a given output shard, precompute exactly which tensors to load from which source files — including required partner tensors from other shards. This is necessary because some converters require that a set of tensors are accessible in order for them to be processed correctly. :param shard_name: the shard filename this job will process and save :param weight_map: tensor name -> shard filename (from safetensors.index.json) :param model_files: shard filename -> resolved absolute path :return: {resolved_file_path: [tensor_names_to_load]} """ def get_dependencies_recursive( weight_name: str, converters: list[Converter], current_deps: set[str] ) -> set[str]: for converter in converters: deps = converter.get_dependencies(weight_name) for dep in deps: if dep not in current_deps: current_deps.add(dep) get_dependencies_recursive(dep, converters, current_deps) return current_deps # map of weight name -> set of dependency names weight_deps_dict: dict[str, set[str]] = dict() for weight_name in weight_map: weight_deps_dict[weight_name] = get_dependencies_recursive( weight_name, converters, set() ) assert ( weight_name not in weight_deps_dict[weight_name] ), f"{weight_name} found in dependencies {weight_deps_dict[weight_name]}" # set of all dependencies (i.e. all weight names required by another) all_dependencies: set[str] = set().union(*weight_deps_dict.values()) inverse_weight_maps: dict[str, InverseWeightMap] = defaultdict( lambda: defaultdict(list) ) for weight_name, weight_shard_name in weight_map.items(): if weight_name in all_dependencies: # weight is a partner to some other primary tensor, skip it continue # weight is purely a primary weight, is not a dependency of anything # add it and all its dependencies current_iwm: InverseWeightMap = inverse_weight_maps[weight_shard_name] dependency_weights = weight_deps_dict[weight_name] for weight_to_add_name in [ weight_name, *dependency_weights, ]: if weight_to_add_name not in weight_map: raise ValueError( f"Dependency weight {weight_to_add_name} not found in weight map" ) weight_to_add_shard_name = weight_map[weight_to_add_name] resolved_path = model_files[weight_to_add_shard_name] current_iwm[resolved_path].append(weight_to_add_name) # return dicts, not defaultdicts, to avoid silent errors return {k: dict(v) for k, v in inverse_weight_maps.items()} ct_dequantizer.py000066400000000000000000000152551521257237700367710ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/entrypoints/convert/converters# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os from typing import Iterable import pydantic import torch from compressed_tensors.compressors import BaseCompressor from compressed_tensors.compressors.format import infer_module_format from compressed_tensors.config import CompressionFormat from compressed_tensors.entrypoints.convert.converters import Converter from compressed_tensors.quantization import KVCacheScaleType, QuantizationConfig from compressed_tensors.utils.match import match_name, match_quantizable_tensors from compressed_tensors.utils.safetensors_load import ( get_checkpoint_files, get_quantization_config, ) from transformers.file_utils import CONFIG_NAME class CompressedTensorsDequantizer(Converter): """ Dequantize a checkpoint in the compressed-tensors quant format The resultant weights will be stored in user-provided dtype """ def __init__( self, model_stub: str | os.PathLike, ignore: Iterable[str] = tuple(), dtype=torch.bfloat16, ): self.dtype = dtype # load quantization config from model_stub model_files = get_checkpoint_files(model_stub) if CONFIG_NAME in model_files: config_resolved_path = model_files[CONFIG_NAME] elif "params.json" in model_files: config_resolved_path = model_files["params.json"] else: raise ValueError("Could not find config.json file") quant_config_data = get_quantization_config(config_resolved_path) if quant_config_data is None: raise ValueError("Could not find quantization_config in config.json") try: self.quant_config = QuantizationConfig.model_validate(quant_config_data) except pydantic.ValidationError as e: raise ValueError( "Model quantization config was found, but it does not match expected " "compressed-tensors quantization format" ) from e # hydrate with additional ignore and inferred scheme formats self.quant_config.ignore += list(ignore) for scheme in self.quant_config.config_groups.values(): scheme.format = CompressionFormat( infer_module_format(torch.nn.Linear, scheme) ) def process(self, tensors: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """ Dequantize compressed tensors to full-precision weight tensors in dtype provided to constructor """ dequantized_tensors = {} for scheme in self.quant_config.config_groups.values(): compressor = BaseCompressor.get_value_from_registry(scheme.format) param_names = compressor.compression_param_names(scheme) for module_name, tensor_name in match_quantizable_tensors( tensors, ignore=self.quant_config.ignore, targets=scheme.targets, param_targets=[param_names[0]], ): # Create state dict of param_name -> torch.Tensor state_dict = { f"{param_name}": tensors.pop(f"{module_name}.{param_name}") for param_name in param_names } dequantized_state_dict = compressor.decompress(state_dict, scheme) # Add only weight param to dequantized tensors dequantized_tensors[f"{module_name}.weight"] = dequantized_state_dict[ "weight" ].to(self.dtype) # Copy over any remaining ignored/untargeted tensors, skipping kv cache qparams kv_cache_param_names = [v.value for v in KVCacheScaleType] for name, tensor in tensors.items(): if any([name.endswith(param_name) for param_name in kv_cache_param_names]): continue dequantized_tensors[name] = tensor return dequantized_tensors def validate(self, tensors: dict[str, torch.Tensor]): """ Ensure all tensor names of targeted layers are expected and no untargeted layers have unexpected tensor names """ consumed_keys = set() matched_modules = set() for scheme in self.quant_config.config_groups.values(): compressor = BaseCompressor.get_value_from_registry(scheme.format) param_names = compressor.compression_param_names(scheme) for module_name, _ in match_quantizable_tensors( tensors, self.quant_config.ignore, scheme.targets, param_targets=[param_names[0]], ): matched_modules.add(module_name) for param_name in param_names: expected_key = f"{module_name}.{param_name}" if expected_key not in tensors: raise ValueError(f"Expected key {expected_key} not found") consumed_keys.add(expected_key) unconsumed_tensor_names = [ name for name in tensors if name not in consumed_keys and name.rpartition(".")[0] in matched_modules ] if len(unconsumed_tensor_names) != 0: raise ValueError( f"Found {len(unconsumed_tensor_names)} unconsumed keys -- " f"{unconsumed_tensor_names}" ) return def create_config(self) -> QuantizationConfig | None: return None def get_dependencies(self, weight_name: str) -> set[str]: """ Dependencies are determined by the associated compressor's compression_param_names. The first param name in the returned list is treated as the root param, and is usually "weight" or "weight_packed" If weight_name is untargeted or ignored, an empty set is returned """ module_name, _, param_name = weight_name.rpartition(".") if any( [match_name(module_name, ignore) for ignore in self.quant_config.ignore] ): return set() for scheme in self.quant_config.config_groups.values(): compressor = BaseCompressor.get_value_from_registry(scheme.format) compression_param_names = compressor.compression_param_names(scheme) if "Linear" in scheme.targets or any( [match_name(module_name, target) for target in scheme.targets] ): if param_name == compression_param_names[0]: return set( f"{module_name}.{param_name}" for param_name in compression_param_names[1:] ) else: return set() return set() fp8block_dequantizer.py000066400000000000000000000133721521257237700400710ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/entrypoints/convert/converters# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Iterable import torch from compressed_tensors.entrypoints.convert.converters import Converter from compressed_tensors.quantization import QuantizationConfig from compressed_tensors.quantization.utils.helpers import ( maybe_pad_tensor_for_block_quant, ) from compressed_tensors.utils.match import match_name, match_quantizable_tensors class FP8BlockDequantizer(Converter): """ Dequantize a checkpoint that has been block-quantized with FP8 quant_method The resultant weights will be stored in user-provided dtype """ def __init__( self, ignore: Iterable[str] = tuple(), targets: Iterable[str] = tuple(), weight_block_size: tuple[int] = (128, 128), dtype=torch.bfloat16, ): self.ignore = ignore self.targets = targets self.weight_block_size = weight_block_size self.dtype = dtype self.param_names = ["weight", "weight_scale_inv"] def process(self, tensors: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """ Dequantize the fp8 block tensors (weight, weight_scale_inv) to full-precision weight tensors in dtype provided to constructor """ for module_name, name in match_quantizable_tensors( tensors, self.ignore, self.targets, param_targets=self.param_names ): param_name = name.rpartition(".")[-1] if param_name == "weight": # weight * weight_scale_inv -> dequantized weight tensors[f"{module_name}.weight"] = self._create_dequantized_weight( tensors[f"{module_name}.weight"], tensors[f"{module_name}.weight_scale_inv"], ) del tensors[f"{module_name}.weight_scale_inv"] return tensors def validate(self, tensors: dict[str, torch.Tensor]): """ Ensure all tensor names of targeted layers are expected and no untargeted layers have unexpected tensor names """ targeted_names = [ name for _, name in match_quantizable_tensors( tensors, self.ignore, self.targets, param_targets=self.param_names ) ] for name in targeted_names: module_name, _, param_name = name.rpartition(".") if ( param_name == "weight" and f"{module_name}.weight_scale_inv" not in tensors ): raise ValueError( f"Found weight without corresponding weight_scale_inv {name}" ) if ( param_name == "weight_scale_inv" and f"{module_name}.weight" not in tensors ): raise ValueError( f"Found weight_scale_inv without corresponding weight {name}" ) disallowed_names = ["weight_scale_inv"] untargeted_names = [ name for name in tensors.keys() if name not in targeted_names ] for name in untargeted_names: param_name = name.rsplit(".", 1)[-1] if param_name in disallowed_names: raise ValueError(f"Found unexpected non-targeted tensor {name}") def create_config(self) -> QuantizationConfig | None: return None def get_dependencies(self, weight_name: str) -> set[str]: module_name, _, param_name = weight_name.rpartition(".") if ( any([match_name(module_name, target) for target in self.targets]) and not any([match_name(module_name, ignore) for ignore in self.ignore]) and param_name == "weight" ): return {f"{module_name}.weight_scale_inv"} return set() def _create_dequantized_weight( self, weight: torch.Tensor, weight_scale_inv: torch.Tensor ) -> torch.Tensor: """ Convert fp8 weight and fp32 weight_scale_inv tensors into corresponding dequantized weight tensor. Tensors are upscaled to fp32 before scaling :return: dequantized tensor in self.dtype and same shape as input weight tensor """ original_shape = weight.shape block_height, block_width = self.weight_block_size # Pad tensor if dimensions are not evenly divisible by block size weight = maybe_pad_tensor_for_block_quant(weight, tuple(self.weight_block_size)) padded_shape = weight.shape # Reshape into blocks of shape: # (num_rows_blocks, block_height, num_cols_blocks, block_width) num_rows_blocks = padded_shape[0] // block_height num_cols_blocks = padded_shape[1] // block_width weight_blocks = weight.reshape( num_rows_blocks, block_height, num_cols_blocks, block_width, ).transpose( 1, 2 ) # (num_rows_blocks, num_cols_blocks, block_height, block_width) # Expand scale_inv for broadcasting over block dimensions # weight_scale_inv shape: (num_rows_blocks, num_cols_blocks) # Expand to: (num_rows_blocks, num_cols_blocks, 1, 1) scale_inv_expanded = weight_scale_inv.unsqueeze(-1).unsqueeze(-1) # Dequantize: weight_bf16 = weight_fp8 * weight_scale_inv dequantized_blocks = ( weight_blocks.to(torch.float32) * scale_inv_expanded.to(torch.float32) ).to(self.dtype) # Restore padded shape dequantized = dequantized_blocks.transpose(1, 2).reshape(padded_shape) # Truncate to original dimensions if padding was applied if original_shape != padded_shape: dequantized = dequantized[tuple([slice(v) for v in original_shape])] return dequantized modelopt_nvfp4.py000066400000000000000000000141251521257237700367030ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/entrypoints/convert/converters# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Iterable import torch from compressed_tensors.config import CompressionFormat from compressed_tensors.entrypoints.convert.converters import Converter from compressed_tensors.quantization import ( QuantizationArgs, QuantizationConfig, QuantizationScheme, QuantizationStatus, ) from compressed_tensors.quantization.quant_scheme import NVFP4 from compressed_tensors.utils.match import match_name, match_quantizable_tensors class ModelOptNvfp4Converter(Converter): """ Convert params from modelopt NVFP4 to CT NVFP4 convention, and optionally the kv_cache_scheme """ def __init__( self, ignore: Iterable[str] = tuple(), targets: Iterable[str] = tuple(), kv_cache_scheme: QuantizationArgs | None = None, ): self.ignore = ignore self.targets = targets self.kv_cache_scheme = kv_cache_scheme self.param_names = ["input_scale", "weight", "weight_scale", "weight_scale_2"] if self.kv_cache_scheme is not None: self.param_names += ["k_scale", "v_scale"] def process(self, tensors: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """ Map the modelopt NVFP4 tensors to the appropriate compressed-tensors NVFP4 format. Some tensors require rename, some require inversion - 1 / input_scale -> input_global_scale - weight -> weight_packed - 1 / weight_scale_2 -> weight_global_scale """ for module_name, name in match_quantizable_tensors( tensors, self.ignore, self.targets, param_targets=self.param_names ): param_name = name.rpartition(".")[-1] match param_name: # input_scale -> input_global_scale F32 case "input_scale": # convert modelopt input_scale x -> 1/x # https://github.com/vllm-project/vllm/blob/v0.13.0/vllm/model_executor/layers/quantization/modelopt.py#L1070-L1073 # https://github.com/vllm-project/vllm/blob/v0.13.0/vllm/model_executor/layers/quantization/modelopt.py#L1134 # https://github.com/vllm-project/vllm/blob/v0.13.0/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py#L190 tensors[f"{module_name}.input_global_scale"] = 1 / tensors[name] del tensors[name] # weight -> weight_packed U8 case "weight": tensors[f"{module_name}.weight_packed"] = tensors[name] del tensors[name] # weight_scale -> weight_scale F8_E4M3 case "weight_scale": pass # weight_scale_2 -> weight_global_scale F32 case "weight_scale_2": # convert modelopt weight_scale_2 x -> 1/x # https://github.com/vllm-project/vllm/blob/v0.13.0/vllm/model_executor/layers/quantization/modelopt.py#L1066-L1068 # https://github.com/vllm-project/vllm/blob/v0.13.0/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py#L163-L166 tensors[f"{module_name}.weight_global_scale"] = 1 / tensors[name] del tensors[name] case "k_scale" | "v_scale": # convert kv cache scales to appropriate dtype # often F32 in modelopt, defaults BF16 in compressed-tensors tensors[name] = tensors[name].to( self.kv_cache_scheme.scale_dtype or torch.bfloat16 ) return tensors def validate(self, tensors: dict[str, torch.Tensor]): """ Ensure all tensor names of targeted layers are expected and no untargeted layers have unexpected tensor names """ targeted_names = [ name for _, name in match_quantizable_tensors( tensors, self.ignore, self.targets, param_targets=self.param_names ) ] for name in targeted_names: param_name = name.rpartition(".")[-1] disallowed_names = [ "input_scale", "weight_scale", "weight_scale_2", "k_scale", "v_scale", ] untargeted_names = [ name for name in tensors.keys() if name not in targeted_names ] for name in untargeted_names: param_name = name.rpartition(".")[-1] if param_name in disallowed_names: raise ValueError(f"Hit unexpected non-targeted tensor {name}") def get_dependencies(self, weight_name: str) -> set[str]: module_name, _, param_name = weight_name.rpartition(".") if ( any([match_name(module_name, target) for target in self.targets]) and not any([match_name(module_name, ignore) for ignore in self.ignore]) and param_name == "weight" ): deps = { f"{module_name}.input_scale", f"{module_name}.weight_scale", f"{module_name}.weight_scale_2", } if self.kv_cache_scheme: if module_name.endswith("k_proj"): deps.add(f"{module_name}.k_scale") if module_name.endswith("v_proj"): deps.add(f"{module_name}.v_scale") return deps return set() def create_config(self) -> QuantizationConfig: return QuantizationConfig( config_groups={ "config_group_0": QuantizationScheme( **NVFP4, targets=self.targets, format=CompressionFormat.nvfp4_pack_quantized.value, ) }, ignore=self.ignore, kv_cache_scheme=self.kv_cache_scheme, format=CompressionFormat.nvfp4_pack_quantized.value, quantization_status=QuantizationStatus.COMPRESSED.value, ) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/linear/000077500000000000000000000000001521257237700264475ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/linear/__init__.py000066400000000000000000000001531521257237700305570ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/linear/compressed_linear.py000066400000000000000000000011751521257237700325230ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.quantization import QuantizationScheme class CompressedLinear(torch.nn.Linear): """ Wrapper module for running a compressed forward pass of a quantized Linear module. The wrapped layer will decompressed on each forward call. """ @classmethod def from_linear( cls, module: torch.nn.Linear, quantization_scheme: QuantizationScheme, quantization_format: str, ): raise ValueError("`CompressedLinear` is no longer supported") vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/logger.py000066400000000000000000000105611521257237700270310ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ Logger configuration for Compressed Tensors. """ import os import sys from dataclasses import dataclass from typing import Any, Dict, Optional from loguru import logger __all__ = ["LoggerConfig", "configure_logger", "logger"] # used by `support_log_once`` _logged_once = set() @dataclass class LoggerConfig: disabled: bool = False clear_loggers: bool = False console_log_level: Optional[str] = None log_file: Optional[str] = None log_file_level: Optional[str] = None def configure_logger(config: Optional[LoggerConfig] = None): """ Configure the logger for Compressed Tensors. This function sets up the console and file logging as per the specified or default parameters. Note: Environment variables take precedence over the function parameters. By default, this function does NOT clear existing loggers or add new handlers, making it safe to use in library code. :param config: The configuration for the logger to use. :type config: LoggerConfig """ logger_config = config or LoggerConfig() # env vars get priority disabled_env = parse_bool_env(os.getenv("COMPRESSED_TENSORS_LOG_DISABLED")) if disabled_env is not None: logger_config.disabled = disabled_env clear_loggers_env = parse_bool_env(os.getenv("COMPRESSED_TENSORS_CLEAR_LOGGERS")) if clear_loggers_env is not None: logger_config.clear_loggers = clear_loggers_env if (console_log_level := os.getenv("COMPRESSED_TENSORS_LOG_LEVEL")) is not None: logger_config.console_log_level = console_log_level.upper() if (log_file := os.getenv("COMPRESSED_TENSORS_LOG_FILE")) is not None: logger_config.log_file = log_file if (log_file_level := os.getenv("COMPRESSED_TENSORS_LOG_FILE_LEVEL")) is not None: logger_config.log_file_level = log_file_level.upper() if logger_config.disabled: logger.disable("compressed_tensors") return logger.enable("compressed_tensors") if logger_config.clear_loggers: logger.remove() if logger_config.console_log_level: # log as a human readable string with the time, function, level, and message logger.add( sys.stdout, level=logger_config.console_log_level.upper(), format="{time} | {function} | {level} - {message}", filter=support_log_once, ) if logger_config.log_file or logger_config.log_file_level: log_file = logger_config.log_file or "compressed_tensors.log" log_file_level = logger_config.log_file_level or "INFO" # log as json to the file for easier parsing logger.add( log_file, level=log_file_level.upper(), serialize=True, filter=support_log_once, ) def parse_bool_env(value: Optional[str]) -> Optional[bool]: """ Parse a boolean environment variable value. Returns: - None if the value is unset or unrecognized - True for recognized truthy tokens - False for recognized falsy tokens Accepts: "1", "true", "True", "TRUE", "yes", "Yes", "YES" as True Accepts: "0", "false", "False", "FALSE", "no", "No", "NO", "" as False """ if value is None: return None value_lower = value.lower().strip() if value_lower in ("1", "true", "yes"): return True elif value_lower in ("0", "false", "no", ""): return False else: return None def support_log_once(record: Dict[str, Any]) -> bool: """ Support logging only once using `.bind(log_once=True)` ``` logger.bind(log_once=False).info("This will log multiple times") logger.bind(log_once=False).info("This will log multiple times") logger.bind(log_once=True).info("This will only log once") logger.bind(log_once=True).info("This will only log once") # skipped ``` """ log_once = record["extra"].get("log_once", False) level = getattr(record["level"], "name", "none") message = hash(str(level) + record["message"]) if log_once and message in _logged_once: return False if log_once: _logged_once.add(message) return True # invoke logger setup on import with default values enabling console logging with INFO # and disabling file logging configure_logger(config=LoggerConfig()) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/modeling/000077500000000000000000000000001521257237700267735ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/modeling/__init__.py000066400000000000000000000002701521257237700311030ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa # isort: off from .kvcache import * from .attention import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/modeling/attention.py000066400000000000000000000111301521257237700313460ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import inspect from typing import Callable, Optional from compressed_tensors.modeling.kvcache import initialize_hooked_kv_cache from compressed_tensors.quantization.lifecycle.forward import forward_quantize from compressed_tensors.utils import getattr_chain from compressed_tensors.utils.internal import InternalModule from torch import Tensor from torch.nn import Module from torch.utils.hooks import RemovableHandle from transformers import PretrainedConfig, PreTrainedModel from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS __all__ = [ "QuantizedAttentionImpl", "initialize_hooked_attention", "register_query_hook", "IMPL_ATTR", ] IMPL_ATTR = "impl" HOOKED_ATTENTION_NAME = "ct_hooked_attention" class QuantizedAttentionImpl(InternalModule): """ QuantizedAttentionImpl module which wraps the functionality of the original attention implementation. Unlike the original attention function, this implementation is a `torch.nn.Module` which can be hooked to trigger transforms and calibration hooks. This module works by being registered as a submodule to attention modules via `initialize_hooked_attention`, registering a new attention implementation function which calls this module, then setting the model attention implementation to the new function. After triggering hooks and quantization, this module calls the original attention implementation function. """ _original_impl = "eager" def __init__(self, config: PretrainedConfig): super().__init__() self.config = config def forward( self, module: Module, query: Tensor, key: Tensor, value: Tensor, *args, **kwargs, ): # quantization quant_args_attr = "quantization_scheme.input_activations" quant_args = getattr_chain(module, quant_args_attr, None) quant_enabled = getattr(module, "quantization_enabled", True) if quant_args is not None and quant_enabled: query = forward_quantize(module, query, "q", quant_args) # original attention return ALL_ATTENTION_FUNCTIONS[QuantizedAttentionImpl._original_impl]( module, query, key, value, *args, **kwargs, ) # ----- initialize ----- # def _hooked_attention(module: Module, *args, **kwargs): assert hasattr(module, IMPL_ATTR), ( f"Using {HOOKED_ATTENTION_NAME} attention implementation, " f"but attention module does not have {IMPL_ATTR} submodule." ) return getattr(module, IMPL_ATTR)(module, *args, **kwargs) def initialize_hooked_attention(model: PreTrainedModel, module: Module): """ Initialize `QuantizedAttentionImpl` and `QuantizedKVCache` instances attached to attention. Assumes that only one model is hooked at a time. :param model: parent model of attention module :param module: attention module to initialize with """ if not hasattr(module, IMPL_ATTR): module.register_module(IMPL_ATTR, QuantizedAttentionImpl(model.config)) if model.config._attn_implementation != HOOKED_ATTENTION_NAME: QuantizedAttentionImpl._original_impl = model.config._attn_implementation original_mask = ALL_MASK_ATTENTION_FUNCTIONS[model.config._attn_implementation] ALL_ATTENTION_FUNCTIONS.register(HOOKED_ATTENTION_NAME, _hooked_attention) ALL_MASK_ATTENTION_FUNCTIONS.register(HOOKED_ATTENTION_NAME, original_mask) model.set_attn_implementation(HOOKED_ATTENTION_NAME) assert model.config._attn_implementation == HOOKED_ATTENTION_NAME initialize_hooked_kv_cache(model, module) # ----- hooks ----- # def register_query_hook( module: Module, hook: Callable[[Module, Tensor], Optional[Tensor]] ) -> RemovableHandle: """ Register a hook which takes post-rope query states as an argument and returns the modified query states or `None` :param module: attention module to add hook to :param hook: query hook function """ impl: QuantizedAttentionImpl = getattr(module, IMPL_ATTR) def _hook(impl: QuantizedAttentionImpl, args, kwargs): bound = inspect.signature(impl.forward).bind(*args, **kwargs) value = hook(module, bound.arguments["query"]) if value is not None: bound.arguments["query"] = value return bound.args, bound.kwargs return impl.register_forward_pre_hook(_hook, with_kwargs=True) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/modeling/kvcache.py000066400000000000000000000135121521257237700307530ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import inspect from typing import Any, Callable, Dict, List, Optional, Tuple from weakref import ReferenceType, ref from compressed_tensors.quantization.lifecycle.forward import forward_quantize from compressed_tensors.utils import getattr_chain from compressed_tensors.utils.internal import InternalModule from torch import Tensor from torch.nn import Module from torch.utils.hooks import RemovableHandle from transformers import Cache, PretrainedConfig, PreTrainedModel __all__ = [ "QuantizedKVCache", "initialize_hooked_kv_cache", "register_key_hook", "register_value_hook", "KV_CACHE_ATTR", ] KV_CACHE_ATTR = "kv_cache" class QuantizedKVCache(InternalModule): """ QuantizedKVCache module which wraps the functionality of any existing kvcache args. Unlike transform Cache instances, this cache is a `torch.nn.Module` which can be hooked to trigger transforms and calibration hooks. This module works by being registered as a submodule to attention modules via `initialize_hooked_kv_cache`, then adding a hook which replaces `past_key_values` kwargs with this module. This module adopts the functionality of the replaced cache, preserving caching functionality such as sliding window attention, ect. :param attn_module: parent attention module """ def __init__(self, config: PretrainedConfig, attn_module: Module): super().__init__() self.config = config self.attn_module = ref(attn_module) # avoid circular reference self.past_key_values: Optional[ReferenceType[Cache]] = None def update(self, *args, **kwargs) -> Tuple[Tensor, Tensor]: return self(*args, **kwargs) def forward( self, key_states: Tensor, value_states: Tensor, *args, **kwargs, ) -> Tuple[Tensor, Tensor]: # quantization module = self.attn_module() quant_args_attr = "quantization_scheme.input_activations" quant_args = getattr_chain(module, quant_args_attr, None) quant_enabled = getattr(module, "quantization_enabled", True) if quant_args is not None and quant_enabled: key_states = forward_quantize(module, key_states, "k", quant_args) value_states = forward_quantize(module, value_states, "v", quant_args) # original cache if self.past_key_values is not None: ret = self.past_key_values().update( key_states, value_states, *args, **kwargs ) else: ret = (key_states, value_states) self.past_key_values = None return ret def add_past_key_values(self, past_key_values: Optional[Cache]): if past_key_values is not None: self.past_key_values = ref(past_key_values) else: self.past_key_values = None # ----- initialize ----- # def _kv_cache_attention_hook( module: Module, args: List[Any], kwargs: Dict[str, Any] ) -> Tuple[List[Any], Dict[str, Any]]: """ Hook which should be called before each quantized attention forward pass. This hook dynamically replaces the `past_key_values` kwarg to the attention forward function. The original kvcache object is assigned to QuantizedKVCache().past_key_values as a weakref to maintain original cache functionality and compute savings """ _past_kv_name = ( "past_key_values" # transformers#39956 if "past_key_values" in inspect.signature(module.forward).parameters else "past_key_value" ) past_key_values: Optional[Cache] = kwargs.get(_past_kv_name, None) cache: QuantizedKVCache = getattr(module, KV_CACHE_ATTR) cache.add_past_key_values(past_key_values) kwargs[_past_kv_name] = cache return args, kwargs def initialize_hooked_kv_cache(model: PreTrainedModel, module: Module): """ Initialize a `QuantizedKVCache` instance attached to attention :param model: parent model of attention module :param module: attention module to initialize with """ if not hasattr(module, KV_CACHE_ATTR): module.register_module(KV_CACHE_ATTR, QuantizedKVCache(model.config, module)) module.register_forward_pre_hook(_kv_cache_attention_hook, with_kwargs=True) # ----- hooks ----- # def register_key_hook( module: Module, hook: Callable[[Module, Tensor], Optional[Tensor]] ) -> RemovableHandle: """ Register a hook which takes post-rope key states as an argument and returns the modified key states or `None` :param module: attention module to add hook to :param hook: key hook function """ kv_cache: QuantizedKVCache = getattr(module, KV_CACHE_ATTR) def _hook(cache: QuantizedKVCache, args, kwargs): bound = inspect.signature(cache.forward).bind(*args, **kwargs) value = hook(module, bound.arguments["key_states"]) if value is not None: bound.arguments["key_states"] = value return bound.args, bound.kwargs return kv_cache.register_forward_pre_hook(_hook, with_kwargs=True) def register_value_hook( module: Module, hook: Callable[[Module, Tensor], Optional[Tensor]] ) -> RemovableHandle: """ Register a hook which takes value states as an argument and returns the modified value states or `None` :param module: attention module to add hook to :param hook: value hook function """ kv_cache: QuantizedKVCache = getattr(module, KV_CACHE_ATTR) def _hook(cache: QuantizedKVCache, args, kwargs): bound = inspect.signature(cache.forward).bind(*args, **kwargs) value = hook(module, bound.arguments["value_states"]) if value is not None: bound.arguments["value_states"] = value return bound.args, bound.kwargs return kv_cache.register_forward_pre_hook(_hook, with_kwargs=True) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/000077500000000000000000000000001521257237700266075ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/README.md000066400000000000000000000755221521257237700301010ustar00rootroot00000000000000# `compressed_tensors` Offload Module ## Overview The `compressed_tensors.offload` module provides a transparent, flexible system for offloading model weights between devices (CPU, GPU, and disk) to enable inference with models larger than available GPU memory. It was designed with the following goals: - **Transparency**: offloading logic is invisible to the model's forward pass — no code changes needed in model implementations. - **Flexibility**: supports CPU offload, disk offload, multi-GPU dispatch, and distributed (multi-process) setups via a single unified interface. - **Performance**: lazy onloading of individual parameters and buffers means that device movement workload is minimized, while the provided utilities encourage keeping offloaded tensors on their offloaded device. - **Interoperability**: can convert to and from HuggingFace `accelerate` offloading, making it compatible with `transformers` loading and saving workflows. --- ## Core Philosophy ### The OffloadCache: Replacing `_parameters` and `_buffers` The central design decision is that offloading is implemented by **replacing a `torch.nn.Module`'s `_parameters` and `_buffers` dictionaries** with `OffloadCache` instances. `OffloadCache` is a `MutableMapping` that behaves like a dict, but with two key differences: - **On write (`__setitem__`):** the tensor is immediately offloaded to its storage device (CPU, disk, etc.). - **On read (`__getitem__`):** the tensor is onloaded (brought back to the execution device) before being returned. Because PyTorch accesses module parameters through `module._parameters[name]` during a forward pass, this substitution is fully transparent — PyTorch naturally triggers onloading whenever it accesses a weight, and offloading occurs automatically when weights are assigned or after the forward pass completes. ``` Normal module: Offloaded module: ┌────────────────────────────────┐ ┌────────────────────────────────┐ │ module._parameters │ │ module._parameters │ │ = { "weight": } │ → │ = OffloadCache { │ └────────────────────────────────┘ │ offload: cpu │ │ onload: cuda:0 │ │ offloaded_values: { │ │ "weight": │ │ } │ │ } │ └────────────────────────────────┘ ``` ### The Forward Wrapper In addition to replacing `_parameters` and `_buffers`, `offload_module` wraps the module's `forward` method. The wrapper moves all input tensors to the `onload_device` before the forward call. This ensures that activations and weights are on the same device without requiring the model's own code to manage device placement. ### Onload/Offload Lifecycle During a normal forward pass: 1. Input tensors arrive on any device. 2. The forward wrapper moves inputs to `onload_device`. 3. PyTorch accesses parameters via `_parameters[name]` → `OffloadCache.__getitem__` → `onload()` is called, moving weights to device. 4. The forward computation runs entirely on `onload_device`. 5. After the forward pass, onloaded weights are released (garbage collected) unless `disable_offloading` is active. ### Global Control Flags Two class-level flags on `OffloadCache` allow callers to suppress device movement across all modules simultaneously: - `offloading_disabled`: weights that have been onloaded are kept in memory; re-accesses are cache hits. - `onloading_disabled`: reads and writes bypass device movement entirely (useful for inspecting raw offloaded tensors). These are exposed as context managers: `disable_offloading()` and `disable_onloading()`. --- ## Module Reference ### `cache/` — OffloadCache Implementations #### `OffloadCache` (base class) — `cache/base.py` The abstract `MutableMapping` that all cache implementations extend. ```python class OffloadCache(MutableMapping, ABC): onload_device: torch.device | str offload_device: torch.device | Literal["disk"] ``` **Key methods:** | Method | Description | |---|---| | `cls_from_device(device)` | Returns the correct `OffloadCache` subclass for the given offload device and distributed state. Automatically selects distributed variants when `dist` is initialized. | | `from_mapping(mapping, onload_device, **kwargs)` | Class method. Constructs an `OffloadCache` from an existing dict (e.g., `module._parameters`), offloading all values immediately. This is the canonical way to attach a cache to a module. | | `onload(offloaded)` | *(abstract)* Given an offloaded tensor, returns the onloaded version. | | `offload(tensor)` | *(abstract)* Given a tensor, offloads it to the storage device and returns a reference. | | `update_offload(offloaded, data)` | *(abstract)* In-place update of an already-offloaded tensor without creating a new storage location. | | `__getitem__(key)` | Onloads the tensor. Uses the onloaded cache if `offloading_disabled` is set. | | `__setitem__(key, value)` | If the key exists and sizes match, calls `update_offload` (in-place update). Otherwise calls `offload` and stores the new reference. | | `disable_offloading()` | *Context manager.* Onloaded tensors are cached in memory; subsequent reads are cache hits. All cached values are released on exit. | | `disable_onloading()` | *Context manager.* All reads return offloaded tensors directly; writes assign directly without triggering device movement. Mainly for debugging. | **When to use `disable_offloading`:** ```python # Without: each access triggers a CPU→GPU copy for _ in range(3): result = module.weight # 3 separate copies # With: first access copies, subsequent reads hit cache with disable_offloading(): for _ in range(3): result = module.weight # 1 copy, 2 cache hits ``` **When to use `disable_onloading`:** ```python # Inspect the raw offloaded tensor (e.g., CPU tensor) without triggering a copy with disable_onloading(): cpu_tensor = module.weight # returns the CPU tensor directly module.weight = new_tensor # assigns without triggering offload ``` --- #### `CPUCache` — `cache/cpu.py` Offloads tensors to CPU RAM. Onloading is a standard `.to(device)` call from CPU to the configured `onload_device` (typically a CUDA device). - **offload**: moves tensor to `torch.device("cpu")`. - **onload**: moves tensor from CPU to `onload_device`. - **update_offload**: `offloaded.copy_(data)` in-place on the CPU tensor. **Use when:** GPU VRAM is insufficient and CPU RAM is available. This is the most common offload strategy. --- #### `DeviceCache` — `cache/device.py` Offloads tensors to a CUDA device. Onloading is typically a no-op (the tensor is already on device), but handles the case where `onload_device` is changed after initialization (e.g., during `set_onload_device` reconfiguration). - **offload**: moves tensor to the device (`self.offload_device = self.onload_device` at init). - **onload**: `send_tensors(offloaded, device=self.onload_device)`. - **update_offload**: in-place copy. **Use when:** a module is permanently resident on a device and you want consistent management via the `OffloadCache` interface (e.g., when `dispatch_model` keeps some modules fully on-device). --- #### `DiskCache` — `cache/disk.py` Offloads tensors to disk as `safetensors` files. In-memory, offloaded tensors are represented as **meta tensors** (tensors with no data, only shape/dtype/stride). A separate `index` dict maps meta tensor identity → disk location. - **offload**: writes tensor to a `.safetensors` file in `offload_dir`, returns a meta tensor. - **onload**: reads the file at the stored path with `safe_open`, reconstructs the tensor. - **update_offload**: overwrites the file with new data. - **`__delitem__`**: removes from index and deletes the file if it was created (not the original model file). **Use when:** neither CPU RAM nor GPU VRAM can hold all weights. Disk offload is the slowest but allows running arbitrarily large models. > **Note:** `DiskCache` requires weights to be stored in `safetensors` format. When loading from non-safetensors formats, weights are onloaded and re-saved. --- #### Distributed Cache Variants Each of the three non-distributed cache types has a distributed counterpart that synchronizes data across ranks during `offload`. ##### `DistributedCPUCache` — `cache/dist_cpu.py` Extends `CPUCache`. On `offload`: 1. Rank 0 creates a CPU tensor and calls `.share_memory_()` to place it in POSIX shared memory (`/dev/shm`). 2. Rank 0 broadcasts the shared memory file handle to all other ranks. 3. All other ranks reconstruct the tensor by attaching to the same shared memory. 4. A `dist.barrier()` ensures rank 0 does not GC the memory before others attach. **Result:** all ranks share the same physical CPU memory for offloaded weights — no redundant copies. **Use when:** running multi-process inference and offloading to CPU. The shared memory avoids `N × model_size` RAM usage. --- ##### `DistributedDeviceCache` — `cache/dist_device.py` Extends `DeviceCache`. On `offload`: 1. Rank 0 moves the tensor to its device. 2. All other ranks create an empty tensor on their local device. 3. `dist.broadcast(as_broadcastable(tensor), src=0)` sends the data to all ranks. **Result:** each rank has its own copy of the tensor on its local GPU. The model is **replicated** across devices. **Use when:** running data-parallel multi-GPU inference where all ranks need identical weights. --- ##### `DistributedDiskCache` — `cache/dist_disk.py` Extends `DiskCache`. On `offload`: 1. Rank 0 writes the safetensors file and broadcasts the file path + weight name + dtype. 2. Other ranks construct a meta tensor and populate their `index` with the broadcasted path. 3. `dist.barrier()` waits for the write to complete before all ranks proceed. **Result:** all ranks share the same disk files for offloaded weights. --- ### `module.py` — Module-Level Offloading #### `offload_module(module, onload_device, offload_device, **kwargs)` The primary function for attaching offloading to a single `torch.nn.Module`. It: 1. Selects the correct `OffloadCache` subclass via `OffloadCache.cls_from_device(offload_device)`. 2. Replaces `module._parameters` and `module._buffers` with cache instances (offloading all existing tensors). 3. Wraps `module.forward` to move input tensors to `onload_device` before the forward call. 4. Stores the original forward function as `module._original_forward_func` so it can be restored. ```python # Offload a single linear layer to CPU, execute on cuda:0 offload_module(layer, onload_device="cuda:0", offload_device="cpu") ``` **When to use:** when you want fine-grained control over which specific modules are offloaded. For model-wide dispatch, prefer `dispatch_model` or `set_onload_device`. > **Note:** Raises `ValueError` if the module is already offloaded. Call `remove_module_offload` first. --- #### `remove_module_offload(module, onload_tensors=False)` Removes offloading from a single module, restoring plain dicts for `_parameters` and `_buffers` and restoring the original forward function. - `onload_tensors=True`: restores weights to the `onload_device` (GPU). Use before running a forward pass without offloading. - `onload_tensors=False` (default): keeps weights on the offload device. Use when freeing memory or re-dispatching. --- #### `unwrap_offload_forward(module)` *(context manager)* Temporarily removes the offload forward wrapper so the underlying forward function can be modified (e.g., by another library). Upon exiting, the offload wrapper is re-applied around any modifications made to `module.forward`. ```python with unwrap_offload_forward(module): module.forward = my_patched_forward # patching the real forward # on exit: offload wrapper is applied around my_patched_forward ``` **Use when:** other code (e.g., PEFT, quantization hooks) needs to wrap the module's real forward without being wrapped around or interfering with the offload wrapper. --- ### `dispatch.py` — Model-Level Dispatch #### `dispatch_model(model, device_memory=None, extra_memory=None, no_split_modules=None)` The highest-level dispatch function. Automatically determines how to distribute a model across all available CUDA devices, maximizing GPU utilization with a CPU offload fallback. **Algorithm:** 1. Queries all available CUDA devices for their total memory. 2. Uses `get_module_sizes` to measure each non-splittable module. 3. Performs a binary search to find the maximum `extra_memory` (reserved for activations/KV cache) such that all modules fit on device. 4. If the model fits entirely on GPU, dispatches with no offloading. 5. If not, falls back to placing as many modules as possible on GPU and offloading the rest to CPU. ```python # Dispatch to all GPUs, reserve memory for KV cache automatically model = dispatch_model(model) # Dispatch with explicit device memory constraints model = dispatch_model(model, device_memory={torch.device("cuda:0"): 16e9}) ``` **Parameters:** - `device_memory`: optional mapping of device → available bytes. Defaults to querying all CUDA devices. - `extra_memory`: bytes to reserve for activations. If `None`, estimated from model config. - `no_split_modules`: names of module classes that cannot be split across devices. Defaults to `model._no_split_modules` if available. **When to use:** the primary entry point for production deployment — automatic, memory-aware, multi-GPU dispatch. --- #### `set_onload_device(model, onload_device)` A lighter-weight dispatch that moves all modules in a model to the same `onload_device`, without changing where weights are stored. For modules not yet offloaded, it offloads them to their current device. ```python # Move all execution to cuda:0, keeping offloads unchanged model = set_onload_device(model, onload_device="cuda:0") ``` **When to use:** when you have already loaded a model with weights in the right place (e.g., via `load_offloaded_model`) and just need to set the execution device. Less powerful than `dispatch_model` but simpler. > **Note:** `offload_model` is a deprecated alias for this function. --- #### `dispatch_with_map(model, device_map, offload_dir=None)` Dispatches a model according to an explicit `DeviceMap` — a dict mapping module name to `(onload_device, offload_device)` tuples. ```python device_map = { "model.embed_tokens": ("cuda:0", "cuda:0"), "model.layers.0": ("cuda:0", "cpu"), "model.layers.1": ("cuda:0", "disk"), } dispatch_with_map(model, device_map, offload_dir="/tmp/offload") ``` **When to use:** when you have a specific, manually-determined dispatch plan. Also used internally by `from_accelerate`. --- #### `get_device_map(model, default_device=cpu) → DeviceMap` Introspects a dispatched model and returns a `DeviceMap` describing each module's current `(onload_device, offload_device)`. ```python device_map = get_device_map(model) # -> {"model.layers.0": (cuda:0, cpu), "model.layers.1": (cuda:0, disk), ...} ``` **Use when:** you need to inspect, serialize, or replicate the dispatch configuration of a model. --- #### `remove_dispatch(module, onload_tensors=False) → module` Removes all offloading from every submodule of `module`. ```python remove_dispatch(model, onload_tensors=True) # bring all weights to GPU before saving ``` --- #### `get_device_memory() → dict[torch.device, int]` Returns a mapping of all available CUDA devices to their total memory (in bytes). In a distributed context, returns only the local rank's GPU. Returns an empty dict if no CUDA devices are available. --- ### `load.py` — Loading with Offloading #### `load_offloaded_model(model_class=None, extra_cpu_mem=5e9)` *(context manager)* A context manager that patches `from_pretrained` on transformers model classes. Within the context, calling `from_pretrained` will: 1. On rank 0: load the model using `accelerate`'s device offloading. 2. On other ranks: load the model on the meta device (no actual weights). 3. After loading: call `from_accelerate` to convert to `compressed_tensors` offloading and share weights across ranks. ```python from transformers import AutoModelForCausalLM with load_offloaded_model(): model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-3-8B", device_map="auto_offload", # CT-specific: only cpu/disk, no GPU ) ``` **Special `device_map` value: `"auto_offload"`** In addition to standard `device_map` options (`"auto"`, `"cpu"`, etc.), `load_offloaded_model` supports `device_map="auto_offload"`. This restricts accelerate's auto-mapping to only use CPU and disk (no GPU VRAM), which is useful when you want the GPU to be fully available for activations during inference. **Parameters:** - `model_class`: explicit class to patch (e.g., `AutoModelForCausalLM`). If `None`, patches all `transformers` model classes or `PreTrainedModel` subclasses visible in the caller's global scope. Providing an explicit class is more efficient and recommended when the target class is known. - `extra_cpu_mem`: bytes to reserve in CPU RAM for non-weight operations (default 5 GB). Memory estimates are automatically computed for both distributed and non-distributed setups. **When to use:** the recommended entry point for loading offloaded models from pretrained checkpoints. Handles distributed loading correctly, sharing weights across ranks with minimal memory overhead. --- ### `convert/` — Accelerate Interoperability The `convert` submodule provides bidirectional conversion between `compressed_tensors` offloading and HuggingFace `accelerate` offloading. This is needed because the `transformers` ecosystem (loading, saving, PEFT, etc.) expects accelerate's `AlignDevicesHook`-based offloading. #### `from_accelerate(model) → (device_map, offload_dir)` — `convert/from_accelerate.py` Converts a model from accelerate offloading to `compressed_tensors` offloading. Called automatically by `load_offloaded_model`. Process: 1. Calls `remove_accelerate(model)` to strip accelerate hooks and collect device/offload information. 2. In a distributed context, broadcasts the `device_map` and `offload_dir` from rank 0 to all ranks. 3. Calls `dispatch_with_map` to apply `compressed_tensors` offloading. For disk-offloaded modules, the accelerate disk index is copied into `DiskCache.index` so tensors can be read via `safe_open` without re-writing files. **Returns:** `(device_map, offload_dir)` — the device map and optional disk offload directory. --- #### `remove_accelerate(model) → (device_map, offload_dir)` Strips all `AlignDevicesHook` instances from every module in the model. For each module, calls `remove_accelerate_from_module`. Removes `model.hf_device_map` if present. --- #### `remove_accelerate_from_module(module) → (onload_device, offload_device, offload_dir)` Removes the accelerate `AlignDevicesHook` from a single module. Handles three offload cases: - **CPU/device offload**: tensors in `hook.weights_map` are reassigned directly. - **Disk offload**: meta tensors are kept as-is; the accelerate disk index is copied into `DiskCache.index`. - **No offload**: no-op; returns the module's current device for both onload and offload. This is a zero-copy operation — no new tensors are allocated. --- #### `to_accelerate(model) → hf_device_map` — `convert/to_accelerate.py` Converts a `compressed_tensors`-offloaded model back to accelerate offloading. This is necessary before calling `model.save_pretrained()`, since `transformers`'s saving logic understands `AlignDevicesHook` but not `OffloadCache`. ```python # Before saving: to_accelerate(model) model.save_pretrained("./output") # After saving, convert back for inference: from_accelerate(model) ``` For disk-offloaded modules, creates an `OffloadedWeightsLoader` backed by the `DiskCache` index. For CPU/device-offloaded modules, uses an in-memory dict as the weights map. Sets `model.hf_device_map` as expected by `transformers`. --- #### `to_accelerate_module(module, name=None, hf_disk_index=None) → str` Lower-level version of `to_accelerate` that converts a single module. Returns the string representation of the module's offload device. --- ### `__init__.py` — Top-Level API The top-level `compressed_tensors.offload` namespace re-exports the most important functions and adds several convenience utilities. #### `disable_offloading()` *(context manager)* Disables offloading globally. Within this context, onloaded tensors are cached so that repeated accesses do not trigger additional device-to-device copies. All cached values are released on exit. ```python with disable_offloading(): # All subsequent accesses to offloaded parameters are cache hits for token in sequence: output = model(token) ``` **Use when:** processing multiple tokens or batches without wanting to offload between each call. --- #### `disable_onloading()` *(context manager)* Disables onloading globally. Parameter reads return the raw offloaded tensors; assignments do not trigger offloading. Primarily used for debugging and internal utilities. ```python with disable_onloading(): raw_cpu_tensor = module.weight # returns the CPU tensor module.weight = new_tensor # assigns without triggering copy ``` --- #### `update_offload_parameter(module, name, data)` Updates an offloaded parameter or buffer in-place. Works for both offloaded and non-offloaded modules. For offloaded modules: calls the cache's `update_offload` (copies into shared CPU memory, overwrites safetensors file, etc.). For non-offloaded modules: uses `param.copy_(data)`. ```python update_offload_parameter(layer, "weight", quantized_weight) ``` > **Caution:** does not guard against multiple ranks writing simultaneously. The caller is responsible for ensuring single-rank writes. Also does not broadcast onloaded values to other ranks. **Use when:** performing quantization, weight updates, or any post-load modification to model weights. --- #### `get_execution_device(module, default=None) → torch.device | "disk"` Returns the device that input tensors should be moved to before running the module's forward pass. For offloaded modules: returns `module._parameters.onload_device`. For non-offloaded modules: returns the device of the first parameter/buffer. --- #### `get_offloaded_device(module, default=None) → torch.device | "disk"` Returns the device where the module's weights reside when not in use (between forward passes). For offloaded modules: returns `module._parameters.offload_device`. For non-offloaded modules: returns the device of the first parameter/buffer. --- #### `register_offload_module(base, name, module)` Registers a new submodule on a parent module, propagating offloading if the parent is offloaded. ```python # Attaches `new_layer` with the same offload config as `model.layers[0]` register_offload_module(model, "new_layer", new_layer) ``` **Use when:** dynamically adding submodules to an already-dispatched model (e.g., adding LoRA adapters after dispatch). --- #### `align_modules(modules, execution_device=None)` *(deprecated)* #### `align_module_device(module, execution_device=None)` *(deprecated)* Legacy context managers for temporarily moving a module's parameters to a specific device. Kept for backwards compatibility. For offloaded modules: temporarily overrides `onload_device` and disables offloading. For non-offloaded modules: moves tensors manually and restores them on exit. > Use `disable_offloading()` combined with `get_execution_device()` for new code. --- ### `dist_utils.py` — Distributed Utilities #### `is_distributed() → bool` Returns `True` if `torch.distributed` is available and has been initialized. #### `is_rank0() → bool` Returns `True` if not in a distributed context, or if the current process is the source rank (typically rank 0). !!! warning This argument is deprecated, please use `compressed_tensors.distributed.utils.is_source_process`. !!! #### `init_dist()` Initializes the distributed process group using NCCL and environment variables set by `torchrun`. Sets the CUDA device to the local rank. Requires the `TORCHELASTIC_RUN_ID` environment variable (set automatically by `torchrun`). ```bash torchrun --nproc-per-node=4 my_script.py ``` ```python # At the top of my_script.py: from compressed_tensors.offload import init_dist init_dist() ``` #### `as_broadcastable(tensor) → tensor` Returns a view of the tensor compatible with `dist.broadcast`. Works around an NCCL limitation: FP8 dtypes (`float8_e4m3fn`, etc.) cannot be broadcast on pre-Hopper hardware. FP8 tensors are reinterpreted as `uint8` for the broadcast and share the same underlying storage. --- ### `utils.py` — Internal Utilities These utilities are used internally throughout the module. #### `send_tensors(value, *args, **kwargs) → value` Recursively traverses a nested structure (tensor, list, tuple, dict, or dataclass) and calls `.to(*args, **kwargs)` on all tensors found. Returns a new structure with the moved tensors, preserving the original's type and `__dict__`. Used by the forward wrapper to move input tensors to the execution device. #### `get_module_device(module, default=None) → torch.device` Returns the device of the first parameter or buffer of a module. Falls back to `default` if the module has no tensors. #### `move_module_tensor(module, name, device)` Moves a single named parameter or buffer to a new device. #### `module_size(module, recurse=True) → int` Returns the total size in bytes of all parameters and buffers in a module. Uses `disable_onloading` to avoid triggering device movement. #### `get_module_sizes(model, no_split_modules=()) → list[(module, int)]` Returns a flat list of `(module, size_bytes)` for all non-splittable modules in the model. A module is considered non-splittable if it has direct parameters or if its class name is in `no_split_modules`. Used by `dispatch_model` to pack modules onto devices. #### `to_empty(tensor, **kwargs) → tensor` Creates an empty tensor with the same shape, dtype, and subclass as the input tensor. Equivalent to `torch.empty_like` but preserves tensor subclass and `__dict__`. #### `to_tensor(dst, src) → dst` Copies the subclass, `__dict__`, and `requires_grad` from `src` into `dst`. Used during accelerate conversion to "convert" an accelerate tensor into the corresponding `compressed_tensors` tensor without copying data. --- ## Common Usage Patterns ### 1. Loading a Large Model for Inference ```python from transformers import AutoModelForCausalLM from compressed_tensors.offload import load_offloaded_model # Recommended: explicit class patching with load_offloaded_model(model_class=AutoModelForCausalLM): model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-3.1-70B", device_map="auto", # or "auto_offload" to restrict to cpu/disk only torch_dtype="bfloat16", ) # model is now using compressed-tensors offloading # Alternative: automatic patching (patches all model classes in scope) with load_offloaded_model(): model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-3.1-70B", device_map="auto", torch_dtype="bfloat16", ) ``` ### 2. Multi-GPU Dispatch ```python from compressed_tensors.offload import dispatch_model model = ... # load on CPU or meta device model = dispatch_model(model) # automatically fills available GPUs, offloads remainder to CPU ``` ### 3. Distributed Inference with torchrun ```python # run with: torchrun --nproc-per-node=4 script.py from transformers import AutoModelForCausalLM from compressed_tensors.offload import init_dist, load_offloaded_model init_dist() # initialize NCCL with load_offloaded_model(model_class=AutoModelForCausalLM): model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-3.1-70B", device_map="auto", torch_dtype="bfloat16", ) # rank 0 loads weights; other ranks get weights via broadcast # all ranks share CPU memory via /dev/shm ``` ### 4. Updating Weights After Loading (e.g., Post-Quantization) ```python from compressed_tensors.offload import update_offload_parameter for name, module in model.named_modules(): if hasattr(module, "weight"): quantized = quantize(module.weight) update_offload_parameter(module, "weight", quantized) ``` ### 5. Saving an Offloaded Model ```python from compressed_tensors.offload import to_accelerate, from_accelerate # Convert to accelerate for saving to_accelerate(model) model.save_pretrained("./output") # Convert back for continued inference from_accelerate(model) ``` ### 6. Manual Dispatch with a Device Map ```python from compressed_tensors.offload import dispatch_with_map device_map = { "model.embed_tokens": ("cuda:0", "cuda:0"), "model.layers.0": ("cuda:0", "cpu"), "model.layers.1": ("cuda:0", "disk"), # ... } dispatch_with_map(model, device_map, offload_dir="/tmp/weights") ``` ### 7. Inspecting Dispatch Configuration ```python from compressed_tensors.offload import get_device_map, get_execution_device, get_offloaded_device device_map = get_device_map(model) exec_dev = get_execution_device(model.layers[0]) # cuda:0 offload_dev = get_offloaded_device(model.layers[0]) # cpu ``` --- ## Architecture Diagram ``` compressed_tensors.offload ├── load.py load_offloaded_model() │ └── calls from_accelerate after loading │ ├── dispatch.py dispatch_model(), set_onload_device(), dispatch_with_map() │ └── calls offload_module() for each module │ ├── module.py offload_module(), remove_module_offload() │ └── replaces _parameters/_buffers with OffloadCache │ wraps module.forward to move inputs │ ├── cache/ │ ├── base.py OffloadCache (abstract MutableMapping) │ ├── cpu.py CPUCache │ ├── device.py DeviceCache │ ├── disk.py DiskCache │ ├── dist_cpu.py DistributedCPUCache (shared memory) │ ├── dist_device.py DistributedDeviceCache (broadcast) │ └── dist_disk.py DistributedDiskCache (shared files) │ ├── convert/ │ ├── from_accelerate.py from_accelerate() — converts from HF accelerate │ ├── to_accelerate.py to_accelerate() — converts to HF accelerate │ └── helpers.py norm_device(), get_tensors() │ ├── dist_utils.py is_distributed(), init_dist(), as_broadcastable() ├── utils.py send_tensors(), module_size(), get_module_sizes() └── __init__.py public API + context managers + parameter utils ``` vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/__init__.py000066400000000000000000000230501521257237700307200ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib from collections.abc import Iterable from typing import Literal import torch from compressed_tensors.distributed.utils import set_source_process from compressed_tensors.offload.cache import OffloadCache from compressed_tensors.offload.convert import from_accelerate, to_accelerate from compressed_tensors.offload.dispatch import ( # noqa: F401 dispatch_model, dispatch_with_map, get_device_map, offload_model, remove_dispatch, set_onload_device, ) from compressed_tensors.offload.dist_utils import ( as_broadcastable, init_dist, is_distributed, is_rank0, ) from compressed_tensors.offload.load import load_offloaded_model from compressed_tensors.offload.module import offload_module, unwrap_offload_forward from compressed_tensors.offload.utils import ( as_single_threaded, get_module_device, move_module_tensor, to_meta, ) from compressed_tensors.utils.helpers import deprecated, patch_attr __all__ = [ # dispatch models "set_onload_device", "offload_model", # deprecated, use set_onload_device "dispatch_model", "remove_dispatch", "dispatch_with_map", "get_device_map", # dispatch modules "offload_module", "get_cache_kwargs", # deprecated, use get_cache_init_kwargs "get_cache_init_kwargs", # accelerate conversion "load_offloaded_model", "from_accelerate", "to_accelerate", # control movement "disable_onloading", "disable_offloading", # manipulate parameters "update_offload_parameter", "get_execution_device", "get_offloaded_device", "register_offload_module", # manipulate forward "unwrap_offload_forward", # backwards compatibility: should be deprecated "align_modules", "align_module_device", # utilities "is_distributed", "is_rank0", "init_dist", "as_broadcastable", "as_single_threaded", "set_source_process", "to_meta", "get_cache_init_kwargs", ] @contextlib.contextmanager def disable_offloading(): """ When offloading is disabled, onloaded tensors remain onloaded in memory until exit ``` with OffloadCache.disable_offloading(): ... = cache["weight"] ... = cache["weight"] # cache hit ... = cache["weight"] # cache hit # upon exit, all onloaded weights are released ``` """ with OffloadCache.disable_offloading(): yield @contextlib.contextmanager def disable_onloading(): """ When onloading is disabled, tensors are not offloaded on access, and assignments do not trigger offloading. This is mostly used to disable device movement for debugging ``` with OffloadCache.disable_onloading(): tensor = ... cache["weight"] = tensor # assignments do not trigger onloading cache["weight"] is tensor # tensor remains offloaded ``` """ with OffloadCache.disable_onloading(): yield def update_offload_parameter(module: torch.nn.Module, name: str, data: torch.Tensor): """ Update the offload and onload data of an existing parameter/buffer. Supports both parameters of both offloaded modules and non-offloaded modules. NOTE: This function does not guard against multiple processes writing to offload at the same time. It is the responsibility of the caller to ensure that, for any parameter/buffer, only one rank calls this function at a time. NOTE: This function does not update onloaded values across ranks. The caller is responsible for broadcasting any updates to other ranks, if they are onloaded. :param module: module containing the parameter to update :param name: name of module parameter to update :param data: tensor to update parameter with """ if isinstance(module._parameters, OffloadCache): # | Component | Update Implementation | # | --------- | --------------------------- | # | CPU | Copy into shared cpu memory | # | Disk | Write file to disk | # | Device | Copy into local device | # | --------- | --------------------------- | # all implementations update onloaded data if applicable if name in module._parameters: cache = module._parameters elif name in module._buffers: cache = module._buffers else: raise AttributeError(f"{type(module)} has no attribute {name}") # triggers update if shapes match cache[name] = data else: with torch.no_grad(): getattr(module, name).copy_(data) def get_execution_device( module: torch.nn.Module, default: torch.device | None = None ) -> torch.device | Literal["disk"]: """ Get the device which inputs should be moved to before module execution. :param module: module to check, may be offloaded :return: onload device of module """ if isinstance(module._parameters, OffloadCache): return module._parameters.onload_device else: return get_module_device(module, default) def get_offloaded_device( module: torch.nn.Module, default: torch.device | None = None ) -> torch.device | Literal["disk"]: """ :param module: module to check :return: device module is offloaded to onto after forward pass """ if isinstance(module._parameters, OffloadCache): return module._parameters.offload_device else: return get_module_device(module, default) @deprecated("compressed_tensors.offload::get_cache_init_kwargs") def get_cache_kwargs(module: torch.nn.Module, default: dict | None = None) -> dict: """ Get any ancillary kwargs needed for the module OffloadCache :param module: module to check :param default: dictionary of kwargs to update with additional arguments :return: dict of cache kwargs """ kwargs = default.copy() if default is not None else {} if isinstance(module._parameters, OffloadCache) and hasattr( module._parameters, "offload_dir" ): kwargs["offload_dir"] = module._parameters.offload_dir return kwargs def get_cache_init_kwargs( module: torch.nn.Module, default: dict | None = None, ) -> dict: """ Get all kwargs needed to initialize an OffloadCache with the same settings as the module. :param module: module to extract cache initialization kwargs from :param default: default kwargs to use as a base (can include onload_device, offload_device, etc.) :return: dict of kwargs for offload_module or cache constructor, including onload_device, offload_device, and any additional cache-specific kwargs """ kwargs = default.copy() if default is not None else {} kwargs["onload_device"] = get_execution_device(module, kwargs.get("onload_device")) kwargs["offload_device"] = get_offloaded_device( module, kwargs.get("offload_device") ) if hasattr(module._parameters, "offload_dir"): kwargs["offload_dir"] = module._parameters.offload_dir return kwargs def register_offload_module(base: torch.nn.Module, name: str, module: torch.nn.Module): """ Register a submodule with offloading if the parent module is offloaded :param base: module to attach submodule to :param name: name of submodule :param module: submodule to attach """ cache = base._parameters if isinstance(cache, OffloadCache): kwargs = get_cache_init_kwargs(base) offload_module(module, **kwargs) base.register_module(name, module) """ Implemented for backwards compatibility """ @contextlib.contextmanager def align_modules( modules: torch.nn.Module | Iterable[torch.nn.Module], execution_device: torch.device | None = None, ): """ Context manager for onloading modules to a device, and disabling onload and offload attempts triggered by forward calls. Used for sequential onloading of layers :param modules: `torch.nn.Module` or iterable of `torch.nn.Module`s to onload :param execution_device: device to onload to """ with contextlib.ExitStack() as stack: for module in modules: stack.enter_context(align_module_device(module, execution_device)) yield @contextlib.contextmanager def align_module_device( module: torch.nn.Module, execution_device: torch.device | None = None ): """ Context manager that moves a module's parameters to the specified execution device. :param module: Module with parameters to align :param execution_device: If provided, overrides the module's execution device within the context. Otherwise, use hook execution device or pass """ if isinstance(module._parameters, OffloadCache): assert isinstance(module._buffers, OffloadCache) with module._parameters.disable_offloading(): if execution_device is not None: with ( patch_attr(module._parameters, "onload_device", execution_device), patch_attr(module._buffers, "onload_device", execution_device), ): yield else: yield else: original_device = {} for name, param in module.named_parameters(recurse=False): original_device[name] = param.device move_module_tensor(module, name, execution_device) try: yield finally: for name, param in module.named_parameters(recurse=False): device = original_device[name] move_module_tensor(module, name, device) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/cache/000077500000000000000000000000001521257237700276525ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/cache/__init__.py000066400000000000000000000006361521257237700317700ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .base import OffloadCache from .cpu import CPUCache from .device import DeviceCache from .disk import DiskCache from .dist_cpu import DistributedCPUCache from .dist_device import DistributedDeviceCache from .dist_disk import DistributedDiskCache from .utils import catch_cpu_mem_error vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/cache/base.py000066400000000000000000000244501521257237700311430ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib from abc import ABC, abstractmethod from collections.abc import Hashable, MutableMapping from typing import ClassVar, Literal import torch import torch.distributed as dist from compressed_tensors.utils import is_accelerator_type class OffloadCache(MutableMapping, ABC): """ Base class for offload caches. Subclasses must implement `offload` and `onload`. Instances have similar behavior to dicts, except that tensors are offloaded when assigned and onloaded when accessed. Typical usage: ``` module._parameters = cache_cls.from_mapping( module._parameters, onload_device, offload_device ) tensor = ... module._parameters["name"] = tensor # tensor is offloaded onloaded_tensor = module._parameters["name"] # tensor is onloaded ``` This class implements two contexts for more fine-grained control of device movement: `OffloadCache.disable_offloading` and `OffloadCache.disable_onloading`. For more info, see `compressed_tensors.offload::(disable_offloading|disable_onloading)` """ onload_device: torch.device | str offload_device: torch.device | Literal["disk"] # global flags for disabling offloading_disabled: ClassVar[bool] = False onloading_disabled: ClassVar[bool] = False # names -> offloaded tensors (populated from _parameters or _buffers) offloaded_values: dict[Hashable, torch.Tensor] # offloaded tensors -> onloaded tensors (only when offloading is disabled) keep_onloaded_values: ClassVar[dict[torch.Tensor, torch.Tensor]] = dict() @classmethod def cls_from_device( cls, device: torch.device | str | Literal["disk"] | None = None, ) -> type["OffloadCache"]: """ Get the subclass which implements offloading for the given `offload_device`. Use `torch.distributed` to detect if the environment is distributed :param device: offload device used to find subclass :return: subclass of `OffloadCache` """ from compressed_tensors.offload.cache.cpu import CPUCache from compressed_tensors.offload.cache.device import DeviceCache from compressed_tensors.offload.cache.disk import DiskCache from compressed_tensors.offload.cache.dist_cpu import DistributedCPUCache from compressed_tensors.offload.cache.dist_device import DistributedDeviceCache from compressed_tensors.offload.cache.dist_disk import DistributedDiskCache device_type = torch.device(device).type if device != "disk" else "disk" distributed = dist.is_available() and dist.is_initialized() match (device_type, distributed): case ("cpu", False): return CPUCache case ("cpu", True): return DistributedCPUCache case (device, False) if is_accelerator_type(device): return DeviceCache case (device, True) if is_accelerator_type(device): return DistributedDeviceCache case ("disk", False): return DiskCache case ("disk", True): return DistributedDiskCache case _: raise NotImplementedError( f"Offload of type {device_type} and " f"distributed={distributed} has not been implemented" ) @classmethod def from_mapping( cls, mapping: MutableMapping[Hashable, torch.Tensor | None], onload_device: torch.device | str, offload_device: "torch.device | str | Literal['disk'] | None" = None, **kwargs, ): """ Initialize an instance from a given mapping, typically `Module._parameters` or `Module._buffers`. Mapping values will be offloaded :param mapping: mapping used to populate cache :param onload_device: device which tensors will be onloaded to :param offload_device: device to offload tensors to. For DeviceCache, this sets the offload target (defaults to onload_device if not provided). For CPUCache and DiskCache, this is validated against the fixed offload_device if provided. :param \\**kwargs: keyword arguments for cache constructor """ instance = cls( onload_device=onload_device, offload_device=offload_device, **kwargs ) instance.offloaded_values = { name: instance.offload(tensor) for name, tensor in mapping.items() } return instance def __init__( self, onload_device: torch.device | str, offload_device: torch.device | str | Literal["disk"] | None = None, ): super().__init__() self.onload_device = onload_device self.offloaded_values = dict() # Validate offload_device for subclasses with a fixed offload_device # (CPUCache, DiskCache). DeviceCache sets offload_device after super().__init__ # so this check only applies when offload_device is a class attribute. if offload_device is not None and hasattr(type(self), "offload_device"): assert str(offload_device) == str(self.offload_device) @abstractmethod def onload(self, offloaded: torch.Tensor | None) -> torch.Tensor | None: """ Given an offloaded tensor, returns that tensor after onloading :param offloaded: offloaded tensor :return: onloaded tensor """ raise NotImplementedError() @abstractmethod def offload(self, tensor: torch.Tensor | None) -> torch.Tensor | None: """ Given a tensor, returns that tensor after offloading :param tensor: tensor to offload :return: offloaded tensor """ raise NotImplementedError() @abstractmethod def update_offload(self, offloaded: torch.Tensor, data: torch.Tensor | None): """ Update the data of an offloaded tensor NOTE: Operation is performed asynchronously. If you need the offloaded value to update across all ranks, call `dist.barrier()` after calling this function :param tensor: offloaded tensor to update :param data: new tensor data """ raise NotImplementedError() def __getitem__(self, key: Hashable) -> torch.Tensor: """ Onload a tensor If called within the `disable_offloading` context, a strong reference of the onloaded tensor is kept so that future accesses will not require device movement :param key: name of tensor to access :return: onloaded tensor """ offloaded = self.offloaded_values[key] # when onloading is disabled, offloaded tensors can be accessed directly if offloaded is None or self.onloading_disabled: return offloaded # check for cache hit if offloaded in self.keep_onloaded_values: return self.keep_onloaded_values[offloaded] # onload value onloaded = self.onload(offloaded) # when offloading is disabled, populate cache if self.offloading_disabled: self.keep_onloaded_values[offloaded] = onloaded return onloaded def __setitem__(self, key: Hashable, value: torch.Tensor | None): """ Update the offloaded and onloaded values if the key exists, otherwise offload the value and add it to the cache. If called within the `disable_onloading` context, the value is assigned directly :param key: name of tensor :param value: tensor value to offload """ # when onloading is disabled, parameters can be access and assigned directly if self.onloading_disabled: self.offloaded_values[key] = value return # if the key already exists, update with the new value offloaded = self.offloaded_values.get(key, None) if offloaded is not None and torch.is_same_size(offloaded, value): self.update_offload(offloaded, value) onloaded = self.keep_onloaded_values.get(offloaded, None) if onloaded is not None and onloaded is not offloaded: onloaded.copy_(value) # if the key does not exist (or the value is None), offload the new value else: self.offloaded_values[key] = self.offload(value) def __delitem__(self, key: Hashable): """ Remove the offloaded tensor associated with `key`. Any references to its onloaded tensors held by this class are invalidated. :param key: name of tensor to invalidate """ offloaded = self.offloaded_values[key] del self.offloaded_values[key] # remove strong ref if offloaded in self.keep_onloaded_values: del self.keep_onloaded_values[offloaded] def __contains__(self, key) -> bool: return key in self.offloaded_values def __iter__(self): return iter(self.offloaded_values) def __len__(self): return len(self.offloaded_values) @classmethod @contextlib.contextmanager def disable_offloading(cls): """ Context to disable all offloading for offloaded modules which share this cache. After a weight has been fetched once, that onloaded value is cached and subsequent fetches will leverage the cache, reducing device movement """ if not OffloadCache.offloading_disabled: restore_value = OffloadCache.offloading_disabled OffloadCache.offloading_disabled = True yield OffloadCache.offloading_disabled = restore_value OffloadCache.keep_onloaded_values.clear() else: yield @classmethod @contextlib.contextmanager def disable_onloading(cls): """ Context to disable all onloading for offloaded modules which share this cache. This is mostly used for debugging purposes, and allows the caller to directly inspect offloaded tensors and directly assign offloaded tensors without copying """ if not OffloadCache.onloading_disabled: restore_value = OffloadCache.onloading_disabled OffloadCache.onloading_disabled = True yield OffloadCache.onloading_disabled = restore_value else: yield vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/cache/cpu.py000066400000000000000000000030461521257237700310160ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.offload.cache.base import OffloadCache from compressed_tensors.offload.cache.utils import catch_cpu_mem_error from compressed_tensors.offload.utils import send_tensors class CPUCache(OffloadCache): """ Handles offloading and onloading tensors from/to cpu memory """ offload_device = torch.device("cpu") def __init__(self, onload_device: torch.device | str, offload_device=None): super().__init__(onload_device, offload_device=offload_device) def onload(self, offloaded: torch.Tensor | None) -> torch.Tensor | None: """ Onload a tensor from cpu to device :param key: cpu tensor to onload :return: device tensor """ return send_tensors(offloaded, device=self.onload_device, copy=False) @catch_cpu_mem_error def offload(self, tensor: torch.Tensor | None) -> torch.Tensor | None: """ Offload a tensor from any device to cpu :param value: tensor on any device :return: tensor on cpu """ return send_tensors(tensor, device=self.offload_device, copy=False) @torch.no_grad() def update_offload(self, offloaded: torch.Tensor, data: torch.Tensor | None): """ Update the offloaded cpu value with new data :param offloaded: cpu tensor to update :param data: new data to copy from """ if data is not None: offloaded.copy_(data) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/cache/device.py000066400000000000000000000037011521257237700314640ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import TYPE_CHECKING, Literal, Optional import torch from compressed_tensors.offload.cache.base import OffloadCache from compressed_tensors.offload.utils import send_tensors if TYPE_CHECKING: from torch._prims_common import DeviceLikeType class DeviceCache(OffloadCache): """ Handles offloading and onloading tensors from/to device memory. Onloading tensors is typically a no-op (except onload device has been modified). """ def __init__( self, onload_device: "DeviceLikeType", offload_device: Optional["DeviceLikeType | Literal['disk']"] = None, ): super().__init__(onload_device, offload_device=offload_device) if offload_device is not None: self.offload_device = torch.device(offload_device) else: self.offload_device = self.onload_device def onload(self, offloaded: torch.Tensor | None) -> torch.Tensor | None: """ Typically a no op, except when onload device has been modified :param key: device tensor to onload :return: device tensor """ # move because onload_device might be modified after init return send_tensors(offloaded, device=self.onload_device, copy=False) def offload(self, tensor: torch.Tensor | None) -> torch.Tensor | None: """ Offload a tensor to the device :param value: tensor on any device :return: tensor on device """ return send_tensors(tensor, device=self.offload_device, copy=False) @torch.no_grad() def update_offload(self, offloaded: torch.Tensor, data: torch.Tensor | None): """ Update the device value with new data :param offloaded: device tensor to update :param data: new data to copy from """ if data is not None: offloaded.copy_(data) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/cache/disk.py000066400000000000000000000174051521257237700311650ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os from pathlib import Path from typing import TYPE_CHECKING, Literal, Optional import torch import torch.distributed as dist from compressed_tensors.distributed import is_source_process from compressed_tensors.offload.cache import OffloadCache from compressed_tensors.offload.utils import send_tensors, to_tensor from compressed_tensors.utils import is_accelerator_type from safetensors import safe_open from safetensors.torch import save_file if TYPE_CHECKING: from torch._prims_common import DeviceLikeType class DiskCache(OffloadCache): """ Handles offloading and onloading tensors from/to disk. Tensors usually start as a key in safetensors file, converted by (TODO NAME). New or updated tensors are written to new safetensors files in `offload_dir`. Tensors are stored in memory as meta tensors. The mapping between offloaded meta tensors and their locations on disk is defined by `index`. """ offload_device = "disk" # offloaded tensors -> weight info index: dict[torch.Tensor, dict[str, str]] = dict() # directory where new tensors are written to offload_dir: str _ct_file_prefix = "ct_disk_cache" def __init__( self, onload_device: torch.device, offload_device: Optional["DeviceLikeType | Literal['disk']"] = None, offload_dir: Optional[str] = None, ): super().__init__(onload_device, offload_device=offload_device) if offload_device is not None: assert str(offload_device) == str(self.offload_device) if offload_dir is None: raise ValueError( "Must provide an `offload_dir` to perform disk offloading " "(add `offload_folder` argument to `from_pretrained`)" ) # Resolve relative paths to absolute paths for symlink creation self.offload_dir = Path(offload_dir).resolve() def onload(self, offloaded: torch.Tensor | None) -> torch.Tensor | None: """ Onload a tensor from disk/meta to device :param offloaded: meta tensor to onload :return: device tensor, read from disk """ if offloaded is None: return None weight_info = self.index[offloaded] device = _get_safe_open_device(self.onload_device) with safe_open( weight_info["safetensors_file"], framework="pt", device=device ) as file: onloaded = file.get_tensor(weight_info["weight_name"]) onloaded = to_tensor(onloaded, offloaded) onloaded = onloaded.to(getattr(torch, weight_info["dtype"])) return onloaded def offload( self, tensor: torch.Tensor | None, offloaded: Optional[torch.Tensor] = None ) -> torch.Tensor | None: """ Offload a tensor to disk by writing a new safetensors file :param tensor: tensor on any device :param offloaded: optional meta tensor used to look up an existing file :return: meta tensor representing the offloaded tensor """ if tensor is None: return None if tensor.device.type == "meta": assert tensor in self.index return tensor if offloaded is None: offloaded = send_tensors(tensor, device="meta") file_path = self._get_ct_file_path(self.offload_dir, offloaded) self.index[offloaded] = { "safetensors_file": file_path, "weight_name": "weight", "dtype": str(tensor.dtype).removeprefix("torch."), } assert self._is_ct_file_path(file_path), f"Attempted to write to {file_path}" save_file({"weight": tensor}, file_path) return offloaded def __delitem__(self, key: str): """ Remove the offload associated with `key`. If a new file was created to store updated tensor data, that new tensor data file is deleted. Any references to onloaded tensors held by this class are invalidated. :param key: name of tensor to invalidate """ offloaded = self.offloaded_values[key] file_path = self.index[offloaded]["safetensors_file"] if self._is_ct_file_path(file_path): os.remove(file_path) del self.index[offloaded] super().__delitem__(key) def update_offload(self, offloaded: torch.Tensor, data: torch.Tensor | None): """ Write new param data to file that already exists. :param offloaded: meta tensors representating parameter to update :param data: new data """ # get weight info from index assert offloaded in self.index, "Cannot find offload to update" weight_info = self.index[offloaded] file_path = weight_info["safetensors_file"] weight_name = weight_info["weight_name"] dtype = getattr(torch, weight_info["dtype"]) # create new file if old file was a symlink to a checkpoint file if os.path.islink(file_path): assert self._is_ct_file_path(file_path), f"Attempted to remove {file_path}" os.unlink(file_path) # save with data using original weight_name assert self._is_ct_file_path(file_path), f"Attempted to write to {file_path}" save_file({weight_name: data.reshape_as(offloaded).to(dtype=dtype)}, file_path) @classmethod def create_checkpoint_symlink( cls, offloaded: torch.Tensor, weight_info: dict, offload_dir: str | os.PathLike | None, ) -> None: assert ( is_source_process() ), "Must call on rank 0 to avoid id collisions between ranks" if offload_dir is None: raise ValueError( "Must provide an `offload_dir` to perform disk offloading " "(add `offload_folder` argument to `from_pretrained`)" ) # Resolve relative paths to absolute paths for symlink creation source_path = Path(weight_info["safetensors_file"]).resolve() file_path = cls._get_ct_file_path(offload_dir, offloaded) os.symlink(source_path, file_path) cls.index[offloaded] = { "safetensors_file": file_path, "weight_name": weight_info["weight_name"], "dtype": weight_info["dtype"], } @classmethod def _is_ct_file_path(cls, file_path: str) -> bool: """Only write and delete files that DiskCache has created""" return os.path.basename(file_path).startswith(cls._ct_file_prefix) @classmethod def _get_ct_file_path(cls, offload_dir: str, offloaded: torch.Tensor) -> str: """Create file path with a prefix marking it as modifiable""" file_name = f"{cls._ct_file_prefix}_{_get_rank()}_{id(offloaded)}.safetensors" return os.path.join(offload_dir, file_name) def _get_safe_open_device(device: "DeviceLikeType") -> str: """ `safetensors.safe_open` does not accept `torch.device` as argument, so we must convert from torch.device to a string, while considering accelerator device index resolution. :param device: torch device to convert :return: device string for `safetensors.safe_open` """ device = torch.device(device) if is_accelerator_type(device.type): # TODO: check if this case can be applied for all non-index accelerators if device.type == "mps": return f"{device.type}" if device.index is None: index = torch.accelerator.current_device_index() else: index = device.index return f"{device.type}:{index}" else: return device.type def _get_rank() -> int: """Get rank, value is zero if not distributed""" if dist.is_initialized(): return dist.get_rank() else: return 0 vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/cache/dist_cpu.py000066400000000000000000000043271521257237700320440ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch import torch.distributed as dist from compressed_tensors.distributed import get_source_rank, is_source_process from compressed_tensors.offload.cache.cpu import CPUCache from compressed_tensors.offload.cache.utils import catch_cpu_mem_error from compressed_tensors.offload.utils import send_tensors, to_empty class DistributedCPUCache(CPUCache): """ Handles offloading and onloading tensors from/to cpu memory shared across processes """ @catch_cpu_mem_error def offload(self, tensor: torch.Tensor | None) -> torch.Tensor | None: """ Synchronously create shared cpu memory for offload :param tensor: tensor on any device :return: cpu tensor whose data is located in shared memory """ if tensor is None: return None # slight runtime cost for views tensor = tensor.contiguous() if is_source_process(): # create shared memory cpu tensor tensor = super().offload(tensor).share_memory_() handle, filename, nbytes = tensor.untyped_storage()._share_filename_cpu_() broadcast_obj = [handle, filename, nbytes] else: broadcast_obj = [None, None, None] # receive shared memory file handle dist.broadcast_object_list(broadcast_obj, src=get_source_rank()) if not is_source_process(): # materialize meta tensor only if necessary if tensor.device.type == "meta": tensor = to_empty(tensor, device=self.offload_device) else: tensor = send_tensors(tensor, device=self.offload_device) # reconstruct tensor from shared memory file handle with torch.no_grad(): tensor.set_( torch.UntypedStorage._new_shared_filename_cpu(*broadcast_obj), storage_offset=tensor.storage_offset(), size=tensor.size(), stride=tensor.stride(), ) # ensure that rank 0 does not garbage collect before other ranks reconstruct dist.barrier() return tensor vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/cache/dist_device.py000066400000000000000000000026751521257237700325200ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch import torch.distributed as dist from compressed_tensors.distributed import ( as_broadcastable, get_source_rank, is_source_process, ) from compressed_tensors.offload.cache.device import DeviceCache from compressed_tensors.offload.utils import send_tensors, to_empty class DistributedDeviceCache(DeviceCache): """ Handles offloading and onloading tensors from/to device memory. Onloading tensors is typically a no-op (except when onload device has been modified). The device offload is not shared between ranks. When dispatching with this cache, the model is replicated across devices. """ def offload(self, tensor: torch.Tensor | None) -> torch.Tensor | None: """ Move a tensor to device, then broadcast data to all other ranks :param value: tensor on any device :return: tensor on device """ if tensor is None: return None if is_source_process(): tensor = super().offload(tensor) # materialize meta tensor only if necessary elif tensor.device.type == "meta": tensor = to_empty(tensor, device=self.offload_device) else: tensor = send_tensors(tensor, device=self.offload_device) dist.broadcast(as_broadcastable(tensor), src=get_source_rank()) return tensor vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/cache/dist_disk.py000066400000000000000000000043251521257237700322050ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch import torch.distributed as dist from compressed_tensors.distributed import get_source_rank, is_source_process from compressed_tensors.offload.cache.disk import DiskCache from compressed_tensors.offload.utils import send_tensors class DistributedDiskCache(DiskCache): """ Handles offloading and onloading tensors from/to disk. For more information, see `compressed_tensors.offload.cache.disk_cache::DiskCache`. """ def offload(self, tensor: torch.Tensor | None) -> torch.Tensor | None: """ Synchronously write tensor data to disk :param tensor: tensor on any device :return: meta tensor representing disk offloaded parameter """ if tensor is None: return None if is_source_process(): # write to disk offloaded = super().offload(tensor) broadcast_obj = [ self.index[offloaded]["safetensors_file"], self.index[offloaded]["weight_name"], self.index[offloaded]["dtype"], ] else: offloaded = send_tensors(tensor, device="meta") broadcast_obj = [None, None, None] dist.broadcast_object_list(broadcast_obj, src=get_source_rank()) if not is_source_process(): self.index[offloaded] = { "safetensors_file": broadcast_obj[0], "weight_name": broadcast_obj[1], "dtype": broadcast_obj[2], } # wait for write to finish dist.barrier() return offloaded def __delitem__(self, key: str): """ Remove the offload associated with `key`. If a new file was created to store updated tensor data, that new tensor data file is deleted. Any references to onloaded tensors held by this class are invalidated. :param key: name of tensor to invalidate """ if is_source_process(): super().__delitem__(key) else: offloaded = self.offloaded_values[key] del self.index[offloaded] super(DiskCache, self).__delitem__(key) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/cache/utils.py000066400000000000000000000027301521257237700313660ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import errno from functools import wraps from loguru import logger _CPU_MEMORY_KEYWORDS = ( "defaultcpuallocator", "can't allocate memory", "cannot allocate memory", "failed to allocate", "out of memory", "mmap", "shm_open", ) _CPU_MEMORY_REMEDIATION = ( "CPU offloading ran out of host RAM or mmap descriptors. " "Switch to disk offloading (`offload_device='disk'`) or " "increase the OS mmap limit." ) def _is_cpu_memory_error(exc: BaseException) -> bool: errno_value = getattr(exc, "errno", None) if errno_value == errno.ENOMEM: return True message = str(exc).lower() return any(kw in message for kw in _CPU_MEMORY_KEYWORDS) def catch_cpu_mem_error(func): """ Decorator to catch CPU memory errors and log a remediation warning. Prevents duplicate logs if nested functions also use this decorator. """ @wraps(func) def wrapper(*args, **kwargs): try: return func(*args, **kwargs) except (RuntimeError, OSError) as exc: # Prevent duplicate logs when DistributedCPUCache calls super().offload() if getattr(exc, "_cpu_mem_logged", False): raise if _is_cpu_memory_error(exc): logger.warning(_CPU_MEMORY_REMEDIATION) exc._cpu_mem_logged = True raise return wrapper vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/convert/000077500000000000000000000000001521257237700302675ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/convert/__init__.py000066400000000000000000000003221521257237700323750ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .from_accelerate import from_accelerate from .to_accelerate import to_accelerate vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/convert/from_accelerate.py000066400000000000000000000222421521257237700337560ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os from typing import TYPE_CHECKING, Literal import torch import torch.distributed as dist from compressed_tensors.distributed import ( get_source_rank, is_distributed, is_source_process, ) from compressed_tensors.offload.cache import DiskCache, OffloadCache from compressed_tensors.offload.convert.helpers import ( DEFAULT_OFFLOAD_DEVICE, get_tensors, norm_device, ) from compressed_tensors.offload.dispatch import dispatch_with_map from compressed_tensors.offload.utils import to_tensor from loguru import logger if TYPE_CHECKING: from accelerate.utils import OffloadedWeightsLoader from compressed_tensors.offload.dispatch import DeviceMap __all__ = ["from_accelerate", "remove_accelerate", "remove_accelerate_from_module"] def from_accelerate(model: torch.nn.Module) -> tuple["DeviceMap", str | None]: """ Convert a model from accelerate offloading to compressed-tensors offloading. Often called by `load_offloaded_model` to load offloaded models across ranks. Rank 0 is always expected to provide an accelerate-offloaded model. Other ranks (if they exist) may provide a model on any device, with or without accelerate offloading. - If called after `load_offloaded_model()`, other ranks will provide a meta model with no accelerate offloading - If called after `to_accelerate`, other ranks will provide an accelerate-offloaded model shared cpu tensors/file paths. :param model: accelerate-offloaded model if source process, no constraint otherwise """ # Early exit if the model has already been converted to compressed-tensors # offloading. This guards against redundant conversion when `from_pretrained` # is patched by multiple wrappers in the same MRO chain (e.g. # `AutoModelForCausalLM` delegating to `Qwen3ForCausalLM`). if _is_offloaded(model): return {}, None device_map, offload_dir = remove_accelerate(model) broadcast_obj = [device_map, offload_dir] if is_distributed(): dist.broadcast_object_list(broadcast_obj, src=get_source_rank()) dispatch_with_map(model, *broadcast_obj) return tuple(broadcast_obj) def remove_accelerate(model: torch.nn.Module) -> tuple["DeviceMap", str | None]: """ Remove accelerate offloading from a model, if applicable :param model: model containing accelerate offloaded modules :returns: `(device_map, offload_dir)` """ offload_dir = None device_map = {} for name, module in model.named_modules(remove_duplicate=False): onload_dev, offload_dev, _offload_dir = remove_accelerate_from_module(module) if _offload_dir is not None: if offload_dir is not None and _offload_dir != offload_dir: raise ValueError( "Expected model to only have one `offload_dir`, " f"instead got {offload_dir} and {_offload_dir}" ) offload_dir = _offload_dir device_map[name] = (onload_dev, offload_dev) if hasattr(model, "hf_device_map"): delattr(model, "hf_device_map") return device_map, offload_dir def remove_accelerate_from_module( module: torch.nn.Module, ) -> tuple[torch.device | None, torch.device | Literal["disk"] | None, str | None]: """ Remove accelerate offloading from a module, if present. Absolutely no device movement occurs, and parameters/buffers pointers from state dicts coming from `to_accelerate` remain unchanged so as to avoid memory duplication :param module: module to remove offloading from :returns: `(onload_device, offload_device, disk_offload_dir)` """ try: from accelerate.hooks import AlignDevicesHook, remove_hook_from_module from accelerate.utils import OffloadedWeightsLoader, PrefixedDataset except ImportError: device = _infer_module_device(module) return device, device, None hook = getattr(module, "_hf_hook", None) direct_tensors = _direct_tensors(module) # No AlignDevicesHook: treat as "not offloaded" if not isinstance(hook, AlignDevicesHook): device = _infer_device_from_tensors(direct_tensors) return device, device, None # Hook exists but no active offload (or nothing to consider) if not hook.offload or not direct_tensors: hook.offload = False remove_hook_from_module(module, recurse=False) device = _infer_device_from_tensors(direct_tensors) return device, device, None # Unwrap PrefixedDataset chain so we can look up real tensor keys prefix, dataset = _unwrap_prefixed_dataset(hook.weights_map, PrefixedDataset) assert isinstance(dataset, (OffloadedWeightsLoader, dict)) offload_dev: str | None = None for local_name, tensor in direct_tensors.items(): full_name = prefix + local_name # Device/CPU offload if isinstance(dataset, dict) and local_name in dataset: offload = dataset[local_name] offload_dev = _set_or_validate_offload(offload_dev, offload.device.type) # Device/CPU offload elif ( isinstance(dataset, OffloadedWeightsLoader) and full_name in dataset.state_dict ): offload = dataset.state_dict[full_name] offload_dev = _set_or_validate_offload(offload_dev, offload.device.type) # Disk offload elif isinstance(dataset, OffloadedWeightsLoader) and full_name in dataset.index: offload = tensor offload_dev = _set_or_validate_offload(offload_dev, "disk") assert offload.device.type == "meta" assert isinstance(offload, (torch.nn.Parameter, torch.nn.Buffer)) # Copy accelerate's disk index into DiskCache for our later use if is_source_process(): _save_ct_index_entry(dataset, full_name, tensor) # Not offloaded, likely a buffer else: offload = tensor # Replace meta tensor with offloaded value (no ptr rematerialization occurs) # In the disk case, the tensor remains as the meta tensor if not isinstance(offload, (torch.nn.Parameter, torch.nn.Buffer)): to_tensor(offload, tensor) setattr(module, local_name, offload) # Prevent onloading disk tensors while removing hook hook.offload = False remove_hook_from_module(module, recurse=False) return ( norm_device(hook.execution_device), norm_device(offload_dev if offload_dev is not None else DEFAULT_OFFLOAD_DEVICE), dataset.save_folder if isinstance(dataset, OffloadedWeightsLoader) else None, ) def _save_ct_index_entry( dataset: "OffloadedWeightsLoader", name: str, offloaded: torch.Tensor, ): # already indexed from a previous round-trip (e.g. to_accelerate -> from_accelerate) if offloaded in DiskCache.index: return entry: dict = dataset.index[name] if "safetensors_file" in entry: # typical case: model is loaded from safetensors file # create a symlink that points to the model safetensor file # if the value is ever updated, the symlink is broken and a real file # is written to that location DiskCache.create_checkpoint_symlink(offloaded, entry, dataset.save_folder) else: # unfortunately, ct's implementation does not support loading non-safetensors # we must onload and save as safetensors. This typically only occurs in testing onloaded = dataset[name] DiskCache("cpu", offload_dir=dataset.save_folder).offload( onloaded, offloaded=offloaded ) logger.warning( "Attempting to disk offload a model which was not saved with safetensors. " "compressed-tensors only supports disk onload from safetensors files, so " "weights must be onloaded and re-saved as safetensors files.", log_once=True, ) # remove original weight_file original_weight_file = os.path.join(dataset.save_folder, f"{name}.dat") if os.path.exists(original_weight_file): os.remove(original_weight_file) def _direct_tensors(module: torch.nn.Module) -> dict[str, torch.Tensor]: return {name: t for name, t in get_tensors(module) if t is not None} def _infer_device_from_tensors(tensors: dict[str, torch.Tensor]) -> torch.device | None: t = next(iter(tensors.values()), None) return norm_device(t.device if t is not None else None) def _infer_module_device(module: torch.nn.Module) -> torch.device | None: return _infer_device_from_tensors(_direct_tensors(module)) def _unwrap_prefixed_dataset(weights_map, PrefixedDatasetType): prefix = "" dataset = weights_map while isinstance(dataset, PrefixedDatasetType): prefix += dataset.prefix dataset = dataset.dataset return prefix, dataset def _set_or_validate_offload(current: str | None, new: str) -> str: if current not in (None, new): raise ValueError("Expected all accelerate tensors to share offload") return new def _is_offloaded(model: torch.nn.Module) -> bool: return any(isinstance(m._parameters, OffloadCache) for m in model.modules()) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/convert/helpers.py000066400000000000000000000030651521257237700323070ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from itertools import chain from typing import Iterable, Literal import torch import torch.distributed as dist from compressed_tensors.distributed import is_distributed from compressed_tensors.utils import is_accelerator_type __all__ = [ "get_tensors", "norm_device", "DEFAULT_OFFLOAD_DEVICE", ] DEFAULT_OFFLOAD_DEVICE = torch.device("cpu") def norm_device( device: str | torch.device | None, ) -> torch.device | Literal["disk"] | None: """ Standardize the representation of devices for the purposes of consistency - when running with distributed, represent rank device as the accelerator type - when not running with distributed, bare accelerator type (e.g. "cuda") is resolved to index 0 """ if device in ("disk", None): return device device = torch.device(device) # (dist) "cuda:R" / "xpu:R" -> bare type if is_distributed() and device.index == dist.get_rank(): device = torch.device(type=device.type, index=None) # (non-dist) bare accelerator type -> index 0 if ( not is_distributed() and is_accelerator_type(device.type) and device.index is None ): device = torch.device(type=device.type, index=0) return device def get_tensors( module: torch.nn.Module, recurse: bool = False ) -> Iterable[tuple[str, torch.Tensor | None]]: return chain( module.named_parameters(recurse=recurse), module.named_buffers(recurse=recurse) ) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/convert/to_accelerate.py000066400000000000000000000110411521257237700334300ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import defaultdict from typing import Optional import torch from compressed_tensors.offload.cache import DiskCache, OffloadCache from compressed_tensors.offload.convert.helpers import ( DEFAULT_OFFLOAD_DEVICE, get_tensors, norm_device, ) from compressed_tensors.offload.module import remove_module_offload from compressed_tensors.utils import patch_attr from loguru import logger __all__ = ["to_accelerate", "to_accelerate_module"] def to_accelerate(model: torch.nn.Module) -> dict[str, str]: """ Convert a model from `compressed_tensors` offloading to `accelerate` offloading. This is is often called before `PreTrainedModel.save_pretrained`, as without this conversion, `save_pretrained` will use excessive memory and device movement. :param model: model dispatched with `compressed_tensors` offloading :return: accelerate-style device map """ hf_device_map = {} hf_disk_index = _to_accelerate_disk_index(model, DiskCache.index) for name, module in model.named_modules(): offload_device_str = to_accelerate_module(module, name, hf_disk_index) hf_device_map[name] = offload_device_str setattr(model, "hf_device_map", hf_device_map) return hf_device_map def to_accelerate_module( module: torch.nn.Module, name: Optional[str] = None, hf_disk_index: Optional[dict[str, dict[str, str]]] = None, ) -> str: """ Convert a module from `compressed_tensors` offloading to `accelerate` offloading :param module: module to convert to accelerate offloading :param name: name of module in model :param hf_disk_index: accelerate-style disk index to attach to weight loaders :return: str of offloaded device. Defaults to cpu if module does not have parameters """ has_accelerate = True try: from accelerate.hooks import AlignDevicesHook, add_hook_to_module from accelerate.utils import OffloadedWeightsLoader, PrefixedDataset except ImportError: logger.warning( "Cannot convert module without `accelerate` installed. This may result " "in high memory usage during saving and tied tensors being saved twice", log_once=True, ) has_accelerate = False offload_device = DEFAULT_OFFLOAD_DEVICE if has_accelerate and isinstance(module._parameters, OffloadCache): cache = module._parameters offload_device = norm_device(cache.offload_device) remove_module_offload(module, onload_tensors=False) # create weights map if isinstance(cache, DiskCache): if name is None or hf_disk_index is None: raise ValueError( "Must provide `name` and `hf_disk_index` " "to convert disk offloaded module" ) weights_map = PrefixedDataset( prefix=f"{name}.", dataset=OffloadedWeightsLoader( index=hf_disk_index, save_folder=cache.offload_dir, ), ) else: weights_map = dict(get_tensors(module, recurse=False)) # create hook hook = AlignDevicesHook( execution_device=cache.onload_device, offload=True, io_same_device=True, weights_map=weights_map, offload_buffers=True, place_submodules=False, ) # add hook (skip onloading) with patch_attr(AlignDevicesHook, "init_hook", lambda self, module: module): add_hook_to_module(module, hook) # skipping init hook => need to populate `original_devices` hook.original_devices = { name: offload_device if offload_device != "disk" else torch.device("cpu") for name, _ in get_tensors(module, recurse=False) } return str(offload_device) def _to_accelerate_disk_index( model: torch.nn.Module, index: dict[torch.Tensor, dict[str, str]] ) -> dict[str, dict[str, str]]: from compressed_tensors.offload import disable_onloading # circular dependency with disable_onloading(): offloaded_to_key = _invert_dict(model.state_dict(keep_vars=True)) return { key: weight_info for offloaded, weight_info in index.items() for key in offloaded_to_key[offloaded] } def _invert_dict(d: dict) -> dict: inverted = defaultdict(list) for key, value in d.items(): inverted[value].append(key) return inverted vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/dispatch.py000066400000000000000000000246411521257237700307670ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Container from copy import deepcopy from functools import partial from typing import Any, Optional, TypeVar import torch import torch.distributed as dist from compressed_tensors.offload.cache import OffloadCache from compressed_tensors.offload.module import offload_module, remove_module_offload from compressed_tensors.offload.utils import ( get_module_device, get_module_sizes, module_size, ) from compressed_tensors.utils import getattr_chain from compressed_tensors.utils.binary_search import SearchFailureError, max_binary_search from compressed_tensors.utils.helpers import deprecated from loguru import logger from tqdm import tqdm from transformers import PreTrainedModel __all__ = [ "set_onload_device", "offload_model", "dispatch_with_map", "get_device_map", "dispatch_model", "remove_dispatch", "get_device_memory", "DeviceMap", ] ModelType = TypeVar("ModelType", bound=torch.nn.Module) DeviceMap = dict[str, tuple[torch.device | None, torch.device | str | None]] def set_onload_device( model: ModelType, onload_device: torch.device | str, ) -> ModelType: """ Modify the dispatch of a model to onload to the provided `onload_device`. Existing offloaded tensors will not be modified. If a module is not already offloaded, it will be offloaded to its current device. :param model: model to dispatch :param onload_device: device to move weights to during forward pass :return: dispatched model """ for module in model.modules(): if isinstance(module._parameters, OffloadCache): module._parameters.onload_device = onload_device module._buffers.onload_device = onload_device else: offload_device = get_module_device(module, torch.device("cpu")) offload_module(module, onload_device, offload_device) return model @deprecated("set_onload_device") def offload_model( model: ModelType, onload_device: torch.device | str, offload_device: Any = None, ) -> ModelType: """ .. deprecated:: Use :func:`set_onload_device` instead. """ return set_onload_device(model, onload_device) def dispatch_with_map( model: torch.nn.Module, device_map: DeviceMap, offload_dir: Optional[str] = None, show_progress: bool = True, ): """ Dispatch a model according to the provided device map :param model: model to dispatch :param device_map: device map specifying the onload and offload of each module :param offload_dir: optional directory for disk offloading :param show_progress: show tqdm progress """ for name, (onload_device, offload_device) in tqdm( list(device_map.items()), desc="Dispatching model", disable=(not show_progress) ): module = model.get_submodule(name) if offload_device == "disk": offload_module( module, onload_device, offload_device, offload_dir=offload_dir ) elif offload_device is not None: offload_module(module, onload_device, offload_device) def get_device_map( model: torch.nn.Module, default_device: torch.device = torch.device("cpu") ) -> DeviceMap: """ Get the device map of a CT-offloaded model :param: model: model to get device map of :param default_device: the default onload/offload device when module has no parameters :return: device map specifying the onload and offload device of all modules """ from compressed_tensors.offload import get_execution_device, get_offloaded_device return { name: ( get_execution_device(module, default_device), get_offloaded_device(module, default_device), ) for name, module in model.named_modules(remove_duplicate=False) } def dispatch_model( model: ModelType, device_memory: dict[torch.device, int] | None = None, extra_memory: int | None = None, no_split_modules: Container[str] | None = None, ) -> ModelType: """ Dispatch a model for autoregressive generation. This means that modules are dispatched evenly across available devices and kept onloaded if possible. If onloading the entire model is not possible, some modules may be offloaded. Any existing offloads will be removed. Disclaimers: * Optimal runtime assumes that modules are called in order of `model.modules()` :param model: model to dispatch :param device_memory: optional dictionary mapping torch device to available memory. If none is provided, all available devices will be used :param extra_memory: the amount of memory to be reserved for activations :param no_split_modules: names of module classes which should not be split across multiple devices :return: dispatched model """ # infer no_split_modules if no_split_modules is None: no_split_modules = getattr(model, "_no_split_modules", tuple()) # collect devices if device_memory is None: device_memory: dict[torch.device, int] = get_device_memory() if len(device_memory) <= 0: raise MemoryError("Did not find any devices to dispatch model to") # collect module sizes sizes = get_module_sizes(model, no_split_modules) if len(sizes) <= 0: raise ValueError("Model does not have any modules") # estimate memory requirement if extra_memory is None: # fragmentation, kv cache, embeddings, ect. extra_memory = max(module_size(model) * 0.05, 1e9) # activations if isinstance(model, PreTrainedModel): extra_memory += ( 1 # batch_size * 2048 # seq_len * getattr_chain(model, "config.intermediate_size", 256) * getattr(model, "dtype", torch.bfloat16).itemsize ) # search for the best dispatch which maximizes extra memory across devices try: max_extra_memory = min(device_memory.values()) extra_memory, (dispatch, _) = max_binary_search( fn=partial(_get_greedy_dispatch, sizes, device_memory), cond=(lambda result: len(result[0]) == len(sizes)), start=extra_memory, end=max_extra_memory, ) # fallback: create a cpu dispatch except SearchFailureError: dispatch, device_memory = _get_greedy_dispatch( sizes, device_memory, extra_memory ) assert len(dispatch) < len(sizes) last_device = dispatch[-1][1] if len(dispatch) else list(device_memory)[0] sizes_dict = {module: size for module, size in sizes} largest_offloaded_module = max(size for _, size in sizes[len(dispatch) :]) # pop off modules until all offloaded modules can fit in last device while largest_offloaded_module + extra_memory > device_memory[last_device]: if len(dispatch) <= 0: raise ValueError( f"Cannot fit no_split module of size {largest_offloaded_module} " f"bytes into any device: {device_memory}" ) module, last_device, _ = dispatch.pop(-1) device_memory[last_device] += sizes_dict[module] largest_offloaded_module = max(largest_offloaded_module, sizes_dict[module]) # fill dispatch back with cpu offloading for module, _ in list(sizes[len(dispatch) :]): dispatch.append((module, last_device, "cpu")) logger.warning("Forced to offload modules due to insufficient gpu resources") # dispatch assert len(dispatch) == len(sizes) dispatch_dict = { submodule: (onload, offload) for module, onload, offload in dispatch for submodule in module.modules() } for module in model.modules(): remove_module_offload(module, onload_tensors=True) if module in dispatch_dict: onload, offload = dispatch_dict[module] offload_module(module, onload, offload) logger.debug(f"Dispatched model with {extra_memory} bytes of extra memory") return model def get_device_memory() -> dict[torch.device, int]: """ Get the total memory of all available accelerator devices. Returns accelerator device memory when available, otherwise falls back to CPU with system RAM. :return: mapping from torch device to total memory """ if not torch.accelerator.is_available(): import os total_ram = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") return {torch.device("cpu"): total_ram} accel_type = torch.accelerator.current_accelerator().type if dist.is_available() and dist.is_initialized(): logger.info("Detected distributed context. Dispatching to local rank gpu") device_memory = torch.accelerator.get_memory_info( torch.accelerator.current_device_index() )[1] return {torch.device(accel_type): device_memory} return { torch.device(accel_type, idx): torch.accelerator.get_memory_info(idx)[1] for idx in range(torch.accelerator.device_count()) } def remove_dispatch( module: torch.nn.Module, onload_tensors: bool = False ) -> torch.nn.Module: """ Remove any existing dispatches from module :param onload_tensors: Whether to move tensors to the onloaded device, or keep them on the offload device. Defaults to False. :return: module with offloading functionality removed """ for submodule in module.modules(): remove_module_offload(submodule, onload_tensors) return module def _get_greedy_dispatch( sizes: list[tuple[torch.nn.Module, int]], device_memory: dict[torch.device, int], extra_memory: int = 0, ) -> tuple[ list[tuple[torch.nn.Module, torch.device, torch.device]], dict[torch.device, int] ]: dispatch = list() memory_remaining = deepcopy(device_memory) device_index = 0 devices = list(memory_remaining.keys()) if len(devices) <= 0: raise ValueError() for module, size in sizes: while True: if device_index >= len(devices): return dispatch, memory_remaining device = devices[device_index] if size > memory_remaining[device] - extra_memory: device_index += 1 continue dispatch.append((module, device, device)) memory_remaining[device] -= size break return dispatch, memory_remaining vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/dist_utils.py000066400000000000000000000007621521257237700313510ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from compressed_tensors.distributed import ( as_broadcastable, init_dist, is_distributed, is_source_process, ) from compressed_tensors.utils.helpers import deprecated __all__ = ["is_distributed", "is_rank0", "init_dist", "as_broadcastable"] @deprecated("compressed_tensors.distributed.utils.py::is_source_process") def is_rank0() -> bool: return is_source_process() vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/load.py000066400000000000000000000105401521257237700301000ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib import os import shutil from functools import wraps import psutil import torch from compressed_tensors.distributed import is_distributed, is_source_process from compressed_tensors.offload.convert import from_accelerate from compressed_tensors.utils import patch_attr from loguru import logger from transformers import AutoModelForCausalLM, PreTrainedModel __all__ = ["load_offloaded_model"] @contextlib.contextmanager def load_offloaded_model( model_class: type[PreTrainedModel] = AutoModelForCausalLM, extra_cpu_mem: int = 5e9 ): """ Context manager used to load a transformers model with offloading implemented by compressed-tensors. The model is first loaded with accelerate's offloading, then convereted into offloading implemented by compressed-tensors. If a distributed environment has been initialized, then rank 0 loads the weights while other ranks load on the meta device, then the offload is shared across ranks during conversion. In addition to the standard `device_map` options, this context also supports `device_map="auto_offload"`, which means that the model will load as many parameters can fit onto the cpu, and any extra parameters will be loaded on disk. :param model_class: model class to patch :param extra_cpu_mem: extra cpu memory to reserve for any operations not related to model loading (bytes). Defaults to 5Gb. """ original_from_pretrained = model_class.from_pretrained patched_fn_called = False @classmethod @wraps(original_from_pretrained) def patched(cls, *args, **kwargs): nonlocal patched_fn_called patched_fn_called = True kwargs.setdefault("device_map", None) # Rank 0 does loading, other ranks init on meta device if not is_source_process(): kwargs["device_map"] = "meta" # Workaround: transformers v5 tie_weights() calls torch.equal() on # meta tensors which is unsupported. Since rank 0 broadcasts the real # weights, we can safely skip tying on non-rank workers. kwargs.setdefault("tie_word_embeddings", False) # Intercept `auto_offload`: same as "auto", but only cpu/disk are visible elif kwargs["device_map"] == "auto_offload": kwargs["device_map"] = "auto" if "max_memory" not in kwargs: kwargs["max_memory"] = _get_cpu_memory(extra_cpu_mem) # Unless the user specifies, use our memory estimates, which take into # account distributed setups and extra cpu reserved memory elif "max_memory" not in kwargs: kwargs["max_memory"] = _get_device_memory() | _get_cpu_memory(extra_cpu_mem) model = original_from_pretrained(*args, **kwargs) from_accelerate(model) # rank 0 shares weights with ranks via offload/broadcast return model with patch_attr(model_class, "from_pretrained", patched): try: yield finally: if not patched_fn_called: logger.warning( f"`{model_class.__name__}.from_pretrained` was never called. If " "you are loading with a model class other than " f"{model_class.__name__}, please pass as argument to " "`load_offloaded_model`" ) def _get_device_memory() -> dict[int, int]: if is_distributed(): index = torch.accelerator.current_device_index() return {index: torch.accelerator.get_memory_info(index)[1]} else: return { index: torch.accelerator.get_memory_info(index)[1] for index in range(torch.accelerator.device_count()) } def _get_cpu_memory(extra_cpu_mem: int) -> dict[str, int]: if is_distributed(): return {"cpu": _get_shared_memory() - extra_cpu_mem} else: return {"cpu": psutil.virtual_memory().available - extra_cpu_mem} def _get_shared_memory() -> int: linux_shm_path = "/dev/shm" if os.path.exists(linux_shm_path): total, _used, _free = shutil.disk_usage(linux_shm_path) return total else: logger.warning( "Could not find shared memory at `/dev/shm`. Please add platform suppport" ) return psutil.virtual_memory().available vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/module.py000066400000000000000000000075401521257237700304540ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib from functools import wraps import torch from compressed_tensors.offload.cache.base import OffloadCache from compressed_tensors.offload.utils import send_tensors def offload_module( module: torch.nn.Module, onload_device: torch.device | str, offload_device: torch.device | str, **kwargs, ): """ Offload a module. Any existing parameters or buffers will be offloaded to the offload device specified by the `cache`. Accessing module parameters or buffers will cause them to be onloaded to the `onload_device`. Calling `forward` will result in input tensors being moved to the `onload_device`, and any onloaded parameters or buffers will remain onloaded for the duration of the forward call if `no_split` is set to `True`. :param module: module to offload :param onload_device: device used to onload parameters and buffers :param offload_device: device used to offload parameters and buffers :param \\**kwargs: keyword arguments for cache constructor """ if isinstance(module._parameters, OffloadCache): raise ValueError( "Attempted to offload a module twice. " "Please call `remove_module_offload` first." ) cache_cls = OffloadCache.cls_from_device(offload_device) module._parameters = cache_cls.from_mapping( module._parameters, onload_device, offload_device, **kwargs ) module._buffers = cache_cls.from_mapping( module._buffers, onload_device, offload_device, **kwargs ) original_forward_func = module.forward.__func__ module._original_forward_func = original_forward_func @wraps(original_forward_func) def forward(self, *args, **kwargs): if not OffloadCache.onloading_disabled and isinstance( module._parameters, OffloadCache ): onload_device = module._parameters.onload_device args = send_tensors(args, device=onload_device) kwargs = send_tensors(kwargs, device=onload_device) return self._original_forward_func(self, *args, **kwargs) module.forward = forward.__get__(module) return module def remove_module_offload(module: torch.nn.Module, onload_tensors: bool = False): """ Remove any offloading applied to the module :param onload_tensors: Whether to move tensors to the onloaded device, or keep them on the offload device. Defaults to False. """ if isinstance(module._parameters, OffloadCache): assert isinstance(module._buffers, OffloadCache) if onload_tensors: module._parameters = { name: module._parameters.onload(param) for name, param in module._parameters.offloaded_values.items() } module._buffers = { name: module._buffers.onload(param) for name, param in module._buffers.offloaded_values.items() } else: module._parameters = module._parameters.offloaded_values module._buffers = module._buffers.offloaded_values module.forward = module._original_forward_func.__get__(module) del module._original_forward_func @contextlib.contextmanager def unwrap_offload_forward(module: torch.nn.Module): """ Upon entering, module forward function is unwrapped. Upon exiting the offloading wrapper is added again. Any modifications made to the forward function while within the context will be reflected upon exiting. """ if hasattr(module, "_original_forward_func"): offload_forward = module.forward module.forward = module._original_forward_func.__get__(module) yield module._original_forward_func = module.forward.__func__ module.forward = offload_forward else: yield vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/offload/utils.py000066400000000000000000000165041521257237700303270ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib from collections.abc import Container from dataclasses import fields, is_dataclass from itertools import chain from typing import TypeVar import torch from compressed_tensors.utils.helpers import patch_attr from loguru import logger __all__ = [ "send_tensors", "get_module_device", "move_module_tensor", "module_size", "to_empty", "to_tensor", "to_meta", "as_single_threaded", ] T = TypeVar("T") TensorCls = TypeVar("TensorCls", bound=torch.Tensor) def send_tensors(value: T, *args, **kwargs) -> T: """ Recursively identify and move tensors using `torch.Tensor.to` :param value: value containing tensors to move :param args: arguments to `to` :param kwargs: keyword arguments to `to` :return: value with moved tensors """ match value: case torch.Tensor(): with torch.no_grad(): tensor = value.to(*args, **kwargs) # special case: avoid changing param pointer when possible if tensor.data_ptr() == value.data_ptr(): return value tensor.__class__ = value.__class__ tensor.__dict__ = value.__dict__.copy() return tensor case list(): return [send_tensors(v, *args, **kwargs) for v in value] case tuple(): return tuple(send_tensors(v, *args, **kwargs) for v in value) case dict(): return {k: send_tensors(v, *args, **kwargs) for k, v in value.items()} case _ if is_dataclass(value): return type(value)( **{ f.name: send_tensors(getattr(value, f.name), *args, **kwargs) for f in fields(value) } ) case _: return value def get_module_device( module: torch.nn.Module, default: torch.device | None = None ) -> torch.device: """ Infer the device of a module using the first parameter or buffer registered to the module :param module: module to check :param default: default device if module does not have tensors or buffers :return: device of module """ tensor = next(module.parameters(), next(module.buffers(), None)) if tensor is not None: return tensor.device elif default is not None: return default else: logger.warning( f"Unable to get execution device of {module}, falling back to CPU", log_once=True, ) return torch.device("cpu") def move_module_tensor( module: torch.nn.Module, name: str, device: int | str | torch.device, ): """ Move a module's tensor to a new device :param module: module containing tensors to move :param name: name of tensor to move :param device: new devices """ if name in module._parameters: module._parameters[name] = send_tensors(module._parameters[name], device=device) elif name in module._buffers: module._buffers[name] = send_tensors(module._buffers[name], device=device) def get_module_sizes( model: torch.nn.Module, no_split_modules: Container[str] = tuple() ) -> list[tuple[torch.nn.Module, int]]: """ Returns a list of modules and their sizes. Only non-splittable modules are returned. Non-splittable modules are modules specified by `no_split_modules` or modules with direct parameters. :param model: model to get sizes from :param no_split_modules: module class names which cannot be split :return: list of modules and their sizes """ module_sizes = [] def dfs(module: torch.nn.Module): # modules with direct parameters cannot be split # otherwise, submodules could return a device that is different than params direct_size = module_size(module, recurse=False) no_split = module.__class__.__name__ in no_split_modules or direct_size > 0 total_size = module_size(module, recurse=no_split) if total_size > 0: module_sizes.append((module, total_size)) if not no_split: for child in module.children(): dfs(child) dfs(model) return module_sizes def module_size(module: torch.nn.Module, recurse: bool = True) -> int: """ Get the size of the module's parameters and buffers in bytes :param module: module to check :param recurse: whether calculate recursive size, or only direct tensors :return: total size of module parameters and buffers """ from compressed_tensors.offload import disable_onloading with disable_onloading(): tensors = chain( module.parameters(recurse=recurse), module.buffers(recurse=recurse) ) return sum((tensor.nbytes for tensor in tensors), 0) def to_empty(tensor: TensorCls, **kwargs) -> TensorCls: """ Create an empty tensor like the given tensor, with the same tensor subclass and dict values :param tensor: tensor to create empty like :return: empty tensor """ empty = torch.empty_like(tensor.data, **kwargs) empty.__class__ = tensor.__class__ empty.__dict__ = tensor.__dict__.copy() return empty def to_tensor(dst: torch.Tensor, src: TensorCls) -> TensorCls: """ Copy the subclass, dict, and attributes into `dst` from `src` :param dst: tensor to copy attributes into :param src: tensor to copy attributes from :return: dst tensor with copied attributes """ dst.__class__ = src.__class__ dst.__dict__ = src.__dict__.copy() dst.requires_grad = src.requires_grad return dst def to_meta(module: torch.nn.Module) -> None: """Move all module parameters and buffers to meta device. This removes pointers to offloaded tensors held by non-processing ranks, allowing the processing rank to compress without increasing peak memory. :param module: module whose tensors should be moved to meta device """ from compressed_tensors.offload import disable_onloading from compressed_tensors.utils.module import ( get_direct_state_dict, replace_direct_state_dict, ) with disable_onloading(): state_dict = get_direct_state_dict(module) meta_state_dict = { name: send_tensors(tensor, device="meta") for name, tensor in state_dict.items() } replace_direct_state_dict(module, meta_state_dict) @contextlib.contextmanager def as_single_threaded(): """ Context manager to temporarily use single-threaded offload methods. This context manager patches distributed cache classes to use their non-distributed counterparts' offload methods. This is useful when operations need to be performed without distributed coordination. Example: >>> with as_single_threaded(): ... # Operations here use single-threaded offload ... cache.offload(data) """ from compressed_tensors.offload.cache import ( CPUCache, DeviceCache, DiskCache, DistributedCPUCache, DistributedDeviceCache, DistributedDiskCache, ) with ( patch_attr(DistributedDeviceCache, "offload", DeviceCache.offload), patch_attr(DistributedCPUCache, "offload", CPUCache.offload), patch_attr(DistributedDiskCache, "offload", DiskCache.offload), ): yield vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/000077500000000000000000000000001521257237700277235ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/__init__.py000066400000000000000000000004301521257237700320310ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa # isort: skip_file from .quant_args import * from .quant_config import * from .quant_metadata import * from .quant_scheme import * from .lifecycle import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/lifecycle/000077500000000000000000000000001521257237700316625ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/lifecycle/__init__.py000066400000000000000000000004061521257237700337730ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa # isort: skip_file from .forward import * from .initialize import * from .compressed import * from .apply import * from .helpers import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/lifecycle/apply.py000066400000000000000000000206061521257237700333650ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import OrderedDict from copy import deepcopy import torch from compressed_tensors.modeling import ( initialize_hooked_attention, initialize_hooked_kv_cache, ) from compressed_tensors.offload import update_offload_parameter from compressed_tensors.quantization.lifecycle.initialize import ( initialize_module_for_quantization, is_attention_module, ) from compressed_tensors.quantization.quant_args import QuantizationArgs from compressed_tensors.quantization.quant_config import ( QuantizationConfig, QuantizationStatus, ) from compressed_tensors.quantization.quant_scheme import QuantizationScheme from compressed_tensors.quantization.utils import KV_CACHE_TARGETS from compressed_tensors.utils.match import ( is_narrow_match, match_named_modules, match_targets, ) from compressed_tensors.utils.safetensors_load import get_safetensors_folder from loguru import logger from safetensors import safe_open from torch.nn import Module __all__ = [ "load_pretrained_quantization_parameters", "apply_quantization_config", ] from compressed_tensors.quantization.utils.helpers import is_module_quantized from compressed_tensors.utils.safetensors_load import ( get_quantization_parameter_to_path_mapping, ) def load_pretrained_quantization_parameters( model: Module, model_name_or_path: str | None = None, load_weight_qparams: bool = False, ): """ Loads the quantization parameters (scale and zero point) from model_name_or_path to a model that has already been initialized with a quantization config. NOTE: Will always load inputs/output parameters. Will conditioanlly load weight parameters, if load_weight_qparams is set to True. :param model: model to load pretrained quantization parameters to :param model_name_or_path: Hugging Face stub or local folder containing a quantized model, which is used to load quantization parameters :param load_weight_qparams: whether or not the weight quantization parameters should be loaded """ model_path = get_safetensors_folder(model_name_or_path) mapping = get_quantization_parameter_to_path_mapping(model_path) for name, submodule in model.named_modules(): if not is_module_quantized(submodule): continue if submodule.quantization_scheme.input_activations is not None: base_name = "input" _load_quant_args_from_mapping( base_name=base_name, module_name=name, module=submodule, mapping=mapping, ) if submodule.quantization_scheme.output_activations is not None: base_name = "output" _load_quant_args_from_mapping( base_name=base_name, module_name=name, module=submodule, mapping=mapping, ) if load_weight_qparams and submodule.quantization_scheme.weights: base_name = "weight" _load_quant_args_from_mapping( base_name=base_name, module_name=name, module=submodule, mapping=mapping, ) def apply_quantization_config( model: Module, config: QuantizationConfig | None, run_compressed: bool = False ): """ Initializes the model for quantization in-place based on the given config. Optionally coverts quantizable modules to compressed_linear modules :param model: model to apply quantization config to :param config: quantization config :param run_compressed: Whether the model will be run in compressed mode or decompressed fully on load """ config = deepcopy(config) if config is None: # see PR #180 return dict() # force zero points during initialization # TODO: remove zero points from initialization force_zero_point = config.quantization_status < QuantizationStatus.COMPRESSED # apply and initialize kv cache quantization if config.kv_cache_scheme is not None: _apply_kv_cache_scheme( model, config.kv_cache_scheme, config.quantization_status ) # build mapping of targets to schemes for easier matching # use ordered dict to preserve target ordering in config target_to_scheme = OrderedDict() for scheme in config.config_groups.values(): for target in scheme.targets: target_to_scheme[target] = scheme # mark appropriate layers for quantization by setting their quantization schemes for name, submodule in match_named_modules( model, target_to_scheme, config.ignore, warn_on_fail=True ): # mark modules to be quantized by adding # quant scheme to the matching layers matched_targets = match_targets(name, submodule, target_to_scheme) scheme = _scheme_from_targets(target_to_scheme, matched_targets, name) # target matched - add layer and scheme to target list submodule.quantization_scheme = scheme if is_attention_module(submodule) and is_narrow_match( model, scheme.targets, name ): initialize_hooked_attention(model, submodule) initialize_module_for_quantization(submodule, force_zero_point=force_zero_point) submodule.quantization_status = config.quantization_status def _apply_kv_cache_scheme( model: torch.nn.Module, kv_cache_scheme: QuantizationArgs, status: QuantizationStatus, ): if not kv_cache_scheme.symmetric: raise logger.warning("vLLM does not support asymmetric kv cache quantization") # applies and initializes kv cache quantization # this step cannot come after attention apply/initialize # otherwise it will override the attention qparams scheme = QuantizationScheme( targets=KV_CACHE_TARGETS.copy(), # is never read in practice input_activations=kv_cache_scheme, ) for submodule in model.modules(): if is_attention_module(submodule): submodule.quantization_scheme = scheme initialize_hooked_kv_cache(model, submodule) initialize_module_for_quantization(submodule, force_zero_point=False) submodule.quantization_status = status def _load_quant_args_from_mapping( base_name: str, module_name: str, module: Module, mapping: dict ): # TODO: skip update and just register here, don't do it in initialize """ Loads scale and zero point from a state_dict into the specified module :param base_name: quantization target, one of: weights, input_activations or output_activations :param module_name: pytorch module name to look up in state_dict :module: pytorch module associated with module_name :mapping: mapping to search fetch paths on disk for a given parameter """ scale_name = f"{base_name}_scale" zp_name = f"{base_name}_zero_point" g_idx_name = f"{base_name}_g_idx" state_dict_scale_path = mapping.get(f"{module_name}.{scale_name}", None) state_dict_zp_path = mapping.get(f"{module_name}.{zp_name}", None) state_dict_g_idx_path = mapping.get(f"{module_name}.{g_idx_name}", None) if state_dict_g_idx_path is not None: with safe_open(state_dict_g_idx_path, framework="pt", device="cpu") as f: state_dict_g_idx = f.get_tensor(f"{module_name}.{g_idx_name}") update_offload_parameter(module, g_idx_name, state_dict_g_idx) if state_dict_scale_path is not None: # module is quantized with safe_open(state_dict_scale_path, framework="pt", device="cpu") as f: state_dict_scale = f.get_tensor(f"{module_name}.{scale_name}") update_offload_parameter(module, scale_name, state_dict_scale) if state_dict_zp_path is None: # fill in zero point for symmetric quantization state_dict_zp = torch.zeros_like(state_dict_scale, device="cpu") else: with safe_open(state_dict_zp_path, framework="pt", device="cpu") as f: state_dict_zp = f.get_tensor(f"{module_name}.{zp_name}") update_offload_parameter(module, zp_name, state_dict_zp) def _scheme_from_targets( target_to_scheme: OrderedDict[str, QuantizationScheme], targets: list[str], name: str, ) -> QuantizationScheme: # return the first scheme (the prioritized one, # since the order of target_to_scheme matters) return target_to_scheme[targets[0]] vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/lifecycle/compressed.py000066400000000000000000000034631521257237700344060ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import logging import torch from compressed_tensors.quantization.lifecycle.forward import quantize from compressed_tensors.quantization.quant_config import QuantizationStatus from torch.nn import Module __all__ = [ "compress_quantized_weights", ] _LOGGER = logging.getLogger(__name__) def compress_quantized_weights(module: Module): """ Quantizes the module weight representation to use fewer bits in memory apply to full model with `model.apply(compress_quantized_weights)` :param module: module to compress to quantized representation """ scheme = getattr(module, "quantization_scheme", None) if not scheme or not scheme.weights: # no quantization scheme or weights not quantized, nothing to do return status = getattr(module, "quantization_status", None) if status is QuantizationStatus.COMPRESSED: # module is already compressed, nothing to do return weight = getattr(module, "weight", None) scale = getattr(module, "weight_scale", None) zero_point = getattr(module, "weight_zero_point", None) g_idx = getattr(module, "weight_g_idx", None) if weight is None or scale is None: # no weight, scale, or ZP, nothing to do # mark as compressed here to maintain consistent status throughout the model module.quantization_status = QuantizationStatus.COMPRESSED return module.weight.requires_grad = False # cannot use auto grad after compression module.weight.data = quantize( x=weight, scale=scale, zero_point=zero_point, g_idx=g_idx, args=scheme.weights, dtype=torch.int8, ) module.quantization_status = QuantizationStatus.COMPRESSED vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/lifecycle/forward.py000066400000000000000000000251031521257237700337010ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from functools import wraps import torch from compressed_tensors.quantization.lifecycle.forward_helpers import ( _apply_quantize_op, _process_block, _process_group, ) from compressed_tensors.quantization.quant_args import ( DynamicType, QuantizationArgs, QuantizationStrategy, ) from compressed_tensors.quantization.quant_config import QuantizationStatus from compressed_tensors.quantization.quant_scheme import QuantizationScheme from compressed_tensors.quantization.utils import ( calculate_range, compute_dynamic_scales_and_zp, ) from compressed_tensors.utils import patch_attr from torch.nn import Module __all__ = [ "quantize", "dequantize", "fake_quantize", "set_forward_quantized", "forward_quantize", ] @torch.no_grad() def quantize( x: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, args: QuantizationArgs, dtype: torch.dtype | None = None, g_idx: torch.Tensor | None = None, global_scale: torch.Tensor | None = None, ) -> torch.Tensor: """ Quantize the input tensor x using the QuantizationStrategy specified in args. Quantization can be done per tensor, channel, token or group. For group quantization, the group_size must be divisible by the column size. The input scale and zero_points are reshaped to support vectorization (Assumes 1 is the channel dimension) :param x: Input tensor :param scale: scale tensor :param zero_point: zero point tensor :param args: quantization args dictating how to quantize x :param dtype: optional dtype to cast the quantized output to :param g_idx: optional mapping from column index to group index :param global_scale: optional constant to scale the quantization scale during QDQ :return: fake quantized tensor """ return _process_quantization( x=x, scale=scale, zero_point=zero_point, args=args, dtype=dtype, do_quantize=True, do_dequantize=False, g_idx=g_idx, global_scale=global_scale, ) @torch.no_grad() def dequantize( x_q: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor | None = None, args: QuantizationArgs | None = None, dtype: torch.dtype | None = None, g_idx: torch.Tensor | None = None, global_scale: torch.Tensor | None = None, ) -> torch.Tensor: """ Dequantize a quantized input tensor x_q based on the strategy specified in args. If args is not provided, the strategy will be inferred. :param x: quantized input tensor :param scale: scale tensor :param zero_point: zero point tensor :param args: quantization args used to quantize x_q :param dtype: optional dtype to cast the dequantized output to :param g_idx: optional mapping from column index to group index :param global_scale: optional constant to scale the quantization scale during QDQ :return: dequantized float tensor """ if args is None: if scale.ndim == 0 or scale.ndim == 1: args = QuantizationArgs(strategy=QuantizationStrategy.TENSOR) elif scale.ndim == 2: if scale.shape[1] == 1: args = QuantizationArgs(strategy=QuantizationStrategy.CHANNEL) # Scale height matches input or is 1 -> group quantization across columns # # Example 1: scale.shape[0] == 1 # x_q: (4, 8), scale: (1, 4) -> 2 columns per group # # Example 2: scale.shape[0] == x_q.shape[0] # x_q: (4, 8), scale: (4, 4) -> 2 elements per group (per row) elif (scale.shape[0] == 1) or (scale.shape[0] == x_q.shape[0]): group_size = int(x_q.shape[1] / scale.shape[1]) args = QuantizationArgs( strategy=QuantizationStrategy.GROUP, group_size=group_size ) else: rows, cols = x_q.shape[-2], x_q.shape[-1] block_height = rows // scale.shape[0] # Rows per block block_width = cols // scale.shape[1] # Columns per block args = QuantizationArgs( strategy=QuantizationStrategy.BLOCK, block_structure=[block_height, block_width], ) else: raise ValueError( f"Could not infer a quantization strategy from scale with {scale.ndim} " "dimmensions. Expected 0 or 2 dimmensions." ) if dtype is None: dtype = scale.dtype return _process_quantization( x=x_q, scale=scale, zero_point=zero_point, args=args, do_quantize=False, do_dequantize=True, dtype=dtype, g_idx=g_idx, global_scale=global_scale, ) @torch.no_grad() def fake_quantize( x: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, args: QuantizationArgs, g_idx: torch.Tensor | None = None, global_scale: torch.Tensor | None = None, ) -> torch.Tensor: """ Fake quantize the input tensor x by quantizing then dequantizing with the QuantizationStrategy specified in args. Quantization can be done per tensor, channel, token or group. For group quantization, the group_size must be divisible by the column size. The input scale and zero_points are reshaped to support vectorization (Assumes 1 is the channel dimension) :param x: Input tensor :param scale: scale tensor :param zero_point: zero point tensor :param args: quantization args dictating how to quantize x :param g_idx: optional mapping from column index to group index :param global_scale: optional constant to scale the quantization scale during QDQ :return: fake quantized tensor """ return _process_quantization( x=x, scale=scale, zero_point=zero_point, args=args, do_quantize=True, do_dequantize=True, g_idx=g_idx, global_scale=global_scale, ) @torch.no_grad() def _process_quantization( x: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, args: QuantizationArgs, g_idx: torch.Tensor | None = None, dtype: torch.dtype | None = None, do_quantize: bool = True, do_dequantize: bool = True, global_scale: torch.Tensor | None = None, ) -> torch.Tensor: q_min, q_max = calculate_range(args, x.device) if args.strategy == QuantizationStrategy.BLOCK: return _process_block( x, scale, zero_point, args, q_min, q_max, dtype, do_quantize, do_dequantize, global_scale, ) elif args.strategy in ( QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP, ): return _process_group( x, scale, zero_point, args, q_min, q_max, dtype, do_quantize, do_dequantize, g_idx, global_scale, ) else: # covers tensor, channel, token, and attn_head strategies return _apply_quantize_op( x, scale, zero_point, q_min, q_max, args, dtype, do_quantize, do_dequantize, global_scale, ) def set_forward_quantized(module: torch.nn.Linear | torch.nn.Embedding): """ Replace a linear or embedding module's forward function with one that performs on-the-fly QDQ. Note that weight quantiation will be skipped for compressed modules. All QDQ operations can be skipped by setting `module.quantization_enabled = False` :param module: linear or embedding module whose forward function will be replaced """ @wraps(module.forward.__func__) def quantized_forward( self: torch.nn.Linear | torch.nn.Embedding, input: torch.Tensor ) -> torch.Tensor: """ Quantized forward pass of a linear or embedding module :param self: instance of linear or embedding module :param input: input activations to this module :return: linear or embedding output """ scheme: QuantizationScheme | None = getattr(self, "quantization_scheme", None) status: QuantizationStatus | None = getattr(self, "quantization_status", None) enabled: bool = ( getattr(self, "quantization_enabled", True) and scheme is not None and status is not None ) weight = self.weight # onload once weight_data = weight.data if enabled and scheme.input_activations: input = forward_quantize(self, input, "input", scheme.input_activations) if enabled and scheme.weights and status < QuantizationStatus.COMPRESSED: weight_data = forward_quantize(self, weight_data, "weight", scheme.weights) with patch_attr(weight, "data", weight_data): output = self.__class__.forward(self, input) if enabled and scheme.output_activations: output = forward_quantize(self, output, "output", scheme.output_activations) return output module.forward = quantized_forward.__get__(module) def forward_quantize( module: Module, value: torch.Tensor, base_name: str, args: "QuantizationArgs" ) -> torch.Tensor: # in compressed mode, the weight is already compressed and quantized so we don't # need to run fake quantization # TODO: remove this line, as this is already guarded by `set_forward_quantized` if ( module.quantization_status >= QuantizationStatus.COMPRESSED and base_name == "weight" ): return value if value.numel() == 0: # if the tensor is empty, # skip quantization return value g_idx = getattr(module, "weight_g_idx", None) global_scale = getattr(module, f"{base_name}_global_scale", None) if args.dynamic in (True, DynamicType.LOCAL): # dynamic quantization - determine the scale/zp on the fly scale, zero_point = compute_dynamic_scales_and_zp( value=value, args=args, module=module, global_scale=global_scale ) else: # static quantization - get scale and zero point from layer scale = getattr(module, f"{base_name}_scale") zero_point = getattr(module, f"{base_name}_zero_point", None) return fake_quantize( x=value, scale=scale, zero_point=zero_point, args=args, g_idx=g_idx, global_scale=global_scale, ) forward_helpers.py000066400000000000000000000165051521257237700353520ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/lifecycle# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from math import ceil import torch from compressed_tensors.quantization.quant_args import ( QuantizationArgs, round_to_quantized_type_args, ) from compressed_tensors.quantization.utils import maybe_pad_tensor_for_block_quant def _apply_quantize_op( x: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor | None, q_min: torch.Tensor, q_max: torch.Tensor, args: QuantizationArgs, dtype: torch.dtype | None, do_quantize: bool, do_dequantize: bool, global_scale: torch.Tensor | None, ) -> torch.Tensor: """Dispatch to the appropriate quantization kernel.""" if do_quantize and do_dequantize: return _quantize_dequantize( x=x, scale=scale, zero_point=zero_point, q_min=q_min, q_max=q_max, args=args, global_scale=global_scale, ) elif do_quantize: return _quantize( x=x, scale=scale, zero_point=zero_point, q_min=q_min, q_max=q_max, args=args, dtype=dtype, global_scale=global_scale, ) else: return _dequantize( x_q=x, scale=scale, zero_point=zero_point, global_scale=global_scale, ) def _process_block( x: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, args: QuantizationArgs, q_min: torch.Tensor, q_max: torch.Tensor, dtype: torch.dtype | None, do_quantize: bool, do_dequantize: bool, global_scale: torch.Tensor | None, ) -> torch.Tensor: """Blockwise quantization: pad, reshape into 2D blocks, quantize, restore.""" original_shape = x.shape block_height, block_width = args.block_structure x = maybe_pad_tensor_for_block_quant(x, args.block_structure) padded_shape = x.shape # reshape into blocks and transpose to make each block contiguous num_rows_blocks = padded_shape[0] // block_height num_cols_blocks = padded_shape[1] // block_width x_blocks = x.reshape( num_rows_blocks, block_height, num_cols_blocks, block_width, ).transpose(1, 2) # expand scale/zero_point for block broadcasting sb = scale.unsqueeze(-1).unsqueeze(-1) zb = zero_point.unsqueeze(-1).unsqueeze(-1) if zero_point is not None else None x_blocks = _apply_quantize_op( x_blocks, sb, zb, q_min, q_max, args, dtype, do_quantize, do_dequantize, global_scale, ) # restore padded shape output = x_blocks.transpose(1, 2).reshape(padded_shape) # truncate to original dimensions if padding was applied if original_shape != padded_shape: output = output[tuple([slice(v) for v in original_shape])] return output def _process_group( x: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor | None, args: QuantizationArgs, q_min: torch.Tensor, q_max: torch.Tensor, dtype: torch.dtype | None, do_quantize: bool, do_dequantize: bool, g_idx: torch.Tensor | None, global_scale: torch.Tensor | None, ) -> torch.Tensor: """Group/tensor-group quantization: handle activation ordering, reshape into groups, quantize, restore.""" group_size = args.group_size output_dtype = dtype if dtype is not None else x.dtype columns = x.shape[-1] while scale.ndim < 2: scale = scale.unsqueeze(1) zero_point = zero_point.unsqueeze(1) if zero_point is not None else None if columns >= group_size and columns % group_size != 0: raise ValueError( "tensor column shape must be divisble " f"by the given group_size {group_size} but got {columns}" ) # support column-order (default) quantization as well as other orderings # such as activation ordering. Below checks if g_idx has been initialized is_column_order = g_idx is None or g_idx.device.type == "meta" or -1 in g_idx if not is_column_order: perm = torch.argsort(g_idx) x = x.index_select(-1, perm) # reshape last dim into (num_groups, group_size) reshaped_dims = (ceil(x.shape[-1] / group_size), group_size) x = x.unflatten(-1, reshaped_dims) output = _apply_quantize_op( x, scale.unsqueeze(-1), zero_point.unsqueeze(-1) if zero_point is not None else None, q_min, q_max, args, dtype, do_quantize, do_dequantize, global_scale, ) output = output.flatten(start_dim=-2).to(output_dtype) if not is_column_order: inv_perm = torch.argsort(perm) output = output.index_select(-1, inv_perm) return output @torch.no_grad() def _quantize_dequantize( x: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor | None, q_min: torch.Tensor, q_max: torch.Tensor, args: QuantizationArgs, global_scale: torch.Tensor | None = None, ) -> torch.Tensor: """ Fused quantize-then-dequantize in a single pass, avoiding: - Double scale/global_scale division - Intermediate quantized dtype allocation """ # compute effective scale once if global_scale is not None: scale = scale / global_scale scaled = x / scale if zero_point is not None: scaled += zero_point.to(x.dtype) # clamp and round (stays in float — no int8/fp8 intermediate) quantized = round_to_quantized_type_args( tensor=scaled, args=args, min=q_min, max=q_max ) # dequantize: subtract zero_point and multiply by scale # cast to scale.dtype to match _dequantize behavior dequant = quantized.to(scale.dtype) if zero_point is not None: dequant = dequant - zero_point.to(scale.dtype) return dequant * scale @torch.no_grad() def _quantize( x: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, q_min: torch.Tensor, q_max: torch.Tensor, args: QuantizationArgs, dtype: torch.dtype | None = None, global_scale: torch.Tensor | None = None, ) -> torch.Tensor: # if a global scale is optionally provided, use it # to further scale the local `scale` parameter if global_scale is not None: scale = scale / global_scale scaled = x / scale if zero_point is not None: scaled += zero_point.to(x.dtype) # clamp and round quantized_value = round_to_quantized_type_args( tensor=scaled, args=args, min=q_min, max=q_max ) if dtype is not None: quantized_value = quantized_value.to(dtype) return quantized_value @torch.no_grad() def _dequantize( x_q: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor | None = None, dtype: torch.dtype | None = None, global_scale: torch.Tensor | None = None, ) -> torch.Tensor: # if a global scale is optionally provided, use it # to further scale the local `scale` parameter if global_scale is not None: scale = scale / global_scale dequant_value = x_q.to(scale.dtype) if zero_point is not None: dequant_value = dequant_value - zero_point.to(scale.dtype) dequant_value = dequant_value * scale if dtype is not None: dequant_value = dequant_value.to(dtype) return dequant_value vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/lifecycle/helpers.py000066400000000000000000000006621521257237700337020ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ Miscelaneous helpers for the quantization lifecycle """ from torch.nn import Module __all__ = [ "enable_quantization", "disable_quantization", ] def enable_quantization(module: Module): module.quantization_enabled = True def disable_quantization(module: Module): module.quantization_enabled = False vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/lifecycle/initialize.py000066400000000000000000000262371521257237700344070ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import logging import math import torch from compressed_tensors.modeling import ( IMPL_ATTR, KV_CACHE_ATTR, QuantizedAttentionImpl, QuantizedKVCache, ) from compressed_tensors.offload import disable_onloading, unwrap_offload_forward from compressed_tensors.quantization import ( ActivationOrdering, DynamicType, QuantizationArgs, QuantizationMetadata, QuantizationScheme, QuantizationStatus, QuantizationStrategy, ) from compressed_tensors.quantization.lifecycle.forward import set_forward_quantized from compressed_tensors.quantization.utils import strategy_cdiv from compressed_tensors.utils import ( get_execution_device, get_head_dim, get_num_attn_heads, get_num_kv_heads, ) from torch.nn import Module, Parameter __all__ = [ "initialize_module_for_quantization", "is_attention_module", "initialize_qparams", "initialize_attn_qparams", ] _LOGGER = logging.getLogger(__name__) def initialize_module_for_quantization( module: Module, scheme: QuantizationScheme | None = None, force_zero_point: bool = True, ): """ Attaches appropriate scales, zero points, and observers to a layer given its target quantization scheme. Previously initialized scales and zero points will be removed from module if they no longer apply to the scheme :param module: module to set for calibration :param scheme: scheme to use for quantization. if None is provided, will attempt to use scheme stored in the module under `quantization_scheme`, if not provided, the layer will be skipped :param force_zero_point: whether to force initialization of a zero point for symmetric quantization """ from compressed_tensors.linear.compressed_linear import CompressedLinear # circ dep scheme = scheme or getattr(module, "quantization_scheme", None) if scheme is None: return QuantizationMetadata.clear_all_qparams(module) if is_attention_module(module): initialize_attn_qparams(module, scheme, force_zero_point) elif isinstance(module, (torch.nn.Linear, torch.nn.Embedding)): with disable_onloading(): weight = module.weight if scheme.input_activations is not None: initialize_qparams( module, "input", scheme.input_activations, observed_shape=weight.shape[-1:], observed_dtype=weight.dtype, force_zero_point=force_zero_point, ) if scheme.weights is not None: initialize_qparams( module, "weight", scheme.weights, observed_shape=weight.shape, observed_dtype=weight.dtype, force_zero_point=force_zero_point, ) if scheme.output_activations is not None: initialize_qparams( module, "output", scheme.output_activations, observed_shape=weight.shape[:-1], observed_dtype=weight.dtype, force_zero_point=force_zero_point, ) # CompressedLinear has its own forward method that handles decompression # Don't override it with the quantized forward if not isinstance(module, CompressedLinear): with unwrap_offload_forward(module): set_forward_quantized(module) else: raise ValueError(f"Quantization of module type {type(module)} is not supported") module.quantization_scheme = scheme module.quantization_status = QuantizationStatus.INITIALIZED def is_attention_module(module: Module): return "attention" in module.__class__.__name__.lower() and ( hasattr(module, "k_proj") or hasattr(module, "v_proj") or hasattr(module, "qkv_proj") or hasattr(module, "kv_b_proj") ) def initialize_qparams( module: Module, base_name: str, quantization_args: QuantizationArgs, observed_shape: tuple[int | None, ...], observed_dtype: torch.dtype, force_zero_point: bool = True, ): """ Initialize quantization parameters for a given basename according to the passed quantization args. The shape and dtype of the observed weight/activation must also be provided. Scales will always be initialized. Global scales are initialized depending on args. Zero points will be initialized if not symmetric or if `force_zero_point` is True. :param module: module to register qparams to :param base_name: base name of qparams, for example "input", "weight", "k", "v" :param quantization_args: arguments for quantization :param observed_shape: last (right-most) known dimensions of the observed weight/act :param observed_dtype: dtype of the observed weight/actt :param force_zero_point: force the zero_point parameter to be initialized """ strategy = quantization_args.strategy dynamic = quantization_args.dynamic actorder = quantization_args.actorder device = get_execution_device(module) # avoid performing intialization ops on cpu # Skip all intialization for fully dynamic quantization if dynamic is True: return # 0. Create global scale for tensor-group quantization if strategy == QuantizationStrategy.TENSOR_GROUP: init_global_scale = Parameter( torch.empty(1, dtype=torch.float32, device=device), requires_grad=False, ) module.register_parameter(f"{base_name}_global_scale", init_global_scale) # Skip scale/zp initialization for locally dynamic quantization if dynamic == DynamicType.LOCAL: return # 1. Infer expected scale/zp shape if strategy == QuantizationStrategy.TENSOR: expected_shape = (1,) elif strategy == QuantizationStrategy.TOKEN: raise ValueError("Cannot perform static token quantization") elif strategy == QuantizationStrategy.CHANNEL: if len(observed_shape) < 2: raise ValueError("Channel quant requires at least 2 observed dimensions") expected_shape = (observed_shape[-2], 1) elif strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP): assert quantization_args.group_size is not None if len(observed_shape) < 1: raise ValueError("Group quant requires at least 1 observed dimension") group_size = quantization_args.group_size num_groups = strategy_cdiv(observed_shape[-1], group_size, strategy) expected_shape = (*observed_shape[:-1], num_groups) # initialize activation ordering if applicable if actorder == ActivationOrdering.GROUP: init_g_idx = Parameter( torch.full((observed_shape[-1],), -1, device=device, dtype=torch.int), requires_grad=False, ) module.register_parameter(f"{base_name}_g_idx", init_g_idx) elif strategy == QuantizationStrategy.BLOCK: assert quantization_args.block_structure is not None if len(observed_shape) < 2: raise ValueError("Block quant requires at least 2 observed dimensions") block_structure = quantization_args.block_structure # NOTE: vllm kernels for block-quantization do not require # num_rows to be evenly divisible by block_structure[-2], # but num_cols does need to be evenly divisible by block_structure[-1] num_rows = math.ceil(observed_shape[-2] / block_structure[-2]) num_cols = strategy_cdiv(observed_shape[-1], block_structure[-1], strategy) expected_shape = (num_rows, num_cols) elif strategy == QuantizationStrategy.ATTN_HEAD: # (batch_size, num_attention_heads, seq_len, head_dim) if len(observed_shape) < 3: raise ValueError("Attention quant requires at least 3 observed dimensions") expected_shape = (observed_shape[-3], 1, 1) else: assert False, f"Unknown strategy {strategy}" # 2. Identify quantization scale and zp dtype scale_dtype = observed_dtype if scale_dtype not in [ torch.float16, torch.bfloat16, torch.float32, torch.float64, ]: scale_dtype = torch.float16 # 3. Initializes scale/zp for the module init_scale = Parameter( torch.empty(expected_shape, dtype=scale_dtype, device=device), requires_grad=False, ) module.register_parameter(f"{base_name}_scale", init_scale) if force_zero_point or not quantization_args.symmetric: init_zero_point = Parameter( torch.zeros( expected_shape, device=device, dtype=quantization_args.zp_dtype ), requires_grad=False, ) module.register_parameter(f"{base_name}_zero_point", init_zero_point) def initialize_attn_qparams( module: Module, scheme: QuantizationScheme, force_zero_point: bool ): """Initlaize k_scale, v_scale for self_attn""" impl: QuantizedAttentionImpl | None = getattr(module, IMPL_ATTR, None) kv_cache: QuantizedKVCache | None = getattr(module, KV_CACHE_ATTR, None) if impl is None and kv_cache is None: raise ValueError( f"Attention module has quantization scheme but no {IMPL_ATTR} " f"or {KV_CACHE_ATTR} attributes. Please ensure that these " "attributes are initialized using `apply_quantization_config`." ) _validate_attention_scheme(scheme) # extract shapes from config config = kv_cache.config num_attn_heads = get_num_attn_heads(config) num_kv_heads = get_num_kv_heads(config) head_dim = get_head_dim(config) # (batch_size, num_heads, slen, head_dim) q_observed_shape = (num_attn_heads, None, head_dim) kv_observed_shape = (num_kv_heads, None, head_dim) observed_dtype = next(module.parameters()).dtype if impl is not None: initialize_qparams( module, "q", scheme.input_activations, observed_shape=q_observed_shape, observed_dtype=observed_dtype, force_zero_point=force_zero_point, ) if kv_cache is not None: initialize_qparams( module, "k", scheme.input_activations, observed_shape=kv_observed_shape, observed_dtype=observed_dtype, force_zero_point=force_zero_point, ) initialize_qparams( module, "v", scheme.input_activations, observed_shape=kv_observed_shape, observed_dtype=observed_dtype, force_zero_point=force_zero_point, ) def _validate_attention_scheme(scheme: QuantizationScheme): if scheme.weights is not None: raise ValueError( "Cannot apply weight quantization to attention. " "Instead, target the (q|k|v)_proj submodule layers of attention" ) if scheme.input_activations is None: raise ValueError( "Cannot apply attention quantization without specifying input activations" ) if scheme.output_activations is not None: raise ValueError("Cannot apply output quantization to attention") vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/quant_args.py000066400000000000000000000374611521257237700324540ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import warnings from enum import Enum from typing import Any import torch from compressed_tensors.utils import Aliasable from compressed_tensors.utils.type import TorchDtype from pydantic import ( BaseModel, ConfigDict, Field, field_serializer, field_validator, model_validator, ) __all__ = [ "FP8_E4M3_DATA", "FP4_E2M1_DATA", "BFLOAT16_DATA", "FLOAT16_DATA", "FLOAT32_DATA", "FLOAT64_DATA", "FloatArgs", "QuantizationType", "QuantizationStrategy", "QuantizationArgs", "round_to_quantized_type_args", "round_to_quantized_type_dtype", "ActivationOrdering", "DynamicType", ] class FloatArgs: exponent: int mantissa: int bits: int | None = None max: float | None = None min: float | None = None dtype: torch.dtype | None = None class FP4_E2M1_DATA(FloatArgs): exponent = 2 mantissa = 1 bits = 4 max = 6.0 min = -6.0 @staticmethod @torch.compile def cast_to_fp4(x): sign = torch.sign(x) x = torch.abs(x) x[(x >= 0.0) & (x <= 0.25)] = 0.0 x[(x > 0.25) & (x < 0.75)] = 0.5 x[(x >= 0.75) & (x <= 1.25)] = 1.0 x[(x > 1.25) & (x < 1.75)] = 1.5 x[(x >= 1.75) & (x <= 2.5)] = 2.0 x[(x > 2.5) & (x < 3.5)] = 3.0 x[(x >= 3.5) & (x <= 5.0)] = 4.0 x[x > 5.0] = 6.0 return x * sign class FP8_E4M3_DATA(FloatArgs): exponent = 4 mantissa = 3 bits = 8 max = torch.finfo(torch.float8_e4m3fn).max min = torch.finfo(torch.float8_e4m3fn).min dtype = torch.float8_e4m3fn class BFLOAT16_DATA(FloatArgs): exponent = 8 mantissa = 7 class FLOAT16_DATA(FloatArgs): exponent = 5 mantissa = 10 class FLOAT32_DATA(FloatArgs): exponent = 8 mantissa = 23 class FLOAT64_DATA(FloatArgs): exponent = 11 mantissa = 52 class QuantizationType(str, Enum): """ Enum storing quantization type options """ INT = "int" FLOAT = "float" class QuantizationStrategy(str, Enum): """ Enum storing quantization strategy options """ TENSOR = "tensor" CHANNEL = "channel" GROUP = "group" BLOCK = "block" TOKEN = "token" TENSOR_GROUP = "tensor_group" ATTN_HEAD = "attn_head" class DynamicType(str, Enum): """ Enum storing potential dynamic types. 1. If dynamic is True, all quantization parameters are generated on the fly. 2. If dynamic is False, all quantization parameters generated are static. 3. If "local" is provided, only local quantization parameters are dynamic. Note: "local" is only currently supported for NVFP4. """ LOCAL = "local" class ActivationOrdering(Aliasable, str, Enum): """ Enum storing strategies for activation ordering during GPTQ calibration Group: Columns are permuted by activation order during calibration. Quantization groups are defined based on this permuted order. Weights are saved in original column order with g_idx mapping columns to groups. Runtime requires reordering columns by g_idx (higher latency but improved accuracy compared to no activation ordering).\n Weight: Changes the way calibration occurs but doesn't change the quantization format compared to no activation ordering (normal latency). Compared to Group, it has lower latency and slightly worse accuracy. Compared to no activation ordering during calibration it has slightly better accuracy. \n Dynamic: alias for Group\n Static: alias for Weight\n """ GROUP = "group" WEIGHT = "weight" # aliases DYNAMIC = "dynamic" STATIC = "static" @staticmethod def get_aliases() -> dict[str, str]: return { "dynamic": "group", "static": "weight", } class QuantizationArgs(BaseModel, use_enum_values=True): """ User facing arguments used to define a quantization config for weights or activations :param num_bits: quantization bit depth :param type: dtype to quantized to, either int or float :param symmetric: whether or not quantization scale is symmetric about zero-point :param strategy: string id determining the scope of scale/zero-point to apply :param group_size: group length to use for the group strategy :param block_structure: 2d block structure to use for the block strategy; must be a list of two ints [rows, cols] like [128, 128]. :param dynamic: set True to perform dynamic quantization - values will not be calibrated during calibration phase, instead during inference new quantization ranges will be observed with every sample. Defaults to False for static quantization. Note that enabling dynamic quantization will change the default observer to a memoryless one :param actorder: activation ordering strategy for GPTQ calibration. Options are GROUP (reorder by activation with g_idx mapping, higher accuracy but higher latency), WEIGHT (reorder during calibration only, normal latency with slight accuracy improvement), or None (no activation ordering). See ActivationOrdering enum for detailed explanations. Defaults to None """ num_bits: int = 8 type: QuantizationType = QuantizationType.INT symmetric: bool = True group_size: int | None = None strategy: QuantizationStrategy | None = None block_structure: list[int] | None = None dynamic: DynamicType | bool = False actorder: ActivationOrdering | bool | None = None scale_dtype: TorchDtype | None = None zp_dtype: TorchDtype | None = None observer: str | None = Field( default=None, description=( "Determines the method of computing quantization parameters (scales and " "zero-points). Defaults to min-max when not using dynamic quantization" ), ) observer_kwargs: dict[str, Any] = Field( default_factory=dict, description=( "optional dict of kwargs to be passed directly to torch quantization " "Observers constructor excluding quantization range or symmetry" ), ) @field_serializer("zp_dtype") def serialize_dtype(self, dtype: torch.dtype): if self.symmetric: return None return str(dtype) @field_validator("type", mode="before") def validate_type(cls, value) -> QuantizationType: if isinstance(value, str): return QuantizationType(value.lower()) return value @field_validator("group_size", mode="before") def validate_group(cls, value) -> int | None: if value is None: return value if value < -1: raise ValueError( f"Invalid group size {value}. Use group_size > 0 for " "strategy='group' and group_size = -1 for 'channel'" ) return value @field_validator("block_structure", mode="before") def validate_block_structure(cls, value) -> list[int] | None: if value is None: return value # For backward compatibility, allow string format "2x4", "8x16", etc. if isinstance(value, str): try: return [int(x) for x in value.split("x")] except Exception: raise ValueError( f"Invalid block_structure '{value}'. Must be a list of ints " "[rows, cols]." ) if isinstance(value, (list, tuple)): if len(value) != 2 or not all(isinstance(v, int) for v in value): raise ValueError( f"Invalid block_structure '{value}'. Must be a list of ints " "[rows, cols]." ) return list(value) raise ValueError( f"Invalid block_structure '{value}'. Must be a list of ints [rows, cols]." ) @field_validator("strategy", mode="before") def validate_strategy(cls, value) -> QuantizationStrategy | None: if isinstance(value, str): return QuantizationStrategy(value.lower()) return value @field_validator("actorder", mode="before") def validate_actorder(cls, value) -> ActivationOrdering | None: if isinstance(value, bool): return ActivationOrdering.GROUP if value else None if isinstance(value, str): return ActivationOrdering(value.lower()) return value @field_validator("dynamic", mode="before") def validate_dynamic(cls, value) -> DynamicType | bool: if isinstance(value, str): return DynamicType(value.lower()) return value @model_validator(mode="after") def validate_model_after(model: "QuantizationArgs") -> "QuantizationArgs": # extract user-passed values from dictionary strategy = model.strategy group_size = model.group_size block_structure = model.block_structure actorder = model.actorder dynamic = model.dynamic observer = model.observer dynamic = model.dynamic zp_dtype = model.zp_dtype # infer strategy if strategy is None: if group_size is None: strategy = QuantizationStrategy.TENSOR elif group_size > 0: strategy = QuantizationStrategy.GROUP elif group_size == -1: strategy = QuantizationStrategy.CHANNEL else: raise ValueError( f"Invalid group size {group_size}. Use group_size > 0 for " "strategy='group' and group_size = -1 for 'channel'" ) # validate token strategy if strategy == QuantizationStrategy.TOKEN and not dynamic: raise ValueError( "Cannot perform static token quantization, please use `dynamic=True`" ) # validate group strategy if strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP): if group_size is None or group_size <= 0: raise ValueError( f"strategy {strategy} requires group_size to be " "set to a positive value" ) if ( group_size is not None and group_size > 0 and strategy not in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP) ): raise ValueError("group_size requires strategy to be set to 'group'") # validate block strategy has_block_strategy = strategy == QuantizationStrategy.BLOCK has_block_structure = block_structure is not None if has_block_strategy and not has_block_structure: raise ValueError(f"Block strategy requires block structure\n{model}") if has_block_structure and not has_block_strategy: raise ValueError(f"Block structure requires block strategy\n{model}") # validate activation ordering and strategy if actorder == ActivationOrdering.GROUP and strategy not in ( QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP, ): raise ValueError( "Must use group or tensor_group quantization strategy in " "order to apply group activation ordering" ) # infer observer w.r.t. dynamic if dynamic: supported_strategies = ( QuantizationStrategy.TOKEN, QuantizationStrategy.TENSOR, QuantizationStrategy.TENSOR_GROUP, QuantizationStrategy.GROUP, ) if strategy not in supported_strategies: raise ValueError( f"One of {supported_strategies} must be used for dynamic quant." ) if ( dynamic == DynamicType.LOCAL and strategy != QuantizationStrategy.TENSOR_GROUP ): raise ValueError("local is only supported for strategy tensor_group") if observer is not None: if dynamic is True: # checking if dynamic is True, not "local" if ( observer != "memoryless" ): # avoid annoying users with old configs warnings.warn( "No observer is used for dynamic quant., setting to None" ) observer = None else: if dynamic == DynamicType.LOCAL: observer = "minmax" elif observer is None: # default to minmax for non-dynamic cases observer = "memoryless_minmax" if zp_dtype is None: if model.num_bits == 4 and model.type == QuantizationType.FLOAT: zp_dtype = FP8_E4M3_DATA.dtype else: zp_dtype = model.pytorch_dtype() # write back modified values model.strategy = strategy model.observer = observer model.zp_dtype = zp_dtype return model def pytorch_dtype(self) -> torch.dtype: if self.type == QuantizationType.FLOAT: if self.num_bits == 8: return FP8_E4M3_DATA.dtype else: raise NotImplementedError("Only num_bits in (8) are supported") elif self.type == QuantizationType.INT: if self.num_bits <= 8: return torch.int8 elif self.num_bits <= 16: return torch.int16 else: return torch.int32 else: raise ValueError(f"Invalid quantization type {self.type}") model_config = ConfigDict(extra="forbid") def round_to_quantized_type_dtype( tensor: torch.Tensor, dtype: torch.dtype, cast_to_original_dtype: bool = True, ) -> torch.Tensor: """ Rounds an input tensor to the nearest quantized representation given a dtype. The original dtype is kept post-rounding. :param tensor: tensor to round :param dtype: dtype to use for rounding :param cast_to_original_dtype: whether or not we cast the rounded tensor to the original dtype :return: rounded tensor """ original_dtype = tensor.dtype if torch.is_floating_point(torch.tensor([], dtype=dtype)): finfo = torch.finfo(dtype) rounded = torch.clamp(tensor, finfo.min, finfo.max).to(dtype) else: iinfo = torch.iinfo(dtype) rounded = torch.round(torch.clamp(tensor, iinfo.min, iinfo.max)).to(dtype) if cast_to_original_dtype: return rounded.to(original_dtype) return rounded def round_to_quantized_type_args( tensor: torch.Tensor, args: QuantizationArgs, min: torch.Tensor, max: torch.Tensor, cast_to_original_dtype: bool = True, ) -> torch.Tensor: """ Rounds an input tensor to the nearest quantized representation given qunatization args. The original dtype is kept post-rounding. :param tensor: tensor to round :param args: quantization args to use for rounding :param min: min value to use for clamping :param max: max value to use for clamping :param cast_to_original_dtype: whether or not we cast the rounded tensor to the original dtype :return: rounded tensor """ original_dtype = tensor.dtype tensor = torch.clamp(tensor, min, max) if args.type == QuantizationType.FLOAT: if args.num_bits == 8: rounded = tensor.to(FP8_E4M3_DATA.dtype) elif args.num_bits == 4: rounded = FP4_E2M1_DATA.cast_to_fp4(tensor) else: raise NotImplementedError("Only num_bits in (4, 8) are supported") elif args.type == QuantizationType.INT: rounded = torch.round(tensor) else: raise ValueError(f"Invalid quantization type {args.type}") if cast_to_original_dtype: return rounded.to(original_dtype) return rounded vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/quant_config.py000066400000000000000000000272641521257237700327650ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import defaultdict from enum import Enum from typing import Annotated, Any import torch from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization.quant_args import DynamicType, QuantizationArgs from compressed_tensors.quantization.quant_scheme import ( QuantizationScheme, preset_name_to_scheme, ) from compressed_tensors.quantization.utils import is_module_quantized from pydantic import BaseModel, ConfigDict, Field from torch.nn import Module __all__ = [ "QuantizationStatus", "QuantizationConfig", "LIFECYCLE_ORDER", "DEFAULT_QUANTIZATION_METHOD", "DEFAULT_QUANTIZATION_FORMAT", ] class QuantizationStatus(str, Enum): """ Enum storing the different states a quantized layer can be in - Initialized: Quantization parameters are initialized to empty values - Calibration: Quantization parameters are being calibrated, observers are attached - Frozen: Quantization parameters are fully calibrated, observers are removed - Compressed: All parameters are quantized to target dtype. If the weight param is no longer applicable (i.e. if "weight" has been converted to "weight_packed"), the weight param is pruned from the module. - Decompressed: Parameters are converted back into frozen state. Quantization params remain on the module so that the module can be compressed, but params are pruned if they are no longer applicable or needed for compression (i.e. if "weight_packed" has been converted to "weight"). Additionally, weight qdq is skipped during forward passes for better performance """ INITIALIZED = "initialized" CALIBRATION = "calibration" FROZEN = "frozen" COMPRESSED = "compressed" DECOMPRESSED = "decompressed" @classmethod def lifecycle_order(cls) -> list["QuantizationStatus"]: """ :return: list of correct quantization lifecycle order """ return def __ge__(self, other): if other is None: return True if not isinstance(other, self.__class__): raise NotImplementedError return LIFECYCLE_ORDER.index(self) >= LIFECYCLE_ORDER.index(other) def __gt__(self, other): if other is None: return True if not isinstance(other, self.__class__): raise NotImplementedError return LIFECYCLE_ORDER.index(self) > LIFECYCLE_ORDER.index(other) def __lt__(self, other): if other is None: return False if not isinstance(other, self.__class__): raise NotImplementedError return LIFECYCLE_ORDER.index(self) < LIFECYCLE_ORDER.index(other) def __le__(self, other): if other is None: return False if not isinstance(other, self.__class__): raise NotImplementedError return LIFECYCLE_ORDER.index(self) <= LIFECYCLE_ORDER.index(other) LIFECYCLE_ORDER = [ QuantizationStatus.INITIALIZED, QuantizationStatus.CALIBRATION, QuantizationStatus.FROZEN, QuantizationStatus.COMPRESSED, QuantizationStatus.DECOMPRESSED, ] DEFAULT_QUANTIZATION_METHOD = "compressed-tensors" DEFAULT_QUANTIZATION_FORMAT = "fakequant" class QuantizationConfig(BaseModel): """ Full configuration specifying how a model is quantized. Each quantized layer is mapped to a QuantizationScheme in config_groups. :param config_groups: dict of QuantizationSchemes specifying the quantization settings for each quantized layer. A group could also be a reference to a predefined scheme name, mapped to a list of its target layers/classes :param quant_method: a constant used to differentiate compressed-tensors quantization from other quantization configs :param format: specifies how the quantized model is stored on disk :quantization_status: specifies the current status of all quantized layers. It is assumed all layers are in the same state. :param kv_cache_scheme: optional QuantizationArgs, that specify the quantization of the kv cache. If None, kv cache is not quantized. When applying kv cache quantization to transformer AutoModelForCausalLM, the kv_cache_scheme gets converted into a QuantizationScheme that: - targets the `q_proj` and `k_proj` modules of the model. The outputs of those modules are the keys and values that might be cached - quantizes the outputs of the aformentioned layers, so that keys and values are compressed before storing them in the cache There is an explicit assumption that the model contains modules with `k_proj` and `v_proj` in their names. If this is not the case and kv_cache_scheme != None, the quantization of kv cache will fail :global_compression_ratio: optional informational config to report the model compression ratio acheived by the quantization config :ignore: optional list of layers to ignore from config_groups. Layers in this list are not quantized even if they match up with a target in config_groups """ config_groups: dict[str, QuantizationScheme | list[str]] quant_method: str = DEFAULT_QUANTIZATION_METHOD kv_cache_scheme: QuantizationArgs | None = None format: str = DEFAULT_QUANTIZATION_FORMAT quantization_status: QuantizationStatus = QuantizationStatus.INITIALIZED global_compression_ratio: float | None = None ignore: list[str] | None = Field(default_factory=list) # `run_compressed` is a dummy, unused arg for backwards compatibility # see: https://github.com/huggingface/transformers/pull/39324 run_compressed: Annotated[Any, Field(exclude=True)] = None def model_post_init(self, __context): """ updates any quantization schemes defined as presets to be fully loaded schemes """ for group_name, targets_or_scheme in self.config_groups.items(): if isinstance(targets_or_scheme, QuantizationScheme): continue # scheme already defined self.config_groups[group_name] = preset_name_to_scheme( name=group_name, targets=targets_or_scheme, ) def to_dict(self): # for compatibility with HFQuantizer return self.model_dump() @staticmethod def from_pretrained( model: Module, format: str | list | None = None ) -> "QuantizationConfig | None": """ Converts a model into its associated QuantizationConfig based on the QuantizationScheme attached to each quantized module :param model: model to calculate quantization scheme of :return: filled out QuantizationScheme for the input model """ from compressed_tensors.modeling import IMPL_ATTR from compressed_tensors.quantization.lifecycle.initialize import ( is_attention_module, ) # set of all quantization schemes # TODO: make quant config/scheme/args frozen/hashable and use a set quantization_schemes: list[QuantizationScheme] = list() # use any status from modules (in practice, use the last module) model_status = None # set of all quantized types # this is later used to create the ignore list quantization_type_names: set[str] = set() # maps types to names which are not quantized # this is later used to create the ignore list ignore: dict[str, list[str]] = defaultdict(list) # this keeps track of any kvcache schemes kv_cache_scheme: QuantizationArgs | None = None for name, submodule in model.named_modules(): layer_type: str = get_vllm_module_type(submodule) # add config group if quantized non-attention or attention quant has_config_group = is_module_quantized(submodule) and ( not is_attention_module(submodule) or hasattr(submodule, IMPL_ATTR) ) # only add kvcache if quant attention (which always implies kvcache) has_kv_cache = is_module_quantized(submodule) and is_attention_module( submodule ) if has_config_group: # add to running set of schemes/layer_type_names model_status = getattr(submodule, "quantization_status", model_status) quantization_type_names.add(layer_type) if submodule.quantization_scheme not in quantization_schemes: quantization_schemes.append(submodule.quantization_scheme) if has_kv_cache: model_status = getattr(submodule, "quantization_status", model_status) kv_cache_scheme = submodule.quantization_scheme.input_activations if not has_config_group: # add non-quantized layers to the ignore list if layer_type not in ignore: ignore[layer_type] = [] ignore[layer_type].append(name) if ( len(quantization_schemes) == 0 and kv_cache_scheme is None ): # No quantized layers return None # create ignore list, only include layers whose class has ever been targeted consolidated_ignore = [] for layer_type, ignore_names in ignore.items(): if layer_type in quantization_type_names: # specific layers of a quantized type are ignored consolidated_ignore += ignore_names # else we leave it off the ignore list, doesn't fall under any of the # existing quantization schemes so it won't be quantized # create config groups from all unique schemes config_groups = {} for idx, scheme in enumerate(quantization_schemes): group_name = "group_" + str(idx) config_groups[group_name] = scheme # infer format if format is None: if model_status == QuantizationStatus.COMPRESSED: format = CompressionFormat.int_quantized.value else: format = CompressionFormat.dense.value elif isinstance(format, list): format = ( CompressionFormat.mixed_precision.value if len(format) > 1 else format[0] ) return QuantizationConfig( config_groups=config_groups, quantization_status=model_status, kv_cache_scheme=kv_cache_scheme, global_compression_ratio=None, format=format, ignore=consolidated_ignore, ) def requires_calibration_data(self): if self.kv_cache_scheme is not None: return True for _, scheme in self.config_groups.items(): if scheme.input_activations is not None: if scheme.input_activations.dynamic in (False, DynamicType.LOCAL): return True if scheme.output_activations is not None: if not scheme.output_activations.dynamic: return True return False # TODO set `extra="forbid"` when upstream transformers is compatible model_config = ConfigDict(extra="ignore") def get_vllm_module_type(module: torch.nn.Module) -> str: """ Returns a string representing the module type used when loading in vLLM. This is typically going to be the same as the `torch.nn.Module` type, however specific cases like MoE gate layers need to be treated like "Linear" layers for the purposes of config matching. """ module_type = type(module).__name__ if "Router" in module_type or "Gate" in module_type or "Gating" in module_type: module_type = "Linear" return module_type vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/quant_metadata.py000066400000000000000000000044511521257237700332710ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum from compressed_tensors.offload.module import unwrap_offload_forward from torch.nn import Module __all__ = ["QuantizationMetadata", "KVCacheScaleType"] class KVCacheScaleType(Enum): KEY = "k_scale" VALUE = "v_scale" QUERY = "q_scale" class QuantizationMetadata: """ Container class for metadata related to quantization """ @staticmethod def all_qparam_names(): """ All quantization parameter names that might be registered onto a module during lifecycle (excluding serialized parameters) """ return [KVCacheScaleType.KEY.value, KVCacheScaleType.VALUE.value] + [ f"{base_name}_{suffix}" for base_name in ("input", "weight", "output") for suffix in ( "global_scale", "scale", "shape", "zero_point", "g_idx", ) ] @classmethod def clear_all_qparams(cls, module: Module): """ Remove all parameters related to quantization that might have been registered onto a module previously in lifecycle (excluding serialized parameters) :param module: Module to clear """ for key in cls.all_qparam_names(): if hasattr(module, key): delattr(module, key) @classmethod def clear_quantization(cls, module: Module): """ Remove all artifacts of quantization from module, non-recursively. Artifacts include any qparams, quantization_scheme, or wrapped forward method that might have been altered previously in lifecycle. `quantization_status` and `quantization_enabled` are left unchanged. :param module: Module to clear """ with unwrap_offload_forward(module): # Unwrap forward call if hasattr(module.forward, "__wrapped__"): module.forward = module.forward.__wrapped__.__get__(module) # Clear any qparams cls.clear_all_qparams(module) # Clear quantization_scheme if hasattr(module, "quantization_scheme"): delattr(module, "quantization_scheme") vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/quant_scheme.py000066400000000000000000000270351521257237700327600ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import warnings from copy import deepcopy import torch from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization.quant_args import ( FP8_E4M3_DATA, DynamicType, QuantizationArgs, QuantizationStrategy, QuantizationType, ) from pydantic import BaseModel, ConfigDict, model_validator __all__ = [ "QuantizationScheme", "preset_name_to_scheme", "is_preset_scheme", ] class QuantizationScheme(BaseModel, use_enum_values=True): """ Set of QuantizationArgs defining how the weights, inputs and outputs of target list of modules should be quantized :param targets: list of modules to apply the QuantizationArgs to, can be layer names, layer types or a regular expression, typically ["Linear"] :param weights: quantization config for layer weights :param input_activations: quantization config for layer inputs :param output_activations: quantization config for layer outputs :param format: CompressionFormat for the layer """ targets: list[str] weights: QuantizationArgs | None = None input_activations: QuantizationArgs | None = None output_activations: QuantizationArgs | None = None format: CompressionFormat | None = None @model_validator(mode="after") def validate_model_after(model: "QuantizationScheme") -> "QuantizationScheme": inputs = model.input_activations outputs = model.output_activations weights = model.weights format = model.format if inputs is not None: if inputs.strategy not in ( QuantizationStrategy.TOKEN, QuantizationStrategy.TENSOR, QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP, QuantizationStrategy.ATTN_HEAD, ): if ( inputs.strategy == QuantizationStrategy.GROUP and inputs.dynamic is True ): raise NotImplementedError( "Static and local group-wise activation " "quantization is not supported" ) raise NotImplementedError( f"Using {inputs.strategy} strategy is not supported for " "activation quantization" ) if inputs.actorder is not None: raise ValueError("Cannot apply actorder to input activations") if outputs is not None: if outputs.actorder is not None: raise ValueError("Cannot apply actorder to output activations") if format == CompressionFormat.mixed_precision: raise ValueError( "mixed-precision cannot be set as a format for a QuantizationScheme" ) if ( inputs and weights and weights.strategy == QuantizationStrategy.GROUP and inputs.strategy == QuantizationStrategy.GROUP and weights.group_size != inputs.group_size ): warnings.warn( "Using GROUP strategy for both weights and input_activations " f"with different group sizes ({weights.group_size} vs " f"{inputs.group_size}) may complicate fused kernel implementations. " "Consider using TENSOR_GROUP strategy for both or matching group" " sizes.", UserWarning, stacklevel=2, ) return model model_config = ConfigDict(extra="forbid") """ Pre-Set Quantization Scheme Args """ def preset_name_to_scheme(name: str, targets: list[str]) -> QuantizationScheme: """ :param name: preset quantization settings name. must exist in upper case in PRESET_SCHEMES :param targets: list of quantization targets to be passed to the Scheme :return: new QuantizationScheme for a given name with the given targets """ name = name.upper() if name not in PRESET_SCHEMES: raise KeyError( f"Unknown preset scheme name {name}, " f"available names: {list(PRESET_SCHEMES.keys())}" ) scheme_args = deepcopy(PRESET_SCHEMES[name]) # deepcopy to avoid args references return QuantizationScheme( targets=targets, **scheme_args, ) def is_preset_scheme(name: str) -> bool: """ :param name: preset quantization settings name :return: True if the name is a preset scheme name """ return name.upper() in PRESET_SCHEMES UNQUANTIZED = dict() NVFP4A16 = dict( weights=QuantizationArgs( num_bits=4, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.TENSOR_GROUP, symmetric=True, dynamic=False, group_size=16, scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=FP8_E4M3_DATA.dtype, ) ) NVFP4 = dict( weights=QuantizationArgs( num_bits=4, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.TENSOR_GROUP, symmetric=True, dynamic=False, group_size=16, scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=FP8_E4M3_DATA.dtype, ), input_activations=QuantizationArgs( num_bits=4, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.TENSOR_GROUP, symmetric=True, dynamic=DynamicType.LOCAL, group_size=16, observer="static_minmax", scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=FP8_E4M3_DATA.dtype, ), ) MXFP4A16 = dict( weights=QuantizationArgs( num_bits=4, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.GROUP, symmetric=True, dynamic=False, group_size=32, scale_dtype=torch.uint8, zp_dtype=torch.uint8, ) ) MXFP4 = dict( weights=QuantizationArgs( num_bits=4, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.GROUP, symmetric=True, dynamic=False, group_size=32, scale_dtype=torch.uint8, zp_dtype=torch.uint8, ), input_activations=QuantizationArgs( num_bits=4, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.GROUP, dynamic=True, symmetric=True, group_size=32, scale_dtype=torch.uint8, zp_dtype=torch.uint8, ), ) MXFP8A16 = dict( weights=QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.GROUP, symmetric=True, dynamic=False, group_size=32, scale_dtype=torch.uint8, zp_dtype=torch.uint8, ) ) MXFP8 = dict( weights=QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.GROUP, symmetric=True, dynamic=False, group_size=32, scale_dtype=torch.uint8, zp_dtype=torch.uint8, ), input_activations=QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.GROUP, dynamic=True, symmetric=True, group_size=32, scale_dtype=torch.uint8, zp_dtype=torch.uint8, ), ) # 8 bit integer weights and 8 bit activations quantization INT8_W8A8 = dict( weights=QuantizationArgs( num_bits=8, type=QuantizationType.INT, strategy=QuantizationStrategy.CHANNEL, symmetric=True, dynamic=False, ), input_activations=QuantizationArgs( num_bits=8, type=QuantizationType.INT, strategy=QuantizationStrategy.TOKEN, symmetric=True, dynamic=True, ), ) # 8 bit integer weights only quantization W8A16 = dict( weights=QuantizationArgs( num_bits=8, type=QuantizationType.INT, strategy=QuantizationStrategy.GROUP, group_size=128, symmetric=True, dynamic=False, ), ) # 4 bit integer weights only quantization W4A16 = dict( weights=QuantizationArgs( num_bits=4, type=QuantizationType.INT, strategy=QuantizationStrategy.GROUP, group_size=128, symmetric=True, dynamic=False, ), ) # 4 bit integer weights only asymmetric quantization W4A16_ASYM = dict( weights=QuantizationArgs( num_bits=4, type=QuantizationType.INT, strategy=QuantizationStrategy.GROUP, group_size=128, symmetric=False, dynamic=False, ), ) # 4 bit integer weights and 8 bit activations quantization INT8_W4A8 = dict( weights=QuantizationArgs( num_bits=4, type=QuantizationType.INT, group_size=128, strategy=QuantizationStrategy.GROUP, symmetric=True, dynamic=False, ), input_activations=QuantizationArgs( num_bits=8, type=QuantizationType.INT, strategy=QuantizationStrategy.TOKEN, symmetric=True, dynamic=True, ), ) # 4 bit integer weights and 8 bit FP activations quantization W4AFP8 = dict( weights=QuantizationArgs( num_bits=4, type=QuantizationType.INT, strategy=QuantizationStrategy.GROUP, group_size=128, symmetric=True, dynamic=False, ), input_activations=QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.TOKEN, symmetric=True, dynamic=True, observer=None, ), ) # FP8 weights and FP8 activations quantization FP8 = dict( weights=QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.TENSOR, symmetric=True, dynamic=False, ), input_activations=QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.TENSOR, symmetric=True, dynamic=False, observer="static_minmax", ), ) # FP8 weights and FP8 dynamic activations quantization FP8_DYNAMIC = dict( weights=QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.CHANNEL, symmetric=True, dynamic=False, ), input_activations=QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.TOKEN, symmetric=True, dynamic=True, ), ) # Block‐wise FP8 (deepseekv3-style quantization): # static 128x128 per‐block weights and # dynamic per‐token‐group activations FP8_BLOCK = dict( weights=QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.BLOCK, symmetric=True, dynamic=False, block_structure=[128, 128], ), input_activations=QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.GROUP, symmetric=True, dynamic=True, group_size=128, ), ) PRESET_SCHEMES = { # Unquantized (no-op) "UNQUANTIZED": UNQUANTIZED, # Integer weight only schemes "W8A16": W8A16, "W4A16": W4A16, "W4A16_ASYM": W4A16_ASYM, # Integer weight and activation schemes "W8A8": INT8_W8A8, "INT8": INT8_W8A8, # alias for W8A8 "W4A8": INT8_W4A8, "W4AFP8": W4AFP8, # Float weight and activation schemes "FP8": FP8, "FP8_DYNAMIC": FP8_DYNAMIC, "FP8_BLOCK": FP8_BLOCK, "NVFP4A16": NVFP4A16, "NVFP4": NVFP4, "MXFP4A16": MXFP4A16, "MXFP4": MXFP4, "MXFP8A16": MXFP8A16, "MXFP8": MXFP8, } vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/utils/000077500000000000000000000000001521257237700310635ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/utils/__init__.py000066400000000000000000000002541521257237700331750ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .helpers import * from .mxfp_utils import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/utils/helpers.py000066400000000000000000000340261521257237700331040ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import logging import math import torch from compressed_tensors.quantization.quant_args import ( FP4_E2M1_DATA, FP8_E4M3_DATA, FloatArgs, QuantizationArgs, QuantizationStrategy, QuantizationType, round_to_quantized_type_dtype, ) from compressed_tensors.quantization.utils.mxfp_utils import ( generate_mx_scales, maybe_convert_from_mx_exp, should_generate_mx_scales, ) from loguru import logger from torch import FloatTensor, IntTensor, Tensor from torch.nn import Module __all__ = [ "is_module_quantized", "is_model_quantized", "module_type", "get_torch_bit_depth", "can_quantize", "KV_CACHE_TARGETS", "compute_dynamic_scales_and_zp", "calculate_range", "calculate_qparams", "generate_gparam", "strategy_cdiv", "calculate_block_padding", "maybe_pad_tensor_for_block_quant", ] # target the self_attn layer # QuantizedKVParameterCache is responsible for obtaining the k_scale and v_scale KV_CACHE_TARGETS = ["re:.*(self_attn|attention)$"] _LOGGER: logging.Logger = logging.getLogger(__name__) def calculate_qparams( min_vals: Tensor, max_vals: Tensor, quantization_args: QuantizationArgs, global_scale: Tensor | None = None, ) -> tuple[FloatTensor, IntTensor]: """ :param min_vals: tensor of min value(s) to calculate scale(s) and zero point(s) from :param max_vals: tensor of max value(s) to calculate scale(s) and zero point(s) from :param quantization_args: settings to quantization :param global_scale: additional global scale to scale the locally generated scale currently only applied/supported for Fp4 :return: tuple of the calculated scale(s) and zero point(s). For FP4, the calculated scale is of dtype FP8 """ # based on the implementations for consuming quantized values, # 0.0 must always be representable within the quantized range min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) device = min_vals.device bit_min, bit_max = calculate_range(quantization_args, device) bit_range = bit_max - bit_min # 1. Generate scale and zero-point if quantization_args.symmetric: max_val_pos = torch.max(torch.abs(min_vals), torch.abs(max_vals)) if should_generate_mx_scales(args=quantization_args): scales = generate_mx_scales( x=max_val_pos, num_bits=quantization_args.num_bits ) else: scales = max_val_pos / (float(bit_range) / 2) zero_points = torch.zeros(scales.shape, device=device, dtype=min_vals.dtype) else: if ( quantization_args.num_bits == 4 and quantization_args.type == QuantizationType.FLOAT ): raise NotImplementedError( "Asymmetric Quantization is not supported for FP4" ) scales = (max_vals - min_vals) / float(bit_range) zero_points = bit_min - (min_vals / scales) zero_points = torch.clamp(zero_points, bit_min, bit_max) # 2. Conditionally scale the generated local scale by a global_scale if global_scale is not None: scales = global_scale * scales # 3. Conditionally round the scale to the quantized dtype, if scale_dtype is set if quantization_args.scale_dtype is not None: scales = round_to_quantized_type_dtype( scales, dtype=quantization_args.scale_dtype ) # 4. Optionally remove exponent scales = maybe_convert_from_mx_exp(quantization_args, scales) # 5. Update any 0s with small values to # prevent div by 0 eps = _get_dtype_eps( dtype=( quantization_args.scale_dtype if quantization_args.scale_dtype is not None else scales.dtype ) ) scales = torch.where( scales == 0, torch.tensor(eps, dtype=scales.dtype, device=device), scales, ) # 6. Round the zp to zp_dtype zero_points = round_to_quantized_type_dtype( zero_points, dtype=quantization_args.zp_dtype, cast_to_original_dtype=False ) if scales.ndim == 0: scales = scales.reshape(1) zero_points = zero_points.reshape(1) return scales, zero_points def compute_dynamic_scales_and_zp( value: Tensor, args: QuantizationArgs, module: torch.nn.Module, global_scale: Tensor | None = None, ): """ Returns the computed scales and zero points for dynamic activation quantization. :param value: tensor to calculate quantization parameters for :param args: quantization args :param reduce_dims: optional tuple of dimensions to reduce along, returned scale and zero point will be shaped (1,) along the reduced dimensions :return: tuple of scale and zero point derived from the observed tensor """ keep_dims = True if args.strategy == QuantizationStrategy.TOKEN: dim = {0, 1} reduce_dims = tuple(idx for idx in range(value.ndim) if idx not in dim) elif args.strategy == QuantizationStrategy.TENSOR: reduce_dims = None elif args.strategy in ( QuantizationStrategy.TENSOR_GROUP, QuantizationStrategy.GROUP, ): reduce_dims = -1 keep_dims = False reshaped_dims = ( math.ceil(value.shape[-1] / args.group_size), args.group_size, ) value = value.unflatten(-1, reshaped_dims) else: supported_strategies = ( QuantizationStrategy.TOKEN, QuantizationStrategy.TENSOR, QuantizationStrategy.TENSOR_GROUP, QuantizationStrategy.GROUP, ) raise ValueError( "Dynamic quantization is only supported for ", f"{supported_strategies}", ) if not reduce_dims: min_val, max_val = torch.aminmax(value) else: min_val = torch.amin(value, dim=reduce_dims, keepdims=keep_dims) max_val = torch.amax(value, dim=reduce_dims, keepdims=keep_dims) return calculate_qparams(min_val, max_val, args, global_scale=global_scale) def calculate_range( quantization_args: QuantizationArgs, device: str ) -> tuple[torch.Tensor, torch.Tensor]: """ Calculated the effective quantization range for the given Quantization Args :param quantization_args: quantization args to get range of :param device: device to store the range to :return: tuple endpoints for the given quantization range """ if quantization_args.type == QuantizationType.INT: bit_range = 2.0**quantization_args.num_bits q_max = torch.tensor(bit_range / 2 - 1, device=device) q_min = torch.tensor(-bit_range / 2, device=device) elif quantization_args.type == QuantizationType.FLOAT: if quantization_args.num_bits == 8: q_max = torch.tensor(FP8_E4M3_DATA.max, device=device) q_min = torch.tensor(FP8_E4M3_DATA.min, device=device) elif quantization_args.num_bits == 4: q_max = torch.tensor(FP4_E2M1_DATA.max, device=device) q_min = torch.tensor(FP4_E2M1_DATA.min, device=device) else: raise NotImplementedError( "Range calculation only supported for 4 and 8 bits" ) else: raise ValueError(f"Invalid quantization type {quantization_args.type}") return q_min, q_max def is_module_quantized(module: Module) -> bool: """ Check if a module is quantized, based on the existence of a non-empty quantization scheme :param module: pytorch module to check :return: True if module is quantized, False otherwise """ if not hasattr(module, "quantization_scheme"): return False if module.quantization_scheme.weights is not None: return True if module.quantization_scheme.input_activations is not None: return True if module.quantization_scheme.output_activations is not None: return True return False def is_model_quantized(model: Module) -> bool: """ Check if any modules in a model are quantized, based on the existence of a non-empty quantization scheme in at least one module :param model: pytorch model :return: True if model is quantized, False otherwise """ return any(is_module_quantized(submodule) for submodule in model.modules()) def module_type(module: Module) -> str: """ Gets a string representation of a module type :module: pytorch module to get type of :return: module type as a string """ return type(module).__name__ def get_torch_bit_depth(value: torch.Tensor) -> int: """ Determine the number of bits used to represent the dtype of a tensor :param value: tensor to check bit depth of :return: bit depth of each element in the value tensor """ try: bit_depth = torch.finfo(value.dtype).bits except TypeError: bit_depth = torch.iinfo(value.dtype).bits return bit_depth def can_quantize(value: torch.Tensor, quant_args: "QuantizationArgs") -> bool: # noqa """ Checks if value can be quantized by quant_args. :param value: tensor to check for quantization :param quant_args: QuantizationArgs to use for quantization :return: False if value is already quantized to quant_args or value is incompatible with quant_args, True if value can be quantized with quant_args """ bit_depth = get_torch_bit_depth(value) requested_depth = quant_args.num_bits if bit_depth < quant_args.num_bits: _LOGGER.warn( f"Can't quantize tensor with bit depth {bit_depth} to {requested_depth}." "The QuantizationArgs provided are not compatible with the input tensor." ) return bit_depth > quant_args.num_bits def generate_gparam( updated_min_val: torch.Tensor, updated_max_val: torch.Tensor, scale_data: FloatArgs | None = FP8_E4M3_DATA, quant_data: FloatArgs | None = FP4_E2M1_DATA, dtype: torch.dtype | None = torch.float32, ): """ Generate a global scale for an entire tensor (input_tensor). Goal of the scale is to ensure that the quantization (local) scale falls into the approproiate dtype range. E.g. for NVFP4, group (local) scales are in dtype FP8. The global_scale attempts to use the entire FP8 dtype range while mapping a per-group max to the FP4 max. """ min_vals = torch.min(updated_min_val, torch.zeros_like(updated_min_val)) max_vals = torch.max(updated_max_val, torch.zeros_like(updated_max_val)) max_val_pos = torch.max(torch.abs(min_vals), torch.abs(max_vals)) max_val_pos = torch.clamp(max_val_pos, min=torch.finfo(max_val_pos.dtype).tiny) global_scale = scale_data.max * quant_data.max / max_val_pos # Replace any NaN or Inf with 1.0. NaN arises when max_val_pos was NaN # (clamp does not propagate NaN, so it passes through to the division). # Inf arises when max_val_pos was near zero and the division overflows float32. # global_scale=1 means no global scaling, which is a safe fallback for # uncalibrated experts. global_scale = torch.nan_to_num(global_scale, nan=1.0, posinf=1.0, neginf=1.0) return global_scale.to(dtype).reshape([1]) def strategy_cdiv( value: int, divisor: int, strategy: QuantizationStrategy | None, strict: bool = False, ) -> int: dividend = math.ceil(value / divisor) if dividend * divisor != value: message = ( f"{strategy} quantization strategy requires strict division of " f"weight/activation size {value} and group/block size {divisor}. " "consider reducing the group/block size or ignoring modules with " f"weights not divisible by {divisor}" ) if strict: raise ValueError(message) else: logger.bind(log_once=True).warning(message) return dividend def _get_dtype_eps(dtype: torch.dtype) -> float: if dtype == FP8_E4M3_DATA.dtype: return 0.125 elif dtype == FP4_E2M1_DATA.dtype: return 0.25 elif torch.is_floating_point(torch.tensor([], dtype=dtype)): return torch.finfo(dtype).eps else: return 1 def calculate_block_padding( shape: tuple[int, ...], block_structure: tuple[int, int], ) -> tuple[int, int]: """ Calculate the padding needed to make tensor dimensions divisible by block size. For block quantization, dimensions must be divisible by the block size for proper scale alignment when layers are merged in inference frameworks like vLLM. :param shape: shape of the tensor (at least 2D) :param block_structure: [block_height, block_width] for block quantization :return: tuple of (pad_rows, pad_cols) needed to make dimensions divisible """ if len(shape) < 2: raise ValueError(f"Tensor must be at least 2D, got shape {shape}") rows, cols = shape[-2], shape[-1] block_height, block_width = block_structure pad_rows = (block_height - rows % block_height) % block_height pad_cols = (block_width - cols % block_width) % block_width return pad_rows, pad_cols def maybe_pad_tensor_for_block_quant( tensor: torch.Tensor, block_structure: tuple[int, int], ) -> torch.Tensor: """ Pad a tensor so its dimensions are divisible by the block size. This is essential for FP8 block quantization when dimensions are not divisible by block size. The padding ensures that when weights are merged in inference frameworks (like vLLM's gate_up_proj), the scale tensor blocks align correctly. :param tensor: tensor to pad (at least 2D) :param block_structure: [block_height, block_width] for block quantization :return: padded tensor """ original_shape = tensor.shape pad_rows, pad_cols = calculate_block_padding(original_shape, block_structure) if pad_rows == 0 and pad_cols == 0: return tensor # F.pad uses (left, right, top, bottom) for last two dimensions padded_tensor = torch.nn.functional.pad( tensor, (0, pad_cols, 0, pad_rows), mode="constant", value=0 ) return padded_tensor vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/quantization/utils/mxfp_utils.py000066400000000000000000000117131521257237700336320ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math import torch from compressed_tensors.quantization.quant_args import ( BFLOAT16_DATA, FLOAT16_DATA, FLOAT32_DATA, FLOAT64_DATA, FP4_E2M1_DATA, FP8_E4M3_DATA, QuantizationArgs, QuantizationType, ) __all__ = [ "maybe_convert_from_mx_exp", "generate_mx_scales", "round_to_power_2", "should_generate_mx_scales", ] # Reference: https://github.com/vllm-project/vllm/blob/main/tests/quantization/reference_mxfp4.py # noqa: E501 # The exponent offset maps the group max into the quantized type's # representable range. It equals floor(log2(type_max)): # FP4 E2M1 max=6.0 -> floor(log2(6)) = 2 # FP8 E4M3 max=448.0 -> floor(log2(448)) = 8 _MX_ELEM_OFFSET = { 4: int(math.floor(math.log2(FP4_E2M1_DATA.max))), # 2 8: int(math.floor(math.log2(FP8_E4M3_DATA.max))), # 8 } def should_generate_mx_scales(args: QuantizationArgs): return ( args.num_bits in (4, 8) and args.type == QuantizationType.FLOAT.value and args.group_size == 32 and args.scale_dtype == torch.uint8 ) def maybe_convert_from_mx_exp( args: QuantizationArgs, scale: torch.Tensor ) -> torch.Tensor: """ Conditionally converts MX (MXFP4/MXFP8) scales from their E8M0 exponent format to float scales. If the quantization arguments indicate an MX format, the input `scale` is treated as E8M0 uint8 exponents and converted to float power-of-2 scales. Otherwise, the input `scale` tensor is returned unchanged. :param args: quantization args to check for MX format :param scale: tensor of scale values (uint8 exponents for MX, or float) :return: float scale tensor, or original scale if not MX format """ original_dtype = scale.dtype if should_generate_mx_scales(args): scale_exp = scale.to(torch.int32) - 127 scale = 2.00 ** (scale_exp.to(torch.float)) return scale.to(original_dtype) return scale def round_to_power_2(x: torch.Tensor) -> torch.Tensor: """ Round values to the closest power of 2. This is done by masking the values with SIGN_EXPONENT_MASK, which essentially removes the mantissa and keeps the exponent. i.e the closest power of 2 for the input_value. E.g: 0.0825 = 1.32 (mantissa) x 2**-4 (exponent) 0.0825 ==> -4 (exponent) + 127 = 123 = 01111011 (8 bits for bfloat16) 0.0825 ==> 0.32 (mantissa) = 0101001 (7 bits for bfloat16) 0.0825 == 0b01111011_0101001 (bfloat16) 0b01111011_0101001 & 111111111_0000000 == 0b01111011_0000000 Keep the exponent + sign bit to give you the closest power of 2, 0.0625 :param x: tensor to round to closest power of 2 """ scale_dtype = x.dtype if scale_dtype is torch.bfloat16: int_dtype = torch.uint16 mantissa = BFLOAT16_DATA.mantissa exponent = BFLOAT16_DATA.exponent elif scale_dtype is torch.float16: int_dtype = torch.uint16 mantissa = FLOAT16_DATA.mantissa exponent = FLOAT16_DATA.exponent elif scale_dtype is torch.float32: int_dtype = torch.uint32 mantissa = FLOAT32_DATA.mantissa exponent = FLOAT32_DATA.exponent elif scale_dtype is torch.float64: int_dtype = torch.uint64 mantissa = FLOAT64_DATA.mantissa exponent = FLOAT64_DATA.exponent else: raise TypeError(f"Unsupported dtype {scale_dtype}") if int_dtype in (torch.uint16, torch.uint32): x = x.view(int_dtype).to(torch.int32) else: x = x.view(int_dtype).to(torch.int64) # Find closest power of 2 VAL_TO_ADD = 1 << (mantissa - FP4_E2M1_DATA.mantissa - 1) # Add value to push the value to the next exponent SIGN_EXPONENT_MASK = ((1 << (exponent + 1)) - 1) << mantissa # mask to only keep exponent - we conservatively round down # to better represent smaller numbers / prevent overflow block_max_uint = torch.bitwise_and(x + VAL_TO_ADD, SIGN_EXPONENT_MASK) if int_dtype is torch.uint16: return block_max_uint.to(int_dtype).view(scale_dtype) return block_max_uint.view(scale_dtype) def generate_mx_scales(x: torch.Tensor, num_bits: int = 4) -> torch.Tensor: """ Generate MX scales (for MXFP4 and MXFP8). The scales require the following steps: 1. Round to the closest power of 2 2. Subtract the element-format offset so that the largest group values map into the quantized type's representable range 3. Convert to biased E8M0 exponent (bias 127) Called when calculating qparams using observers. :param x: tensor of per-group max absolute values :param num_bits: quantized element width (4 for MXFP4, 8 for MXFP8) :returns scales as E8M0 exponents (uint8 after rounding) """ offset = _MX_ELEM_OFFSET[num_bits] # Round to closest power of 2 scale_power_2 = round_to_power_2(x) return 127 + torch.floor(torch.log2(scale_power_2)) - offset vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/registry/000077500000000000000000000000001521257237700270455ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/registry/__init__.py000066400000000000000000000002241521257237700311540ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .registry import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/registry/registry.py000066400000000000000000000260531521257237700312750ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ Universal registry to support registration and loading of child classes and plugins """ import importlib from collections import defaultdict from typing import Any, TypeVar __all__ = [ "RegistryMixin", "register", "get_from_registry", "registered_names", "registered_aliases", "standardize_lookup_name", ] _ALIAS_REGISTRY: dict[type, dict[str, str]] = defaultdict(dict) _REGISTRY: dict[type, dict[str, Any]] = defaultdict(dict) T = TypeVar("", bound="RegistryMixin") def standardize_lookup_name(name: str) -> str: """ Standardize the given name for lookup in the registry. This will replace all underscores and spaces with hyphens and convert the name to lowercase. example: ``` standardize_lookup_name("Foo_bar baz") == "foo-bar-baz" ``` :param name: name to standardize :return: standardized name """ return name.replace("_", "-").replace(" ", "-").lower() def standardize_alias_name( name: str | list[str] | None, ) -> str | list[str] | None: if name is None: return None elif isinstance(name, str): return standardize_lookup_name(name) else: # isinstance(name, list) return [standardize_lookup_name(n) for n in name] class RegistryMixin: """ Universal registry to support registration and loading of child classes and plugins Classes that require a registry or plugins may add the `RegistryMixin` and use `register` and `load` as the main entrypoints for adding new implementations and loading requested values from its registry. If a class should only have its child classes in its registry, the class should set the static attribute `registry_requires_subclass` to True example ```python class Dataset(RegistryMixin): pass # register with default name @Dataset.register() class ImageNetDataset(Dataset): pass # load as "ImageNetDataset" imagenet = Dataset.load("ImageNetDataset") # register with custom name @Dataset.register(name="cifar-dataset") class Cifar(Dataset): pass Note: the name will be standardized for lookup in the registry. For example, if a class is registered as "cifar_dataset" or "cifar dataset", it will be stored as "cifar-dataset". The user will be able to load the class with any of the three name variants. # register with multiple aliases @Dataset.register(alias=["cifar-10-dataset", "cifar_100_dataset"]) class Cifar(Dataset): pass # load as "cifar-dataset" cifar = Dataset.load_from_registry("cifar-dataset") # load from custom file that implements a dataset mnist = Dataset.load_from_registry("/path/to/mnnist_dataset.py:MnistDataset") ``` """ # set to True in child class to add check that registered/retrieved values # implement the class it is registered to registry_requires_subclass: bool = False @classmethod def register(cls, name: str | None = None, alias: list[str] | str | None = None): """ Decorator for registering a value (ie class or function) wrapped by this decorator to the base class (class that .register is called from) :param name: name or list of names to register the wrapped value as, defaults to value.__name__ :param alias: alias or list of aliases to register the wrapped value as, defaults to None :return: register decorator """ def decorator(value: Any): cls.register_value(value, name=name, alias=alias) return value return decorator @classmethod def register_value( cls, value: Any, name: str, alias: str | list[str] | None = None ): """ Registers the given value to the class `.register_value` is called from :param value: value to register :param name: name to register the wrapped value as, defaults to value.__name__ :param alias: alias or list of aliases to register the wrapped value as, defaults to None """ register( parent_class=cls, value=value, name=name, alias=alias, require_subclass=cls.registry_requires_subclass, ) @classmethod def load_from_registry(cls: type[T], name: str, **constructor_kwargs) -> T: """ :param name: name of registered class to load :param constructor_kwargs: arguments to pass to the constructor retrieved from the registry :return: loaded object registered to this class under the given name, constructed with the given kwargs. Raises error if the name is not found in the registry """ constructor = cls.get_value_from_registry(name=name) return constructor(**constructor_kwargs) @classmethod def get_value_from_registry(cls: type[T], name: str) -> T: """ :param name: name to retrieve from the registry :return: value from retrieved the registry for the given name, raises error if not found """ return get_from_registry( parent_class=cls, name=name, require_subclass=cls.registry_requires_subclass, ) @classmethod def registered_names(cls) -> list[str]: """ :return: list of all names registered to this class """ return registered_names(cls) @classmethod def registered_aliases(cls) -> list[str]: """ :return: list of all aliases registered to this class """ return registered_aliases(cls) def register( parent_class: type, value: Any, name: str | None = None, alias: list[str] | str | None = None, require_subclass: bool = False, ): """ :param parent_class: class to register the name under :param value: the value to register :param name: name to register the wrapped value as, defaults to value.__name__ :param alias: alias or list of aliases to register the wrapped value as, defaults to None :param require_subclass: require that value is a subclass of the class this method is called from """ if name is None: # default name name = value.__name__ name = standardize_lookup_name(name) alias = standardize_alias_name(alias) register_alias(name=name, alias=alias, parent_class=parent_class) if require_subclass: _validate_subclass(parent_class, value) if name in _REGISTRY[parent_class]: # name already exists - raise error if two different values are attempting # to share the same name registered_value = _REGISTRY[parent_class][name] if registered_value is not value: raise RuntimeError( f"Attempting to register name {name} as {value} " f"however {name} has already been registered as {registered_value}" ) else: _REGISTRY[parent_class][name] = value def get_from_registry( parent_class: type, name: str, require_subclass: bool = False ) -> Any: """ :param parent_class: class that the name is registered under :param name: name to retrieve from the registry of the class :param require_subclass: require that value is a subclass of the class this method is called from :return: value from retrieved the registry for the given name, raises error if not found """ name = standardize_lookup_name(name) if ":" in name: # user specifying specific module to load and value to import module_path, value_name = name.split(":") retrieved_value = _import_and_get_value_from_module(module_path, value_name) else: # look up name in alias registry name = _ALIAS_REGISTRY[parent_class].get(name, name) # look up name in registry retrieved_value = _REGISTRY[parent_class].get(name) if retrieved_value is None: raise KeyError( f"Unable to find {name} registered under type {parent_class}.\n" f"Registered values for {parent_class}: " f"{registered_names(parent_class)}\n" f"Registered aliases for {parent_class}: " f"{registered_aliases(parent_class)}" ) if require_subclass: _validate_subclass(parent_class, retrieved_value) return retrieved_value def registered_names(parent_class: type) -> list[str]: """ :param parent_class: class to look up the registry of :return: all names registered to the given class """ return list(_REGISTRY[parent_class].keys()) def registered_aliases(parent_class: type) -> list[str]: """ :param parent_class: class to look up the registry of :return: all aliases registered to the given class """ registered_aliases_plus_names = list(_ALIAS_REGISTRY[parent_class].keys()) registered_aliases = list( set(registered_aliases_plus_names) - set(registered_names(parent_class)) ) return registered_aliases def register_alias(name: str, parent_class: type, alias: str | list[str] | None = None): """ Updates the mapping from the alias(es) to the given name. If the alias is None, the name is used as the alias. ``` :param name: name that the alias refers to :param parent_class: class that the name is registered under :param alias: single alias or list of aliases that refer to the name, defaults to None """ if alias is not None: alias = alias if isinstance(alias, list) else [alias] else: alias = [] if name in alias: raise KeyError( f"Attempting to register alias {name}, " f"that is identical to the standardized name: {name}." ) alias.append(name) for alias_name in alias: if alias_name in _ALIAS_REGISTRY[parent_class]: raise KeyError( f"Attempting to register alias {alias_name} as {name} " f"however {alias_name} has already been registered as " f"{_ALIAS_REGISTRY[alias_name]}" ) _ALIAS_REGISTRY[parent_class][alias_name] = name def _import_and_get_value_from_module(module_path: str, value_name: str) -> Any: # import the given module path and try to get the value_name if it is included # in the module # load module spec = importlib.util.spec_from_file_location( f"plugin_module_for_{value_name}", module_path ) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) # get value from module value = getattr(module, value_name, None) if not value: raise RuntimeError( f"Unable to find attribute {value_name} in module {module_path}" ) return value def _validate_subclass(parent_class: type, child_class: type): if not issubclass(child_class, parent_class): raise ValueError( f"class {child_class} is not a subclass of the class it is " f"registered for: {parent_class}." ) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/000077500000000000000000000000001521257237700272105ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/__init__.py000066400000000000000000000006151521257237700313230ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa # isort: skip_file from .transform_args import * from .transform_scheme import * from .transform_config import * from .factory.base import * from .factory.hadamard import * from .factory.matrix_multiply import * from .factory.random_hadamard import * from .apply import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/apply.py000066400000000000000000000044631521257237700307160ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors import TRANSFORM_CONFIG_NAME from compressed_tensors.transform import TransformConfig, TransformFactory from compressed_tensors.transform.factory.base import TransformBase __all__ = ["apply_transform_config"] def apply_transform_config(model: torch.nn.Module, config: TransformConfig): """ Apply a transform config to a model. Weight transforms are fused into weights, while activation transforms are attached as submodules and trigger via pytorch hooks :param model: model to apply config to :param config: transform config to apply """ for name, scheme in config.config_groups.items(): factory = TransformFactory.from_scheme(scheme, name=name) factory.apply_to_model(model) # declare shared transform parameters as tied weights for save_pretrained compat _register_tied_transform_weights(model) # attach config to model for compression/serialization setattr(model, TRANSFORM_CONFIG_NAME, config) def _register_tied_transform_weights(model: torch.nn.Module): """ Scan for transform submodules that share parameters and register them as tied weights via ``_tied_weights_keys``. This allows ``save_pretrained`` in transformers v5+ to handle shared tensors without raising an error. """ # Map parameter id -> first full parameter name that owns it first_seen: dict[int, str] = {} for module_name, module in model.named_modules(): if not isinstance(module, TransformBase): continue tied_keys: dict[str, str] = {} for key in getattr(module, "_dynamic_tied_weights_keys", []): param = getattr(module, key, None) if param is None: continue param_id = id(param) full_key = f"{module_name}.{key}" if module_name else key if param_id not in first_seen: first_seen[param_id] = full_key else: # This parameter was already registered under a different module; # mark it as tied so save_pretrained knows to deduplicate tied_keys[key] = first_seen[param_id] if tied_keys: module._tied_weights_keys = tied_keys vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/factory/000077500000000000000000000000001521257237700306575ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/factory/__init__.py000066400000000000000000000001531521257237700327670ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/factory/base.py000066400000000000000000000163711521257237700321530ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod import torch import torch.nn.utils.parametrize as P import tqdm from compressed_tensors.modeling.attention import ( initialize_hooked_attention, register_query_hook, ) from compressed_tensors.modeling.kvcache import ( initialize_hooked_kv_cache, register_key_hook, ) from compressed_tensors.offload import OffloadCache from compressed_tensors.registry.registry import RegistryMixin, T from compressed_tensors.transform import ( TransformArgs, TransformLocation, TransformScheme, ) from compressed_tensors.utils import ( align_module_device, match_named_modules, patch_attr, register_offload_module, update_offload_parameter, ) from compressed_tensors.utils.internal import InternalModule from torch import Tensor from torch.nn import Module, Parameter from transformers import PreTrainedModel __all__ = ["TransformFactory", "TransformBase"] class TransformFactory(RegistryMixin, ABC): """ Abstract factory base used to create and apply transforms to a model :param name: name associated with transform scheme :param scheme: transform scheme which defines how transforms should be created :param seed: random seed used to transform weight randomization """ transforms: list["TransformBase"] def __init__(self, name: str, scheme: TransformScheme, seed: int | None = None): self.name = name self.scheme = scheme self.generator = torch.Generator() if seed is not None: self.generator.manual_seed(seed) @classmethod def from_scheme(cls: type[T], scheme: TransformScheme, **kwargs) -> T: """ Create a transform factory from a scheme :param scheme: defines how transforms should be created :param kwargs: TransformFactory constructor arguments :return: subclass of `TransformFactory` corresponding to the scheme type """ constructor = cls.get_value_from_registry(name=scheme.type) return constructor(scheme=scheme, **kwargs) @abstractmethod def create_transform(self, module: Module, args: TransformArgs) -> "TransformBase": """ Abstract method which defines how a transform should be created. May utilize caching to maximize shared memory :param module: parent module that transform will be applied to :param args: defines how the transform will be applied to the module :return: instance of TransformBase """ raise NotImplementedError() def apply_to_model(self, model: Module, use_tqdm=True): """ Create transforms and apply them to the model :param model: module to apply transforms to """ modules_args = [ (module, arg) for arg in self.scheme.apply for _, module in match_named_modules(model, arg.targets, arg.ignore) ] desc = f"Applying {self.name} transforms" for module, arg in tqdm.tqdm(modules_args, desc=desc, disable=(not use_tqdm)): self._apply_to_module(model, module, arg) def _apply_to_module(self, model: Module, module: Module, args: TransformArgs): """ Create transforms and apply them to the module :param model: model which module belongs to :param module: target module to apply transforms to :param args: defines how the transform will be applied to the target module """ # create transform as submodule transform_name = f"{self.name}_{args.location}" transform = self.create_transform(module, args) register_offload_module(module, transform_name, transform) # register input transformation hook if args.location == TransformLocation.INPUT: def input_hook(_, args): input = args[0] return transform(input) module.register_forward_pre_hook(input_hook, prepend=True) # eagerly apply transformation to weight elif args.location in ( TransformLocation.WEIGHT_INPUT, TransformLocation.WEIGHT_OUTPUT, ): # fuse transform into weight assert hasattr(module, "weight") with torch.no_grad(), align_module_device(module): update_offload_parameter(module, "weight", transform(module.weight)) # For WEIGHT_OUTPUT, the bias must also be transformed: # y' = R @ (W @ x + b) = (R @ W) @ x + R @ b # Without this, models with bias (e.g. Qwen2 attention) # produce incorrect outputs under head-wise rotations (R2). if ( args.location == TransformLocation.WEIGHT_OUTPUT and getattr(module, "bias", None) is not None ): new_bias = transform(module.bias.unsqueeze(-1)).squeeze(-1) update_offload_parameter(module, "bias", new_bias) if self.scheme.requires_grad: # for training, the weight changes with every forward pass # so we can leverage parametrization to propagate the gradient if isinstance(module._parameters, OffloadCache): raise ValueError("Offloaded training is not supported") P.register_parametrization(module, "weight", transform) else: # transform is no longer needed (unfusing is not supported) delattr(module, transform_name) # register output transformation hook elif args.location == TransformLocation.OUTPUT: def output_hook(_, _input, output): return transform(output) module.register_forward_hook(output_hook) # register query hook to attention elif args.location == TransformLocation.Q_ATTN: if not isinstance(model, PreTrainedModel): raise ValueError(f"Cannot hook attention of model: {model}") def query_hook(_, query_states): return transform(query_states) initialize_hooked_attention(model, module) register_query_hook(module, query_hook) # register key hook to kvcache elif args.location == TransformLocation.K_CACHE: if not isinstance(model, PreTrainedModel): raise ValueError(f"Cannot hook attention of model: {model}") def key_hook(_, key_states): return transform(key_states) initialize_hooked_kv_cache(model, module) register_key_hook(module, key_hook) else: raise NotImplementedError() class TransformBase(InternalModule, ABC): """ Represents the application of a transform accord to TransformArgs """ args: TransformArgs weight: Parameter _dynamic_tied_weights_keys: list[str] = ["weight"] @abstractmethod def forward(self, value: Tensor) -> Tensor: raise NotImplementedError() def right_inverse(self, value: Tensor) -> Tensor: with patch_attr(self.args, "inverse", not self.args.inverse): return self.forward(value) def __repr__(self): return f"{self.__class__.__name__}(inverse={self.args.inverse})" vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/factory/hadamard.py000066400000000000000000000076741521257237700330100ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.transform import TransformArgs, TransformScheme from compressed_tensors.transform.factory.base import TransformBase, TransformFactory from compressed_tensors.transform.utils.hadamard import deterministic_hadamard_matrix from compressed_tensors.transform.utils.matrix import ( apply_transform_weight, get_transform_size, ) from compressed_tensors.utils import get_execution_device, get_offloaded_device from compressed_tensors.utils.helpers import ParameterizedDefaultDict from torch import Tensor, device, dtype from torch.nn import Module, Parameter @TransformFactory.register("hadamard") class HadamardFactory(TransformFactory): """ Factory used to apply hadamard transforms to a model :param name: name associated with transform scheme :param scheme: transform scheme which defines how transforms should be created :param seed: random seed used to transform weight randomization """ def __init__(self, name: str, scheme: TransformScheme, seed: int | None = None): super().__init__(name, scheme, seed) self.weights = ParameterizedDefaultDict(self._create_weight) self.perms = ParameterizedDefaultDict(self._create_permutation) def create_transform(self, module: Module, args: TransformArgs): """ Create a HadamardTransform for applying to a module. Transforms with the same size, dtype, and device are cached :param module: parent module that transform will be applied to :param args: defines how the transform will be applied to the module """ size = get_transform_size(module, args.location, self.scheme.head_dim) exec_device = get_execution_device(module) device = get_offloaded_device(module) precision = self.scheme.precision if args.is_online() else torch.float64 factory_kwargs = { "device": device, "construct_device": exec_device, "precision": precision, } weight = self.weights.get(size, factory_kwargs=factory_kwargs) # TODO: permutations should be keyed by fused modules, not weight perm = self.perms[weight] if self.scheme.randomize else None return HadamardTransform(weight, perm, self.scheme, args, type(module)) def _create_weight( self, size: int, device: device, construct_device: device, precision: dtype, ) -> Parameter: data = deterministic_hadamard_matrix(size, precision, construct_device) data = data.to(device=device) return Parameter(data, requires_grad=self.scheme.requires_grad) def _create_permutation(self, weight: Parameter) -> Parameter: data = torch.randperm(weight.size(0), generator=self.generator) return Parameter(data, requires_grad=False) class HadamardTransform(TransformBase): _dynamic_tied_weights_keys: list[str] = ["weight", "perm"] def __init__( self, weight: Parameter, perm: Parameter | None, scheme: TransformScheme, args: TransformArgs, module_type: type[torch.nn.Module], ): super().__init__() self.weight = weight self.perm = perm self.scheme = scheme self.args = args self.module_type = module_type self._scale = torch.tensor(weight.size(0), dtype=torch.float64).sqrt() def forward(self, value: Tensor) -> Tensor: weight = self.weight if self.perm is not None: weight = weight[self.perm][:, self.perm] if self.args.inverse: weight = weight.T return ( apply_transform_weight( weight.to(device=value.device), value.to(dtype=weight.dtype), self.args.location, self.module_type, ) / self._scale ).to(value.dtype) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/factory/matrix_multiply.py000066400000000000000000000110671521257237700345010ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.transform import TransformArgs, TransformScheme from compressed_tensors.transform.factory.base import TransformBase, TransformFactory from compressed_tensors.transform.utils.matrix import ( apply_transform_weight, get_transform_size, ) from compressed_tensors.utils.helpers import ParameterizedDefaultDict from compressed_tensors.utils.offload import get_offloaded_device from torch import Tensor, device, dtype from torch.nn import Module, Parameter @TransformFactory.register("random-matrix") class RandomMatrixFactory(TransformFactory): """ Factory used to apply random matrix transforms to a model :param name: name associated with transform scheme :param scheme: transform scheme which defines how transforms should be created :param seed: random seed used to transform weight randomization """ def __init__(self, name: str, scheme: TransformScheme, seed: int | None = None): super().__init__(name, scheme, seed) self.weights = ParameterizedDefaultDict(self._create_weight) self.inverses = ParameterizedDefaultDict(self._create_inverse) def create_transform(self, module: Module, args: TransformArgs): """ Create a RandomMatrixTransform for applying to a module. Transforms with the same size, dtype, and device are cached :param module: parent module that transform will be applied to :param args: defines how the transform will be applied to the module """ size = get_transform_size(module, args.location, self.scheme.head_dim) device = get_offloaded_device(module) precision = self.scheme.precision if args.is_online() else torch.float64 factory_kwargs = {"device": device, "precision": precision} weight = self.weights.get(size, factory_kwargs=factory_kwargs) if args.inverse: weight = self.inverses[weight] return RandomMatrixTransform(weight, self.scheme, args, type(module)) def _create_weight(self, size: int, device: device, precision: dtype) -> Parameter: # TODO: construct such that weight is invertible (has non-zero determinant) data = torch.rand( (size, size), generator=self.generator, dtype=precision, device=self.generator.device, ).to(device) return Parameter(data, requires_grad=self.scheme.requires_grad) def _create_inverse(self, weight: Parameter) -> Parameter: data = high_precision_invert(weight.data) data = data.contiguous() # ensure proper serialization return Parameter(data, requires_grad=False) class RandomMatrixTransform(TransformBase): def __init__( self, weight: Tensor, scheme: TransformScheme, args: TransformArgs, module_type: type[torch.nn.Module], ): super().__init__() self.weight = weight # is an inverse if args.inverse self.scheme = scheme self.args = args self.module_type = module_type def forward(self, value: Tensor) -> Parameter: return apply_transform_weight( self.weight.to(device=value.device), value.to(dtype=self.weight.dtype), self.args.location, self.module_type, ).to(value.dtype) def right_inverse(self, value: Tensor) -> Tensor: inverse = high_precision_invert(self.weight) return apply_transform_weight( inverse.to(device=value.device), value.to(dtype=inverse.dtype), self.args.location, self.module_type, ).to(value.dtype) def _has_cpu_lapack() -> bool: try: torch.linalg.inv(torch.eye(2, dtype=torch.float64)) return True except RuntimeError: return False _cpu_lapack_available: bool | None = None def high_precision_invert(weight: Tensor) -> Tensor: global _cpu_lapack_available original_device = weight.device compute_device = original_device # If the tensor is on CPU and LAPACK is not available (e.g. ROCm builds), # move to GPU for the inversion if compute_device.type == "cpu": if _cpu_lapack_available is None: _cpu_lapack_available = _has_cpu_lapack() if not _cpu_lapack_available and torch.cuda.is_available(): compute_device = torch.device("cuda") result = torch.linalg.inv(weight.to(device=compute_device, dtype=torch.float64)) return result.to(device=original_device, dtype=weight.dtype) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/factory/random_hadamard.py000066400000000000000000000020531521257237700343320ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from compressed_tensors.transform import HadamardFactory, TransformFactory from compressed_tensors.transform.utils.hadamard import random_hadamard_matrix from torch import device, dtype from torch.nn import Parameter @TransformFactory.register("random-hadamard") class RandomHadamardFactory(HadamardFactory): """ Factory used to apply random hadamard transforms to a model :param name: name associated with transform scheme :param scheme: transform scheme which defines how transforms should be created :param seed: random seed used to transform weight randomization """ def _create_weight( self, size: int, device: device, construct_device: device, precision: dtype, ) -> Parameter: data = random_hadamard_matrix(size, precision, construct_device, self.generator) data = data.to(device=device) return Parameter(data, requires_grad=self.scheme.requires_grad) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/transform_args.py000066400000000000000000000060431521257237700326140ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum from pydantic import BaseModel, ConfigDict, Field, field_validator __all__ = ["TransformArgs", "TransformLocation"] class TransformLocation(str, Enum): """ Enum representing which parameters/activations a transform weight should be applied to on a given module. | -------------------------------------------------------------------------------------------------------- | # noqa: E501 | Name | Runtime | Values | Locations Where Inverse Could Be Applied | # noqa: E501 | --------------- | ----------- | ------------- | -------------------------------------------------------- | # noqa: E501 | `INPUT` | online | activations | `prev.WEIGHT_OUTPUT`, `prev.OUTPUT`, `this.WEIGHT_INPUT` | # noqa: E501 | `WEIGHT_INPUT` | offline | weight | `prev.WEIGHT_OUTPUT`, `prev.OUTPUT`, `this.INPUT` | # noqa: E501 | `WEIGHT_OUTPUT` | offline | weight | `this.OUTPUT`, `next.INPUT`, `next.WEIGHT_INPUT` | # noqa: E501 | `OUTPUT` | online | activations | `this.WEIGHT_OUTPUT`, `next.INPUT`, `next.WEIGHT_INPUT` | # noqa: E501 | `K_CACHE` | online | key_values | `q_proj.Q_ATTN` | # noqa: E501 | `Q_ATTN` | online | query_values | `k_proj.K_CACHE` | # noqa: E501 | -------------------------------------------------------------------------------------------------------- | # noqa: E501 """ INPUT = "input" WEIGHT_INPUT = "weight_input" WEIGHT_OUTPUT = "weight_output" OUTPUT = "output" K_CACHE = "k_cache" Q_ATTN = "q_attn" def is_online(self) -> bool: """ Returns True if the transform location is online (applied at runtime), False otherwise """ return self not in ( TransformLocation.WEIGHT_INPUT, TransformLocation.WEIGHT_OUTPUT, ) class TransformArgs(BaseModel, use_enum_values=True): """ Arguments which define *how* and where a transform should be applied to a model :param targets: list of modules to apply transforms to :param location: where to apply transform on module, one of (`input`, `weight`, `output`, `k_cache`, `q_attn`) :param inverse: whether or not to apply the inverse of a transform :param ignore: any modules which should be ignored from the targets list """ targets: list[str] location: TransformLocation inverse: bool = Field(default=False) ignore: list[str] = Field(default_factory=list) @field_validator("targets", "ignore", mode="before") @classmethod def wrap_singleton(cls, value): if isinstance(value, str): return [value] return value def is_online(self) -> bool: return TransformLocation(self.location).is_online() model_config = ConfigDict(extra="forbid") vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/transform_config.py000066400000000000000000000012351521257237700331230ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from compressed_tensors.transform import TransformScheme from pydantic import BaseModel, ConfigDict __all__ = ["TransformConfig"] class TransformConfig(BaseModel): """ Configuration of transforms to be applied to a model. This config is to be serialized within a model's `config.json` file :param config_groups: A dictionary of `TransformSchemes` that should be applied to a particular model. The keys can be any arbitrary string """ config_groups: dict[str, TransformScheme] model_config = ConfigDict(extra="forbid") vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/transform_scheme.py000066400000000000000000000036071521257237700331270ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.transform import TransformArgs from compressed_tensors.utils import TorchDtype from pydantic import BaseModel, ConfigDict, Field __all__ = ["TransformScheme"] class TransformScheme(BaseModel): """ Scheme used to parameterize a particular transform type and specify how and where it should be applied to the model :param type: string indicating the particular transform type that should be created and applied. This should be one of the registered transform types (see `Transforms.registered_names()`) :param apply: list of TransformationArgs containing the information about the modules that should be targeted by the specified transform :param randomize: True if uniquely randomized transform weights should be used, otherwise use identical transform weights where applicable :param requires_grad: True if weights include gradients for training :param head_dim: If set, the transform matrix will be block diagonal with each block being a square matrix of this size. The name head_dim was chosen because some rotations need to be block-diagonal with block size equal to the head_dim, but research has shown value in applying some rotations with smaller block size, irrespective of head_dim. :param precision: Precision at which this transform should be applied during online rotations. Fused (offline) rotations are always performed in float64 """ type: str apply: list[TransformArgs] = Field(default_factory=list) randomize: bool = Field(default=False) requires_grad: bool = Field(default=False) head_dim: int | None = Field(default=None) precision: TorchDtype = Field(default=torch.float32) model_config = ConfigDict(extra="forbid") vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/utils/000077500000000000000000000000001521257237700303505ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/utils/__init__.py000066400000000000000000000001531521257237700324600ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/utils/hadamard.py000066400000000000000000000121271521257237700324660ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math from pathlib import Path import torch from safetensors import safe_open REPO_PATH = Path(__file__).parent / "hadamards.safetensors" __all__ = ["random_hadamard_matrix", "deterministic_hadamard_matrix", "is_pow2"] # note that hadamard matrix multiplication can be accelerated using a library such as # https://github.com/Dao-AILab/fast-hadamard-transform/tree/master def deterministic_hadamard_matrix( size: int, dtype: torch.dtype = torch.bfloat16, device: torch.device = torch.device("cpu"), ) -> torch.Tensor: """ Construct an n-by-n Hadamard matrix, using Sylvester's construction. `n` must be a power of 2. Adapated from https://github.com/scipy/scipy/blob/v1.15.2/scipy/linalg/_special_matrices.py # noqa: E501 :param size: order of the matrix, must be a power of 2 :param dtype: data type of matrix :param device: device to construct matrix on :return: hadamard matrix of size `size` """ if size <= 0: raise ValueError("Cannot construct deterministic hadamard of size <= 0") log2 = int(math.log2(size)) if size != 2**log2: raise ValueError("Cannot construct deterministic hadamard of size != 2^n") H = torch.tensor([[1]], dtype=dtype, device=device) # Sylvester's construction for _ in range(log2): H = torch.vstack((torch.hstack((H, H)), torch.hstack((H, -H)))) return H def random_hadamard_matrix( size: int, dtype: torch.dtype = torch.bfloat16, device: torch.device = torch.device("cpu"), gen: torch.Generator | None = None, ) -> torch.Tensor: """ Produces a randomly generated Hadamard matrix. Differs from `deterministic_hadamard_matrix` in that this function supports non powers of 2 and randomization using a seeded generator Adapated from https://github.com/facebookresearch/SpinQuant/blob/main/utils/hadamard_utils.py # noqa: E501 Known matrices were retrieved from N. J. A. Sloane's Library of Hadamard Matrices http://www.neilsloane.com/hadamard/ # noqa: E501 :param size: The dimension of the hamadard matrix :param dtype: data type of matrix :param device: device to construct matrix on :param gen: Optional generator random values :return: randomly generated hadamard matrix """ Q = torch.randint(low=0, high=2, size=(size,), generator=gen, dtype=dtype) # cpu Q = Q.to(device=device) Q = Q * 2 - 1 Q = torch.diag(Q) return _matmul_hadU(Q) def is_pow2(n: int) -> bool: """ Check if a number is a power of 2 :param n: number to check :return: True iff `n` is a power of 2 """ return n > 0 and (n & (n - 1) == 0) def _fetch_hadamard_divisor( n: int, dtype: torch.dtype, device: torch.device = torch.device("cpu"), file_path: str = REPO_PATH, ) -> torch.Tensor | None: """ Fetch a known hadamard matrix from the given file path. The returned matrix will be of of size `k` such that `n / k` is a power of two. Return None if no such matrix exists. Note: This function reopens the safetensors file every time it is called. This is technically inefficient, but a very small runtime cost and simpler than forcing callers to manage the file open context :param n: size of known hadamard matrix :param dtype: data type to move fetched hadamard to :param device: device to move fetched hadamard to :return: a known hadamard matrix of size `n` if one exists, else None """ open_device = torch.device("cpu") if device.type == "meta" else device with safe_open(file_path, framework="pt", device=str(open_device)) as file: divisors = sorted((int(key) for key in file.keys()), reverse=True) for divisor in divisors: if n % divisor == 0 and is_pow2(n // divisor): return file.get_tensor(str(divisor)).to(dtype=dtype, device=device) return None def _matmul_hadU(X: torch.Tensor) -> torch.Tensor: size = X.size(0) dtype = X.dtype device = X.device # Check if we have the determined hadamard matrix hadK = _fetch_hadamard_divisor(size, dtype, device=device) if hadK is None: raise ValueError(f"Cannot construct random hadamard matrix of size {size}") K = hadK.size(0) # Reshape diag matrix with randomized -1/+1 input = X.clone().view(-1, size, 1) output = input.clone() while input.shape[1] > K: input = input.view(input.shape[0], input.shape[1] // 2, 2, input.shape[2]) output = output.view(input.shape) output[:, :, 0, :] = input[:, :, 0, :] + input[:, :, 1, :] output[:, :, 1, :] = input[:, :, 0, :] - input[:, :, 1, :] output = output.view(input.shape[0], input.shape[1], -1) input, output = (output, input) assert input.shape[1] == K del output # Do not explicitly repeat - OOM # input = torch.bmm( # hadK.repeat(len(input), 1, 1).to(input.device).to(input.dtype), input) # Use bcast instead input = hadK.view(1, K, K).to(input) @ input # normalize return input.view(X.shape) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/utils/hadamards.safetensors000066400000000000000000053663451521257237700346000ustar00rootroot00000000000000{"__metadata__":{"224":"had.224.pal","256":"had.256.syl","136":"had.136.tpal","200":"had.200.pal","100":"had.100.will","244":"had.244.will","104":"had.104.pal","120":"had.120.tpal","192":"had.192.pal","40":"had.40.tpal","116":"had.116.will","132":"had.132.pal","76":"had.76.pal2","204":"had.204.pal2","60":"had.60.pal","2":"had.2","140":"had.140.pal","128":"had.128.syl","240":"had.240.pal","108":"had.108.pal","176":"had.176.ttpal","220":"had.220.pal2","196":"had.196.pal2","28":"had.28.pal2","212":"had.212.pal","20":"had.20.hall.n","36":"had.36.pal2","92":"had.92.will","164":"had.164.pal","180":"had.180.pal","148":"had.148.pal2","32":"had.32.pal","8":"had.8","48":"had.48.pal","144":"had.144.tpal","124":"had.124.pal2","232":"had.232.twill","216":"had.216.tpal","44":"had.44.pal","168":"had.168.pal","236":"had.236.od","84":"had.84.pal","156":"had.156.will","88":"had.88.tpal","96":"had.96.tpal","112":"had.112.ttpal2","160":"had.160.tpal","56":"had.56.tpal2","248":"had.248.twill","80":"had.80.pal","1":"had.1","64":"had.64.syl","72":"had.72.pal","188":"had.188.tur","228":"had.228.pal","172":"had.172.will","24":"had.24.pal","68":"had.68.pal","208":"had.208.twill","52":"had.52.will","12":"had.12","184":"had.184.twill","152":"had.152.pal","16":"had.16.0","252":"had.252.pal","4":"had.4"},"1":{"dtype":"I8","shape":[1,1],"data_offsets":[0,1]},"100":{"dtype":"I8","shape":[100,100],"data_offsets":[1,10001]},"104":{"dtype":"I8","shape":[104,104],"data_offsets":[10001,20817]},"108":{"dtype":"I8","shape":[108,108],"data_offsets":[20817,32481]},"112":{"dtype":"I8","shape":[112,112],"data_offsets":[32481,45025]},"116":{"dtype":"I8","shape":[116,116],"data_offsets":[45025,58481]},"12":{"dtype":"I8","shape":[12,12],"data_offsets":[58481,58625]},"120":{"dtype":"I8","shape":[120,120],"data_offsets":[58625,73025]},"124":{"dtype":"I8","shape":[124,124],"data_offsets":[73025,88401]},"128":{"dtype":"I8","shape":[128,128],"data_offsets":[88401,104785]},"132":{"dtype":"I8","shape":[132,132],"data_offsets":[104785,122209]},"136":{"dtype":"I8","shape":[136,136],"data_offsets":[122209,140705]},"140":{"dtype":"I8","shape":[140,140],"data_offsets":[140705,160305]},"144":{"dtype":"I8","shape":[144,144],"data_offsets":[160305,181041]},"148":{"dtype":"I8","shape":[148,148],"data_offsets":[181041,202945]},"152":{"dtype":"I8","shape":[152,152],"data_offsets":[202945,226049]},"156":{"dtype":"I8","shape":[156,156],"data_offsets":[226049,250385]},"16":{"dtype":"I8","shape":[16,16],"data_offsets":[250385,250641]},"160":{"dtype":"I8","shape":[160,160],"data_offsets":[250641,276241]},"164":{"dtype":"I8","shape":[164,164],"data_offsets":[276241,303137]},"168":{"dtype":"I8","shape":[168,168],"data_offsets":[303137,331361]},"172":{"dtype":"I8","shape":[172,172],"data_offsets":[331361,360945]},"176":{"dtype":"I8","shape":[176,176],"data_offsets":[360945,391921]},"180":{"dtype":"I8","shape":[180,180],"data_offsets":[391921,424321]},"184":{"dtype":"I8","shape":[184,184],"data_offsets":[424321,458177]},"188":{"dtype":"I8","shape":[188,188],"data_offsets":[458177,493521]},"192":{"dtype":"I8","shape":[192,192],"data_offsets":[493521,530385]},"196":{"dtype":"I8","shape":[196,196],"data_offsets":[530385,568801]},"2":{"dtype":"I8","shape":[2,2],"data_offsets":[568801,568805]},"20":{"dtype":"I8","shape":[20,20],"data_offsets":[568805,569205]},"200":{"dtype":"I8","shape":[200,200],"data_offsets":[569205,609205]},"204":{"dtype":"I8","shape":[204,204],"data_offsets":[609205,650821]},"208":{"dtype":"I8","shape":[208,208],"data_offsets":[650821,694085]},"212":{"dtype":"I8","shape":[212,212],"data_offsets":[694085,739029]},"216":{"dtype":"I8","shape":[216,216],"data_offsets":[739029,785685]},"220":{"dtype":"I8","shape":[220,220],"data_offsets":[785685,834085]},"224":{"dtype":"I8","shape":[224,224],"data_offsets":[834085,884261]},"228":{"dtype":"I8","shape":[228,228],"data_offsets":[884261,936245]},"232":{"dtype":"I8","shape":[232,232],"data_offsets":[936245,990069]},"236":{"dtype":"I8","shape":[236,236],"data_offsets":[990069,1045765]},"24":{"dtype":"I8","shape":[24,24],"data_offsets":[1045765,1046341]},"240":{"dtype":"I8","shape":[240,240],"data_offsets":[1046341,1103941]},"244":{"dtype":"I8","shape":[244,244],"data_offsets":[1103941,1163477]},"248":{"dtype":"I8","shape":[248,248],"data_offsets":[1163477,1224981]},"252":{"dtype":"I8","shape":[252,252],"data_offsets":[1224981,1288485]},"256":{"dtype":"I8","shape":[256,256],"data_offsets":[1288485,1354021]},"28":{"dtype":"I8","shape":[28,28],"data_offsets":[1354021,1354805]},"32":{"dtype":"I8","shape":[32,32],"data_offsets":[1354805,1355829]},"36":{"dtype":"I8","shape":[36,36],"data_offsets":[1355829,1357125]},"4":{"dtype":"I8","shape":[4,4],"data_offsets":[1357125,1357141]},"40":{"dtype":"I8","shape":[40,40],"data_offsets":[1357141,1358741]},"44":{"dtype":"I8","shape":[44,44],"data_offsets":[1358741,1360677]},"48":{"dtype":"I8","shape":[48,48],"data_offsets":[1360677,1362981]},"52":{"dtype":"I8","shape":[52,52],"data_offsets":[1362981,1365685]},"56":{"dtype":"I8","shape":[56,56],"data_offsets":[1365685,1368821]},"60":{"dtype":"I8","shape":[60,60],"data_offsets":[1368821,1372421]},"64":{"dtype":"I8","shape":[64,64],"data_offsets":[1372421,1376517]},"68":{"dtype":"I8","shape":[68,68],"data_offsets":[1376517,1381141]},"72":{"dtype":"I8","shape":[72,72],"data_offsets":[1381141,1386325]},"76":{"dtype":"I8","shape":[76,76],"data_offsets":[1386325,1392101]},"8":{"dtype":"I8","shape":[8,8],"data_offsets":[1392101,1392165]},"80":{"dtype":"I8","shape":[80,80],"data_offsets":[1392165,1398565]},"84":{"dtype":"I8","shape":[84,84],"data_offsets":[1398565,1405621]},"88":{"dtype":"I8","shape":[88,88],"data_offsets":[1405621,1413365]},"92":{"dtype":"I8","shape":[92,92],"data_offsets":[1413365,1421829]},"96":{"dtype":"I8","shape":[96,96],"data_offsets":[1421829,1431045]}}vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/transform/utils/matrix.py000066400000000000000000000123121521257237700322250ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.transform import TransformLocation __all__ = ["get_transform_size", "apply_transform_weight"] def get_transform_size( module: torch.nn.Module, location: TransformLocation, head_dim: int | None = None, ) -> int: """ Determine the size of a transform matrix given its location on the module :param module: module that matrix will be applied to :param location: location on module :param head_dim: size of head when transform is applied to mha :return: size of matrix """ size = None if isinstance(module, torch.nn.Linear): if location in (TransformLocation.INPUT, TransformLocation.WEIGHT_INPUT): size = module.in_features else: size = module.out_features elif isinstance(module, torch.nn.Embedding): if location in (TransformLocation.INPUT, TransformLocation.WEIGHT_INPUT): size = module.num_embeddings else: size = module.embedding_dim elif head_dim is None: raise NotImplementedError( f"Transforms on {type(module)} are not supported without head_dim" ) if head_dim is not None: if size is not None and size % head_dim != 0: raise ValueError( f"{head_dim} must divide {size} for {type(module)} at {location}" ) size = head_dim return size def apply_transform_weight( transform_weight: torch.Tensor, value: torch.Tensor, location: TransformLocation, module_type: type[torch.nn.Module], ) -> torch.Tensor: """ Using the transform location, apply the transform_weight to the given value wrt linear weights. For more info on input and output transforms, see `TransformLocation` The following explains how weights should be applied to values according to location let x be input activation W be weight, yh, xh, Wh be transformed output, input, weight note that y = (x W.T) // torch.nn.Linear Choose values for yh, xh, and Wh which incorporate matrix transforms let V, Vi be transform matrices on input side U, Ui be transform matrices on output side pick xh = (x V) Wh = (U.T W Vi.T) yh = (y U) The following shows that `yh = (xh) (Wh).T` for the chosen values of yh, xh, and Wh (xh) (Wh).T = (x V) (U.T W Vi.T).T = (x V) (Vi W.T U) // transpose matrix product identity = (x W.T) U = y U = yh :param transform_weight: transform weight to apply :param value: value to apply transform_weight to :param location: determines how weight should be applied :param model_type: result of type(module), passed in to determine application of weight transform :return: value after transform_weight has been applied """ assert transform_weight.shape[0] == transform_weight.shape[1] if TransformLocation(location).is_online(): return _multihead_matmul(value, transform_weight) if module_type == torch.nn.Linear: if location == TransformLocation.WEIGHT_INPUT: # equivalent to (transform_weight @ value.T).T return _multihead_matmul(value, transform_weight.T) elif location == TransformLocation.WEIGHT_OUTPUT: # equivalent to (value.T @ transform_weight).T return _multihead_matmul(transform_weight.T, value) # similar derivation to torch.nn.Linear, but `y = (x W)` elif module_type == torch.nn.Embedding: if location == TransformLocation.WEIGHT_INPUT: return _multihead_matmul(transform_weight, value) elif location == TransformLocation.WEIGHT_OUTPUT: return _multihead_matmul(value, transform_weight) raise NotImplementedError( f"Applying transforms to {module_type} {location} is not supported" ) def _multihead_matmul(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor: """ Performs A @ B for last two dims of two matrices A and B that possibly have different shapes, as is the case in multi-headed dimension. If shapes are different, this is equivalent to converting the last two dims of the smaller matrix into a block-diagonal matrix with the same shape as the last two dims of the larger matrix. E.g. if A is half the size of B, this function will perform [[A ] @ B [ A]] If B is a third of the size of A, this function will perform A @ [[B ] [ B ] [ B]] This function will error out if the shapes are not evenly divisble :param A: left-hand tensor :param B: right-hand tensor :return: result """ if A.shape[-1] > B.shape[-2]: head_dim = B.shape[-2] num_heads = A.shape[-1] // head_dim A = A.unflatten(-1, (num_heads, head_dim)) return (A @ B).flatten(-2, -1) elif A.shape[-1] < B.shape[-2]: head_dim = A.shape[-1] num_heads = B.shape[-2] // head_dim B = B.unflatten(-2, (num_heads, head_dim)) return (A @ B).flatten(-3, -2) else: return A @ B vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/utils/000077500000000000000000000000001521257237700263355ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/utils/__init__.py000066400000000000000000000005761521257237700304560ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa from .helpers import * from .internal import * from .match import * from .module import * from .mtp import * from .offload import * from .permutations_24 import * from .safetensors_load import * from .semi_structured_conversions import * from .type import * vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/utils/binary_search.py000066400000000000000000000015121521257237700315170ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Callable, TypeVar T = TypeVar("T") __all__ = ["SearchFailureError", "max_binary_search"] class SearchFailureError(ValueError): pass def max_binary_search( fn: Callable[[int], T], cond: Callable[[T], bool], start: int, end: int, ) -> tuple[int, T]: best_idx = None best_val = None while start <= end: mid = (start + end) // 2 val = fn(mid) if cond(val): # condition is true, search higher best_idx, best_val = mid, val start = mid + 1 else: # condition is false, search lower end = mid - 1 if best_idx is None: raise SearchFailureError() return best_idx, best_val vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/utils/helpers.py000066400000000000000000000362041521257237700303560ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib import warnings from collections.abc import Callable, Iterable, Mapping from functools import wraps from types import MappingProxyType from typing import TYPE_CHECKING, Any, TypeVar import numpy import torch from transformers import AutoConfig, PretrainedConfig T = TypeVar("T", bound="Callable") # used by `deprecated` if TYPE_CHECKING: from compressed_tensors.compressors import ModelCompressor __all__ = [ "infer_compressor_from_model_config", "fix_fsdp_module_name", "tensor_follows_mask_structure", "replace_module", "is_compressed_tensors_config", "getattr_chain", "deprecated", "Aliasable", "combine_shards", "shard_tensor", "pack_bitmasks", "unpack_bitmasks", "patch_attr", "patch_attrs", "ParameterizedDefaultDict", "get_num_attn_heads", "get_num_kv_heads", "get_head_dim", "is_accelerator_type", ] FSDP_WRAPPER_NAME = "_fsdp_wrapped_module" def infer_compressor_from_model_config( pretrained_model_name_or_path: str, ) -> "ModelCompressor | None": # noqa: F821 """ Given a path to a model config, extract a sparsity config if it exists and return the associated ModelCompressor :param pretrained_model_name_or_path: path to model config on disk or HF hub :return: matching compressor if config contains a sparsity config """ from compressed_tensors.compressors import ModelCompressor from compressed_tensors.config import CompressionConfig config = AutoConfig.from_pretrained(pretrained_model_name_or_path) sparsity_config = ModelCompressor.parse_sparsity_config(config) if sparsity_config is None: return None format = sparsity_config.get("format") sparsity_config = CompressionConfig.load_from_registry(format, **sparsity_config) compressor = ModelCompressor.load_from_registry(format, config=sparsity_config) return compressor def fix_fsdp_module_name(name: str) -> str: """ Remove FSDP wrapper prefixes from a module name Accounts for scenario where FSDP_WRAPPER_NAME is at the end of the name, as well as in the middle. :param name: name to strip :return: stripped name """ return name.replace(FSDP_WRAPPER_NAME + ".", "").replace( "." + FSDP_WRAPPER_NAME, "" ) def tensor_follows_mask_structure(tensor, mask: str = "2:4") -> bool: """ :param tensor: tensor to check :param mask: mask structure to check for, in the format "n:m" :return: True if the tensor follows the mask structure, False otherwise. Note, some weights can incidentally be zero, so we check for atleast n zeros in each chunk of size m """ n, m = tuple(map(int, mask.split(":"))) # Reshape the tensor into chunks of size m tensor = tensor.view(-1, m) # Count the number of zeros in each chunk zero_counts = (tensor == 0).sum(dim=1) # Check if the number of zeros in each chunk atleast n # Greater than sign is needed as some weights can incidentally # be zero if not torch.all(zero_counts >= n).item(): raise ValueError() return True def replace_module(model: torch.nn.Module, name: str, new_module: torch.nn.Module): if "." in name: parent_name = name.rsplit(".", 1)[0] child_name = name[len(parent_name) + 1 :] parent = model.get_submodule(parent_name) else: parent_name = "" parent = model child_name = name setattr(parent, child_name, new_module) def is_compressed_tensors_config(compression_config: Any) -> bool: """ Returns True if CompressedTensorsConfig is available from transformers and compression_config is an instance of CompressedTensorsConfig See: https://github.com/huggingface/transformers/pull/31704 """ try: from transformers.utils.quantization_config import CompressedTensorsConfig return isinstance(compression_config, CompressedTensorsConfig) except ImportError: return False def getattr_chain(obj: Any, chain_str: str, *args, **kwargs) -> Any: """ Chain multiple getattr calls, separated by `.` :param obj: base object whose attributes are being retrieved :param chain_str: attribute names separated by `.` :param default: default value, throw error otherwise """ if len(args) >= 1: has_default = True default = args[0] elif "default" in kwargs: has_default = True default = kwargs["default"] else: has_default = False attr_names = chain_str.split(".") res = obj for attr_name in attr_names: if not hasattr(res, attr_name): if has_default: return default else: raise AttributeError(f"{res} object has no attribute {attr_name}") res = getattr(res, attr_name) return res def deprecated( future_name: str | None = None, message: str | None = None ) -> Callable[[T], T]: """ Decorator to mark functions as deprecated :param new_function: Function called in place of deprecated function :param message: Deprecation message, replaces default deprecation message """ def decorator(func: T) -> T: nonlocal message if message is None: message = ( f"{func.__name__} is deprecated and will be removed in a future release" ) if future_name is not None: message += f". Please use {future_name} instead." @wraps(func) def wrapped(*args, **kwargs): warnings.warn(message, DeprecationWarning, stacklevel=2) return func(*args, **kwargs) return wrapped return decorator class Aliasable: """ A mixin for enums to allow aliasing of enum members Example: >>> class MyClass(Aliasable, int, Enum): >>> ... """ @staticmethod def get_aliases() -> dict[str, str]: raise NotImplementedError() def __eq__(self, other): if isinstance(other, self.__class__): aliases = self.get_aliases() return self.value == other.value or ( aliases.get(self.value, self.value) == aliases.get(other.value, other.value) ) else: aliases = self.get_aliases() self_value = aliases.get(self.value, self.value) other_value = aliases.get(other, other) return self_value == other_value def __hash__(self): canonical_value = self.aliases.get(self.value, self.value) return hash(canonical_value) def shard_tensor( tensor: torch.Tensor, shard_sizes: list[int], dim: int = 0 ) -> list[torch.Tensor]: """ Shards a tensor into a list of tensors along a given dimension. raises: ValueError: If the sum of shard_sizes does not match the size of the tensor along the given dimension. :param tensor: The input tensor to shard. :param shard_sizes : List of sizes for each shard along the specified dimension. :param dim : The dimension along which to shard the tensor. :returns: A list of tensors sharded along the specified dimension. """ if sum(shard_sizes) != tensor.size(dim): raise ValueError( "Sum of shard_sizes must equal the size of the tensor " "along the specified dimension." ) shards = [] start_idx = 0 for size in shard_sizes: end_idx = start_idx + size shard = tensor.narrow(dim, start_idx, size) shards.append(shard) start_idx = end_idx return shards def combine_shards(shards: list[torch.Tensor], dim: int = 0) -> torch.Tensor: """ Combine decompressed shards along a given dimension using `narrow`. :param shards: List of decompressed shard tensors. :param dim: Dimension to combine along (default: 0). :return: Combined decompressed tensor. """ if not shards: raise ValueError("The list of shards is empty.") # Assert that all shards have the same dtype shard_dtypes = {shard.dtype for shard in shards} if len(shard_dtypes) > 1: raise ValueError("All shards must have the same dtype.") # Determine the total shape of the combined tensor total_shape = list(shards[0].shape) total_shape[dim] = sum(shard.shape[dim] for shard in shards) # Create the combined tensor combined = torch.zeros(total_shape, dtype=shards[0].dtype, device=shards[0].device) # Fill the combined tensor using narrow shard_offset = 0 for shard in shards: shard_size = shard.shape[dim] combined.narrow(dim, shard_offset, shard_size).copy_(shard) shard_offset += shard_size return combined def pack_bitmasks(bytemasks: torch.Tensor) -> torch.Tensor: """ Converts a bytemask tensor to a bitmask tensor to reduce memory. Shape RxC will be compressed to R x ceil(C/8) :param bytemasks: mask tensor where each byte corresponds to a weight :return: mask tensor where each bit corresounds to a weight """ packed_bits_numpy = numpy.packbits(bytemasks.numpy(), axis=-1, bitorder="little") packed_bits_torch = torch.from_numpy(packed_bits_numpy) return packed_bits_torch def unpack_bitmasks( packed_bitmasks: torch.Tensor, original_shape: list[int] ) -> torch.Tensor: """ Converts a bitmask tensor back to a bytemask tensor for use during decompression :param packed_bitmasks: mask tensor where each bit corresponds to a weight :param original_shape: dense shape to decompress to :return: boolean mask of weights in the original dense shape """ # Unpack the bits unpacked_bits = numpy.unpackbits( packed_bitmasks.cpu().numpy(), axis=-1, count=original_shape[-1], bitorder="little", ) # Reshape to match the original shape unpacked_bitmasks_torch = torch.from_numpy( unpacked_bits.reshape(original_shape).astype(bool) ) return unpacked_bitmasks_torch @contextlib.contextmanager def patch_attr(base: object, attr: str, value: Any): """ Patch the value of an object attribute. Original value is restored upon exit :param base: object which has the attribute to patch :param attr: name of the the attribute to patch :param value: used to replace original value Usage: >>> from types import SimpleNamespace >>> obj = SimpleNamespace() >>> with patch_attr(obj, "attribute", "value"): ... assert obj.attribute == "value" >>> assert not hasattr(obj, "attribute") """ _sentinel = object() original_value = getattr(base, attr, _sentinel) setattr(base, attr, value) try: yield finally: if original_value is not _sentinel: setattr(base, attr, original_value) else: delattr(base, attr) @contextlib.contextmanager def patch_attrs(bases: Iterable[Any], attr: str, values: Iterable[Any]): """ Same as `patch_attr` but for a list of objects to patch Patch attribute for a list of objects with list of values. Original values are restored upon exit :param bases: objects which has the attribute to patch :param attr: name of the the attribute to patch :param values: used to replace original values. Must be same length as bases Usage: >>> from types import SimpleNamespace >>> obj1 = SimpleNamespace() >>> obj2 = SimpleNamespace() >>> with patch_attr([obj1, obj2], "attribute", ["value1", "value2"]): ... assert obj1.attribute == "value1" ... assert obj2.attribute == "value2" >>> assert not hasattr(obj1, "attribute") >>> assert not hasattr(obj2, "attribute") """ with contextlib.ExitStack() as stack: for base, value in zip(bases, values): stack.enter_context(patch_attr(base, attr, value)) yield class ParameterizedDefaultDict(dict): """ Similar to `collections.DefaultDict`, but upon fetching a key which is missing, the key is passed as arguments to the `default_factory` :param default_factory: function which takes a key as input and returns the corresponding default value """ def __init__(self, default_factory: Callable[[Any], Any]): self.default_factory = default_factory self._factory_kwargs = MappingProxyType({}) def __missing__(self, key: Any) -> Any: if isinstance(key, tuple): value = self.default_factory(*key, **self._factory_kwargs) else: value = self.default_factory(key, **self._factory_kwargs) self[key] = value return value def get(self, *args, factory_kwargs: Mapping = MappingProxyType({})) -> Any: """ Similar to `__getitem__`, but allows passing kwargs to factory function :param \\*args: args whose tuple will value will be treated as key :param factory_kwargs: keyword arguments to pass to `default_factory` :return: dictionary entry for given key """ with patch_attr(self, "_factory_kwargs", factory_kwargs): return self[args] def get_num_attn_heads(config: PretrainedConfig) -> int: """ Get the number of attention heads used by a model :param config: model config :return: num_attention_heads of model """ if hasattr(config, "num_attention_heads"): return config.num_attention_heads elif hasattr(config, "hidden_size") and hasattr(config, "head_dim"): return config.hidden_size // config.head_dim else: raise ValueError( "Cannot determine num_attention_heads from config. Config must define " "either `num_attention_heads` or both `hidden_size` and `head_dim`. " f"{config}" ) def get_num_kv_heads(config: PretrainedConfig) -> int: """ Get the number of key-value attention heads used by a model :param config: model config :return: num_key_value_heads of model """ if hasattr(config, "num_key_value_heads"): return config.num_key_value_heads else: raise ValueError( "Cannot determine num_key_value_heads from config. Config must define " f"`num_key_value_heads`. {config}" ) def get_head_dim(config: PretrainedConfig) -> int: """ Get the number of dimensions used by the attention heads of a model :param config: model config :return: head_dim of model """ if hasattr(config, "head_dim"): return config.head_dim elif hasattr(config, "hidden_size") and hasattr(config, "num_attention_heads"): return config.hidden_size // config.num_attention_heads else: raise ValueError( "Cannot determine head_dim from config. Config must define " "either `head_dim` or both `hidden_size` and `num_attention_heads`. " f"{config}" ) def is_accelerator_type(device_type: str) -> bool: """Return ``True`` if *device_type* matches the current accelerator. Works for any backend exposed via :mod:`torch.accelerator` — CUDA, XPU, NPU, etc. Returns ``False`` when no accelerator is present. """ if not torch.accelerator.is_available(): return False return device_type == torch.accelerator.current_accelerator().type vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/utils/internal.py000066400000000000000000000007301521257237700305230ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch __all__ = ["InternalModule"] class InternalModule(torch.nn.Module): """ Abstract base class for modules which are not a part of the the model definition. `torch.nn.Module`s which inherit from this class will not be targeted by configs This is typically used to skip apply configs to `Observers` and `Transforms` """ pass vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/utils/match.py000066400000000000000000000451161521257237700300120ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import logging import os import re from collections import defaultdict from collections.abc import Generator, Iterable, Mapping from typing import Iterator import torch from compressed_tensors.utils.internal import InternalModule _LOGGER: logging.Logger = logging.getLogger(__name__) __all__ = [ "match_name", "match_named_modules", "match_named_parameters", "match_targets", "match_modules_set", "match_quantizable_tensors", "get_lowest_common_ancestor_name", "is_match", "is_narrow_match", ] FusedMappping = Mapping[str, Iterable[str]] def match_named_modules( model: torch.nn.Module, targets: Iterable[str] | None, ignore: Iterable[str] | None = None, fused: FusedMappping | None = None, warn_on_fail: bool = False, ) -> Generator[tuple[str, torch.nn.Module], None, None]: """ Yields names and modules which match `targets` but do not match `ignore`. Values are returned in order of `model.named_modules()` :param model: model containing submodules to match against :param targets: target strings, potentially containing "re:" prefixes :param ignore: targets to ignore, potentially containing "re:" prefixes :fused: optional mapping from suffixes of fused modules to the suffixes of their corresponding shards. See `compressed_tensors.utils.match.is_match` :param warn_on_fail: if True, warns if any targets do not match any modules in model :return: generator of module names and modules """ targets = targets or [] ignore = ignore or [] unmatched_targets = set(targets) for name, module in model.named_modules(): for target in targets: if is_match(name, module, target, fused=fused): unmatched_targets -= {target} if not is_match(name, module, ignore, fused=fused): yield name, module break if warn_on_fail: for target in unmatched_targets: _LOGGER.warning( f"Could not match `{target}` in instance of {model.__class__.__name__}" ) def match_named_parameters( model: torch.nn.Module, targets: Iterable[str] | None, ignore: Iterable[str] | None = None, fused: FusedMappping | None = None, warn_on_fail: bool = False, ) -> Generator[tuple[str, torch.nn.Module, torch.nn.Parameter], None, None]: """ Yields parameters which match `targets` but do not match `ignore`. Values are returned in order of `model.named_modules()` :param model: model containing params to match against :param targets: target strings, potentially containing "re:" prefixes :param ignore: targets to ignore, potentially containing "re:" prefixes :fused: optional mapping from suffixes of fused modules to the suffixes of their corresponding shards. See `compressed_tensors.utils.match.is_match` :param warn_on_fail: if True, warns if any targets do not match any params in model :return: generator of fully-qualified param names, parent modules, and params """ targets = targets or [] ignore = ignore or [] unmatched_targets = set(targets) for module_name, module in model.named_modules(): if isinstance(module, InternalModule): continue for param_name, param in module.named_parameters(recurse=False): param_fqn = f"{module_name}.{param_name}" for target in targets: if match_name(param_fqn, target, fused): unmatched_targets -= {target} if not any(match_name(param_fqn, ign, fused) for ign in ignore): yield param_fqn, module, param if warn_on_fail: for target in unmatched_targets: _LOGGER.warning( f"Could not match `{target}` in instance of {model.__class__.__name__}" ) def match_targets( name: str, module: torch.nn.Module, targets: Iterable[str] | None ) -> list[str]: """ Returns the targets that match the given name and module. :param name: the name of the module :param module: the module to match :param targets: the target strings, potentially containing "re:" prefixes :return: the targets that match the given name and module Outputs are ordered by type: exact name match, regex name match, class name match """ targets = targets or [] if isinstance(module, InternalModule): return [] # The order of the output `matches` list matters, they are arranged from most # specific to least specific, and this order will be used when merging configs. # The entries are sorted in the following order: # 1. matches on exact strings # 2. matches on regex patterns # 3. matches on module names (e.g. "Linear") targets = sorted(targets, key=lambda x: ("re:" in x, x)) matched_targets = [] for target in targets: if match_name(name, target): matched_targets.append(target) for target in targets: if _match_class(module, target) and target not in matched_targets: matched_targets.append(target) return matched_targets def get_lowest_common_ancestor_name(names: list[str | None]) -> str: """ Given a list of names, returns the lowest-scope common name ignoring Nones. Implementation is a small alteration of os.path.commonprefix https://docs.python.org/3/library/os.path.html#os.path.commonprefix ([s1, s2]->prefix->result) # case 0: multiple modules: [abc.a., abc.b.] -> .abc. -> abc # case 1: single module: [abc.] -> .abc. -> abc # case 2: substring modules: [abc., ab.] -> .ab -> "" # case 3: parent & child: [ab., ab.a.] -> .ab. -> ab """ names = [name for name in names if name is not None] if len(names) == 0: return "" # 1) find longest shared prefix s1 = "." + min(names) + "." s2 = "." + max(names) + "." common_prefix = os.path.commonprefix([s1, s2]) # 2) throw away right most dot and name fragment, throw away leftmost char # ".keep.thro" -> "keep", "." -> "" return common_prefix[1 : common_prefix.rfind(".")] def match_modules_set( model: torch.nn.Module, targets: Iterable[str] | None, ignore: Iterable[str] | None = None, error_on_module_rematch: bool = True, ) -> Generator[list[list[torch.nn.Module]], None, None]: """ Yields modules grouped by parent context. We group by parent context so that we can return ALL matches of a specific target that can be paired with another target. This is most relevant in the case of MoE modules with multiple modules for each expert i.e. post_attention_layernorm <-> mlp.expert.N.gate_proj, mlp.expert.N.up_proj for all N. The parent context will differ from one layer to another while being the same for one expert to another. Each returned group is a list (of lists) with the same size and order as `targets` while all matches for each target and the overall order of the groups are ordered in the same way as `model.named_modules` E.g. the following targets would yield modules belonging to the following layers: ```python3 match_modules_set(model, ["q_proj", "k_proj", "v_proj"]) == ( [ [`layers.0.self_attn.q_proj`], [`layers.0.self_attn.k_proj`], [`layers.0.self_attn.v_proj`], ], [ [`layers.1.self_attn.q_proj`], [`layers.1.self_attn.k_proj`], [`layers.1.self_attn.v_proj`], ], ... ) ``` This can be used to match layers to their corresponding downstream counterparts. For example, matching layer norms to their subsequent linear layers ```python3 for norm, q, k, v in match_modules_set(model, (norm_tgt, q_tgt, k_tgt, v_tgt)): fuse_norm_linears(*norm, [*q, *k, *v]) ``` Alternatively for MoE you would get multiple matches per target per group, E.g. ```python3 targets = [ "post_attention_layernorm", "up_proj", "down_proj" ] match_modules_set(model, targets) == ( [ [layers.0.post_attention_layernorm], [ `layers.0.mlp.experts.0.up_proj`, `layers.0.mlp.experts.1.up_proj`, ... ], [ `layers.0.mlp.experts.0.down_proj`, `layers.0.mlp.experts.1.down_proj`, ... ] ], # <- first yield [ [layers.1.post_attention_layernorm], [ `layers.1.mlp.experts.0.up_proj`, `layers.1.mlp.experts.1.up_proj`, ... ], [ `layers.1.mlp.experts.0.down_proj`, `layers.1.mlp.experts.1.down_proj`, ... ] ], ... ) ``` :param model: model containing modules to match against :param targets: target strings, potentially containing "re:" prefixes :param ignore: targets to ignore, potentially containing "re:" prefixes :param error_on_module_rematch: if True, errors when a module gets matched to multiple targets, if False, no error. (Defaults to True) """ targets = targets or [] ignore = ignore or [] # as we iterate through modules and try to match them with targets, # the algorithm can be in 2 possible states: # 0) unmatched_targets > 0, i.e. some of the targets haven't been matched. # Keep matching until all targets have at least one match # 1) unmatched_targets == 0 i.e. we have at least one match for each target. # At this point we are unsure if we have a full set or if we need to add # more matches. # There are 3 things that can happen once were in state 1: # A) found a new match with same parent_context, # (add it to matches and keep going) # B) found a new match with different parent_context, i.e. we found a match # that requires a deeper parent context, this indicates that this match # should be part of a new set. # (yield current set [not including newest match] and go back to state 0) # C) ran out of modules, we will always yield the final remaining set when # we we've iterated through all the modules in the model. # (yield final set then exit.) # Note: its possible to iterate through all the modules in the model while # not having a full matched set if the user specified a bad matching, in # that case something has gone wrong and we error matches = defaultdict(list) parent_context = None unmatched_targets = set(targets) for name, module in model.named_modules(): matched_targets_for_cur_module = set() for target in targets: if is_match(name, module, target, ignore): new_parent_context = get_lowest_common_ancestor_name( [name, parent_context] ) # code for (B) if not unmatched_targets and new_parent_context != parent_context: yield [matches[target] for target in targets] matches = defaultdict(list) new_parent_context = name unmatched_targets = set(targets) matches[target].append(module) parent_context = new_parent_context unmatched_targets -= {target} matched_targets_for_cur_module |= {target} if len(matched_targets_for_cur_module) > 1 and error_on_module_rematch: raise ValueError( f"module: {name} was matched with multiple targets: " f"{matched_targets_for_cur_module} which is unexpected " "disable this check by setting `error_on_module_rematch = False`" ) # never found anything if unmatched_targets == set(targets): return # code for (C) if not unmatched_targets: # have a full matching yield [matches[target] for target in targets] return raise ValueError( f"Found a final incomplete set with matches found for keys: " f"{set(targets) - unmatched_targets} " f"but no matches found for keys: {unmatched_targets}" ) def is_match( name: str, module: torch.nn.Module, targets: str | Iterable[str], ignore: str | Iterable[str] = tuple(), fused: FusedMappping | None = None, ) -> bool: """ Returns true if either module name or module parent classes match against target and the module is not an internal module. The name and module may refer to a fused module defined by vLLM. In these cases, a `fused` mapping must be provided. For example, in `vllm/model_executor/models/llama.py`: ```python packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"] } ``` :param name: name of module :param module: module to match :param target: target which matches name or module, potentially contains regex :fused: optional mapping from suffixes of fused modules to the suffixes of their corresponding shards """ targets = [targets] if isinstance(targets, str) else targets ignore = [ignore] if isinstance(ignore, str) else ignore return not isinstance(module, InternalModule) and ( any( match_name(name, target, fused) or _match_class(module, target) for target in targets ) and not any( match_name(name, ign, fused) or _match_class(module, ign) for ign in ignore ) ) def is_narrow_match( model: torch.nn.Module, targets: str | Iterable[str], name: str, module: torch.nn.Module | None = None, ) -> bool: """ Checks if any of the targets narrowly match the module. A target narrowly matches a module if the target matches the module, but does not match the module's parent :param model: model containing both module and its parent :param targets: target strings, potentially containing "re:" prefixes :param name: name of module to match :param module: module to match. If none is provided, then get module from model :return: True if any of the targets narrow match the module """ targets = [targets] if isinstance(targets, str) else targets module = module if module is not None else model.get_submodule(name) parent_name = name.rsplit(".", 1)[0] parent = model.get_submodule(parent_name) return any( is_match(name, module, target) and not is_match(parent_name, parent, target) for target in targets ) def match_name(name: str, target: str, fused: FusedMappping | None = None) -> bool: """ Returns true if target string begins with "re:" and regex matches or if target string exactly matches name. If the name refers to a fused module defined by vLLM, a `fused` mapping must be provided. :param name: name of module :param target: target name, potentially contains regex :fused: optional mapping from suffixes of fused modules to the suffixes of their corresponding shards """ if fused is not None: for fused_suffix in fused: if name.endswith(fused_suffix): name_stripped = name.removesuffix(fused_suffix) return any( match_name(name_stripped + shard_suffix, target) for shard_suffix in fused[fused_suffix] ) if target.startswith("re:"): return re.match(target.removeprefix("re:"), name) is not None else: return target == name def _match_class(module: torch.nn.Module, target: str) -> bool: """ Returns true if any torch parent class names match the target string exactly. A special exception is made for vllm's `LinearBase` class which matches `Linear` :param module: module to match :param target: target which matches name or module """ # will never match against a regex pattern since `:` is not allowed in class names return any( ( issubclass(cls, torch.nn.Module) and ( cls.__name__ == target or (cls.__name__ == "LinearBase" and target == "Linear") ) ) for cls in module.__class__.__mro__ ) def match_quantizable_tensors( tensors: Mapping[str, torch.Tensor], ignore: Iterable[str], targets: Iterable[str] = tuple(), param_targets: Iterable[str] = ("weight",), allow_nonquantizable: bool = False, ) -> Iterator[tuple[str, str]]: """ Match all quantizable tensors that are targeted and not ignored. Note that inputs "targets" and "ignore" operate on module names, matching their usage elsewhere in `compressed_tensors.utils.match`. :param tensors: Mapping of name in safetensors file to actual tensor :param ignore: ignore individual tensor if any match is found within elements in ignore when compared to module name (regex allowed) :param targets: include if any match is found within elements in targets when compared to module name (regex allowed). Unlike model-based matching utils, this only checks names, not classes. If empty or if "Linear" is included, assume targets is all-inclusive. :param param_targets: same as targets, but for param name instead of module name. Useful when qparam names like "weight_scale", "weight_packed", etc. need to be matched as well (regex allowed). :param allow_nonquantizable: Override to include non-quantizable modules, like layernorms. Note that this list must be kept up-to-date with naming conventions, unless a more concrete check can be found without necessitating loading the model up in transformers, or inspecting the model def code itself. :return: iterator of module_name and full tensor name meeting filtering criterion. """ for name in list(tensors.keys()): module_name, _, param_name = name.rpartition(".") if not allow_nonquantizable and module_name.endswith("norm"): continue is_param_targeted = any( (match_name(param_name, target)) for target in param_targets ) if not is_param_targeted: continue is_module_targeted = ( len(targets) == 0 or "Linear" in targets or any((match_name(module_name, target)) for target in targets) ) if not is_module_targeted: continue is_module_ignored = any(match_name(module_name, ign) for ign in ignore) if is_module_ignored: continue yield module_name, name vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/utils/module.py000066400000000000000000000042731521257237700302020ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from itertools import chain import torch from compressed_tensors.utils.type import TensorStateDict __all__ = ["get_direct_state_dict", "replace_direct_state_dict"] def get_direct_state_dict(module: torch.nn.Module) -> TensorStateDict: """ Extract a state dict directly from a module's parameters and buffers. Returns tensor data (unwrapped from Parameter/Buffer wrappers) for all parameters and buffers in the module. Does not recurse into child modules. :param module: the module to extract state from :return: dict mapping parameter/buffer names to their tensor data """ return { name: ( tensor.data if isinstance(tensor, (torch.nn.Parameter, torch.nn.Buffer)) else tensor ) for name, tensor in chain(module._parameters.items(), module._buffers.items()) } def replace_direct_state_dict(module: torch.nn.Module, new_state_dict: TensorStateDict): """ Replace a module's parameters and buffers with a new state dict. Removes parameters/buffers that exist in the old state but not the new state, and adds/updates parameters from the new state dict. All new tensors are added as non-trainable parameters (not buffers). Skips unchanged values for efficiency. :param module: the module to update :param new_state_dict: dict of new parameter/buffer values """ from compressed_tensors.offload import disable_onloading with disable_onloading(): old_state_dict = get_direct_state_dict(module) for name in old_state_dict: # remove attributes that don't exist in the new state if name not in new_state_dict: delattr(module, name) for name, new_value in new_state_dict.items(): if name in old_state_dict: # skip unchanged values if old_state_dict[name] is new_value: continue # overwrite changed values delattr(module, name) # treat all new tensors as parameters (not buffers) setattr(module, name, torch.nn.Parameter(new_value, requires_grad=False)) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/utils/mtp.py000066400000000000000000000062731521257237700275170ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import os from compressed_tensors.base import QUANTIZATION_CONFIG_NAME from compressed_tensors.utils.safetensors_load import ( _fetch_and_save_prefix_tensors, find_config_path, get_weight_mappings, update_safetensors_index, ) __all__ = ["save_mtp_tensors_to_checkpoint"] def save_mtp_tensors_to_checkpoint( source_model: str, dest_dir: str, mtp_prefix: str = "mtp", shard_name: str = "model_mtp.safetensors", ): """ Extracts MTP (Multi-Token Prediction) tensors from a source model checkpoint and saves them into a destination checkpoint directory. Updates the safetensors index to include the new MTP shard and updates the quantization_config ignore list in config.json so inference engines skip quantization for MTP layers. MTP layers are not quantized and are excluded from the model's state dict during quantization (e.g. via _keys_to_ignore_on_load_unexpected). This function copies them as-is from the original checkpoint so they are present in the final saved checkpoint. :param source_model: local path or HuggingFace stub of the original (unquantized) model to extract MTP tensors from :param dest_dir: path to the quantized checkpoint directory to save MTP tensors into :param mtp_prefix: key prefix used to identify MTP tensors, defaults to "mtp" :param shard_name: filename for the new shard written into dest_dir, defaults to "model_mtp.safetensors" """ # Extract MTP tensors from the original checkpoint and save them as a new # shard in dest_dir. MTP layers are not part of the quantized model so they # must be carried over as-is. mtp_tensors = _fetch_and_save_prefix_tensors( source_model, mtp_prefix, dest_dir, shard_name ) # Build weight_map from existing index or single-shard file, then add MTP entries. # update_safetensors_index will create the index file if it doesn't exist yet. weight_map = { k: os.path.basename(v) for k, v in get_weight_mappings(dest_dir).items() } weight_map.update({key: shard_name for key in mtp_tensors}) total_size = sum( os.path.getsize(os.path.join(dest_dir, s)) for s in set(weight_map.values()) ) update_safetensors_index(dest_dir, total_size, weight_map) # Update quantization_config.ignore in config.json so inference engines # know not to apply quantization to MTP layers config_path = find_config_path(dest_dir) if config_path is not None: with open(config_path, "r") as f: config = json.load(f) quant_config = config.get(QUANTIZATION_CONFIG_NAME) if quant_config is not None: ignore_list = quant_config.get("ignore") or [] mtp_ignore_pattern = f"re:^{mtp_prefix}.*" if mtp_ignore_pattern not in ignore_list: ignore_list.append(mtp_ignore_pattern) quant_config["ignore"] = ignore_list config[QUANTIZATION_CONFIG_NAME] = quant_config with open(config_path, "w") as f: json.dump(config, f, indent=2) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/utils/offload.py000066400000000000000000000047751521257237700303360ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ Utilities associated with offloading functionality | ------------------------------------------------------------------------------------------------------ | # noqa: E501 | Operation | Without offloading support | With offloading support | # noqa: E501 | ---------- | -------------------------------------- | ------------------------------------------------ | # noqa: E501 | Update | module.name.data.copy_(new_data) | update_offload_parameter(module, name, new_data) | # noqa: E501 | ------------------------------------------------------------------------------------------------------ | # noqa: E501 """ from typing import Literal import torch from compressed_tensors.offload import ( align_module_device, align_modules, disable_offloading, get_execution_device, get_offloaded_device, register_offload_module, remove_dispatch, update_offload_parameter, ) from compressed_tensors.utils.helpers import deprecated __all__ = [ "get_execution_device", "get_offloaded_device", "register_offload_parameter", "update_offload_parameter", "delete_offload_parameter", "align_modules", "align_module_device", "register_offload_module", "disable_offloading", "remove_dispatch", ] """ Candidates for Upstreaming """ @deprecated("module.register_parameter(name, parameter)") def register_offload_parameter( module: torch.nn.Module, name: str, parameter: torch.nn.Parameter, offload_device: torch.device | Literal["disk"] | None = None, ): """ Register a parameter to the given module which may be offloaded :param module: maybe offloaded module :param name: name of newly registered parameter :param parameter: parameter being registered :param offload_device: device on which weight will be offloaded to. If None is provided, then infer device from parameters on module """ if offload_device == "disk": raise NotImplementedError("Disk offloading is not currently supported") module.register_parameter(name, parameter) @deprecated("delattr(module, name)") def delete_offload_parameter(module: torch.nn.Module, name: str): """ Delete a parameter from a module which may be offloaded, including any metadata in _hf_hook :param module: maybe offloaded module :param name: name of parameter being deleted """ delattr(module, name) vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/utils/permutations_24.py000066400000000000000000000037431521257237700317550ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import numpy import torch __all__ = ["get_permutations_24"] # Precompute permutations for Marlin24 weight and scale shuffling # Originally implemented in nm-vllm/vllm/model_executor/layers/quantization/utils/marlin_24_perms.py # noqa: E501 # # Marlin works on [16*2,64] tiles. The goal of the permutations is to reorder the weight # data so that it is compatible with the tensor-core format that is described here: # https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type # noqa: E501 # # As a result of this reordering, the vector loads inside the kernel will get the data # as it is needed for tensor-core (without the need to use ldmatrix instructions) def get_permutations_24(num_bits): perm_list = [] for i in range(32): perm1 = [] col = i // 4 col_o = col // 2 for block in [0, 1]: for row in [ 2 * (i % 4), 2 * (i % 4) + 1, 2 * (i % 4 + 4), 2 * (i % 4 + 4) + 1, ]: perm1.append(16 * row + col_o * 256 + 8 * (col % 2) + 4 * block) for j in range(4): perm_list.extend([p + 1 * j for p in perm1]) perm = numpy.array(perm_list) if num_bits == 4: interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7]) elif num_bits == 8: interleave = numpy.array([0, 2, 1, 3]) else: raise ValueError("num_bits must be 4 or 8, got {}".format(num_bits)) perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel() perm = torch.from_numpy(perm) scale_perm = [] for i in range(8): scale_perm.extend([i * 8 + j for j in [0, 4, 1, 5, 2, 6, 3, 7]]) scale_perm_single = [] for i in range(8): scale_perm_single.extend([8 * i + j for j in [0, 1, 2, 3, 4, 5, 6, 7]]) return perm, scale_perm, scale_perm_single vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/utils/safetensors_load.py000066400000000000000000000506231521257237700322500ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import os import re import struct from collections.abc import Iterable import torch from compressed_tensors.base import QUANTIZATION_CONFIG_NAME from huggingface_hub import list_repo_files from safetensors import safe_open from safetensors.torch import save_file from transformers.file_utils import CONFIG_NAME from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, cached_file __all__ = [ "get_safetensors_folder", "get_safetensors_header", "match_param_name", "get_weight_mappings", "get_nested_weight_mappings", "get_quantization_parameter_to_path_mapping", "InverseWeightMap", "load_tensors_from_inverse_weight_map", "is_quantization_param", "find_config_path", "get_quantization_config", "find_safetensors_index_path", "find_safetensors_index_file", "get_weight_map", "update_safetensors_index", "is_weights_file", "get_checkpoint_files", ] WeightMappingType = dict[str, str] NestedWeightMappingType = dict[str, WeightMappingType] str_to_torch_dtype = { "BOOL": torch.bool, "U8": torch.uint8, "I8": torch.int8, "I16": torch.int16, "U16": torch.uint16, "F16": torch.float16, "BF16": torch.bfloat16, "I32": torch.int32, "U32": torch.uint32, "F32": torch.float32, "F64": torch.float64, "I64": torch.int64, "U64": torch.uint64, "F8_E4M3": torch.float8_e4m3fn, "F8_E5M2": torch.float8_e5m2, } def is_weights_file(file_name: str) -> bool: """ Check whether a filename corresponds to a model weights file based on its extension. :param file_name: filename to check :return: True if the file is a recognized weights format, else False """ return any( file_name.endswith(suffix) for suffix in [ ".bin", ".safetensors", ".pth", ".msgpack", ".pt", ] ) def get_checkpoint_files(model_stub: str | os.PathLike) -> dict[str, str]: """ Given a local path or HuggingFace model stub, return a mapping from each file's relative path to its resolved local path. Local directories are walked recursively; HuggingFace stubs are resolved via the Hub API. :param model_stub: local path to a model directory or HuggingFace model stub :return: dict mapping relative file path to resolved local file path """ # In the future, this function can accept and pass download kwargs to cached_file if os.path.exists(model_stub): file_paths = _walk_directory_files( model_stub, ignore=[".cache", ".gitattributes"] ) else: file_paths = list_repo_files(model_stub) return {file_path: cached_file(model_stub, file_path) for file_path in file_paths} def _walk_directory_files(root_dir: str, ignore: Iterable[str]) -> list[str]: """ Return all file paths relative to root_dir, optionally skipping entries whose relative path starts with `ignore`. :param root_dir: root directory to walk :param ignore: optional path prefix to exclude from results :return: list of relative file paths """ all_files = [] for dirpath, _, filenames in os.walk(root_dir): for filename in filenames: rel_path = os.path.relpath(os.path.join(dirpath, filename), root_dir) if not any([rel_path.startswith(i) for i in ignore]): all_files.append(rel_path) return all_files def find_safetensors_index_path(save_directory: str | os.PathLike) -> str | None: """ Search save_directory for a safetensors weight index file. :param save_directory: directory to search :return: absolute path to the index file, or None if not found """ for file_name in os.listdir(save_directory): if file_name.endswith("safetensors.index.json"): return os.path.join(save_directory, file_name) return None def find_config_path(save_directory: str | os.PathLike) -> str | None: """ Search save_directory for a model config file - check first for config.json - then for params.json If still not found, returns None :param save_directory: directory to search :return: absolute path to the config file, or None if not found """ file_names = os.listdir(save_directory) if CONFIG_NAME in file_names: return os.path.join(save_directory, CONFIG_NAME) elif "params.json" in file_names: return os.path.join(save_directory, "params.json") return None def get_quantization_config(config_path: str) -> dict | None: """ Given path to a checkpoint's config.json file, get quantization config Follow cascading pattern in vllm: - check "quantization_config" - check "text_config.quantization_config" - check "compression_config" If still not found, returns None https://github.com/vllm-project/vllm/blob/v0.20.0/vllm/model_executor/model_loader/weight_utils.py#L259-L278 :param config_path: resolved path to config.json :return: quantization config as dictionary pulled from config.json. If not found, returns None. """ with open(config_path, "r") as f: config = json.load(f) quant_config = None if QUANTIZATION_CONFIG_NAME in config: quant_config = config[QUANTIZATION_CONFIG_NAME] elif "text_config" in config and QUANTIZATION_CONFIG_NAME in config["text_config"]: quant_config = config["text_config"][QUANTIZATION_CONFIG_NAME] elif "compression_config" in config: quant_config = config["compression_config"] return quant_config def find_safetensors_index_file(model_files: dict[str, str]) -> str | None: """ Find safetensors index file from full list of model_files :param model_files: mapping of file relative path to absolute path, usually the result of `get_checkpoint_files` :return: absolute path to the safetensors index file, or None if not found """ for file_path, resolved_path in model_files.items(): if file_path.endswith(SAFE_WEIGHTS_INDEX_NAME): return resolved_path # Fallback when names like consolidated.safetensors.index.json are used # as is the case for RedHatAI/Mistral-Small-4-119B-2603-NVFP4 for file_path, resolved_path in model_files.items(): if file_path.endswith(".safetensors.index.json"): return resolved_path return None def get_weight_map(model_files: dict[str, str]) -> dict[str, str]: """ Get weight map from full list of model_files. If safetensors index.json file is found, weight_map can be pulled from there. Otherwise, it is created from the single safetensors weights file. :returns: weight map of form {weight name -> safetensor file name} """ index_file = find_safetensors_index_file(model_files) if index_file is not None: with open(index_file, "r") as f: return json.load(f)["weight_map"] # if no index_file, use model.saftensors instead. if SAFE_WEIGHTS_NAME not in model_files: raise ValueError( f"File {SAFE_WEIGHTS_NAME} expected but not found in {model_files.keys()}" ) # create from model.safetensors with safe_open(model_files[SAFE_WEIGHTS_NAME], framework="pt") as file: return {tensor: SAFE_WEIGHTS_NAME for tensor in file.keys()} def update_safetensors_index( save_directory: str | os.PathLike, total_size: int, weight_map: dict[str, str], ): """ Write (or overwrite) the safetensors weight index file in save_directory. If an existing index file is found it will be replaced in-place; otherwise the standard model.safetensors.index.json filename is used. :param save_directory: directory containing the checkpoint :param total_size: total byte size of all shards, stored in index metadata :param weight_map: mapping from tensor name to shard filename """ file_path = find_safetensors_index_path(save_directory) if file_path is None: file_path = os.path.join(save_directory, SAFE_WEIGHTS_INDEX_NAME) with open(file_path, "w") as file: json.dump( { "metadata": { "total_size": total_size, }, "weight_map": weight_map, }, file, indent=2, sort_keys=True, ) def get_safetensors_folder( pretrained_model_name_or_path: str, cache_dir: str | None = None ) -> str: """ Given a Hugging Face stub or a local path, return the folder containing the safetensors weight files :param pretrained_model_name_or_path: local path to model or HF stub :param cache_dir: optional cache dir to search through, if none is specified the model will be searched for in the default TRANSFORMERS_CACHE :return: local folder containing model data """ if os.path.exists(pretrained_model_name_or_path): # argument is a path to a local folder return os.path.abspath(pretrained_model_name_or_path) safetensors_path = cached_file( pretrained_model_name_or_path, SAFE_WEIGHTS_NAME, cache_dir=cache_dir, _raise_exceptions_for_missing_entries=False, ) index_path = cached_file( pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME, cache_dir=cache_dir, _raise_exceptions_for_missing_entries=False, ) if safetensors_path is not None: # found a single cached safetensors file return os.path.split(safetensors_path)[0] if index_path is not None: # found a cached safetensors weight index file return os.path.split(index_path)[0] # model weights could not be found locally or cached from HF Hub raise ValueError( "Could not locate safetensors weight or index file from " f"{pretrained_model_name_or_path}." ) def get_safetensors_header(safetensors_path: str) -> dict[str, str]: """ Extracts the metadata from a safetensors file as JSON :param safetensors_path: path to a safetensors file :return: dictionary of metadata extracted from the safetensors file """ with open(safetensors_path, "rb") as f: length_of_header = struct.unpack(" str | None: """ Helper function extracting the uncompressed parameterized layer name from a compressed name. Assumes the compressed name was merged using merge_names. :param full_name: full name of parameter in compressed model :param param_name: compression paramater name :return: uncompressed name of the uncompressed parameterized layer """ pattern = r"^(.*)\." + param_name + r"$" regex = re.findall(pattern, full_name) if len(regex) == 0: return None return regex[0] def get_weight_mappings(path_to_model_or_tensors: str) -> dict[str, str]: """ Takes a path to a state dict saved in safetensors format and returns a mapping from parameterized layer name to file location. { layer.weight.bitmask: file_location, layer.weight.row_offsets: file_location, layer.weight.shape: file_location, layer.weight.compressed: file_location } This generalizes to cases where the model is split into multiple safetensors files :param path_to_model_or_tensors: path to directory that contains safetensors (must contain either a single file or multiple files with an index), or a path to a single safetensors file :return: mapping of parameterized layer name to file location """ if os.path.isfile(path_to_model_or_tensors): # we have a single safetensors file to read header = get_safetensors_header(path_to_model_or_tensors) for key in header.keys(): header[key] = path_to_model_or_tensors header.pop("__metadata__", None) else: # we have a directory with multiple safetensors files safetensors_path = os.path.join(path_to_model_or_tensors, SAFE_WEIGHTS_NAME) index_path = os.path.join(path_to_model_or_tensors, SAFE_WEIGHTS_INDEX_NAME) if os.path.exists(safetensors_path): # we have a single safetensors file to read header = get_safetensors_header(safetensors_path) for key in header.keys(): header[key] = SAFE_WEIGHTS_NAME header.pop("__metadata__", None) elif os.path.exists(index_path): # we have multiple safetensors file, read from index with open(index_path, "r", encoding="utf-8") as f: index = json.load(f) header = index["weight_map"] else: raise ValueError( "Could not find a safetensors weight " f"or index file at {path_to_model_or_tensors}" ) # convert weight locations to full paths for key, value in header.items(): header[key] = os.path.join(path_to_model_or_tensors, value) return header def get_nested_weight_mappings( model_path: str, params_to_nest: Iterable[str], return_unmatched_params: bool = False, ) -> NestedWeightMappingType | tuple[NestedWeightMappingType, WeightMappingType]: """ Takes a path to a state dict saved in safetensors format and returns a nested mapping from uncompressed parameterized layer names to the file locations of each layer's compression parameters. Example of the nested mapping: layer: { bitmask: file_location, row_offsets: file_location, shape: file_location, compressed: file_location } If other parameters are found that do not match the nested parameters, they will be returned in a separate dictionary only if return_unmatched_params is True. This dictionary may be needed for cases where compressors are stacked (e.g., quantization compression followed by sparse compression). Example of the unmatched params mapping: { layer.weight_scale: file_location, layer.input_scale: file_location } This generalizes to cases where the model is split into multiple safetensors files. :param model_path: Path to the safetensors state dict, must contain either a single safetensors file or multiple files with an index. :param params_to_nest: Iterable of parameter names to nest. :param return_unmatched_params: If True, return a second dictionary containing the remaining parameters that were not matched to the params_to_nest. :return: - If return_unmatched_params is False: NestedWeightMappingType: A nested mapping of parameterized layer names to file locations of each layer's compression parameters. - If return_unmatched_params is True: Tuple[NestedWeightMappingType, WeightMappingType]: A tuple containing: - NestedWeightMappingType: A nested mapping of parameterized layer names to file locations of each layer's compression parameters. - WeightMappingType: A mapping of the remaining parameter names to their file locations that were not matched to the params_to_nest. """ weight_mappings = get_weight_mappings(model_path) nested_weight_mappings = {} unmatched_params = {} for key, file_location in weight_mappings.items(): matched = False for param_name in params_to_nest: module_path = match_param_name(key, param_name) if module_path: if module_path not in nested_weight_mappings: nested_weight_mappings[module_path] = {} nested_weight_mappings[module_path][param_name] = file_location matched = True if return_unmatched_params and not matched: unmatched_params[key] = file_location if return_unmatched_params: return nested_weight_mappings, unmatched_params return nested_weight_mappings def get_quantization_parameter_to_path_mapping(model_path: str) -> dict[str, str]: """ Given a model path, return a mapping between a parameter and its path on disk """ weight_mappings = get_weight_mappings(model_path) mapping = {} for weight_name, safe_path in weight_mappings.items(): if is_quantization_param(weight_name): mapping[weight_name] = safe_path continue return mapping InverseWeightMap = dict[str, list[str] | None] """ Mapping of absolute path -> list of tensors. Used to pull tensors across different safetensors files that must be loaded/processed together. Used in conjunction with `load_tensors_from_inverse_weight_map` """ def load_tensors_from_inverse_weight_map( inverse_weight_map: InverseWeightMap, device: str | torch.device = torch.device("cpu"), ) -> dict[str, torch.Tensor]: """ Given an inverse_weight_map, which is a dictionary of file name to list of tensor names, load up all listed tensor names :param inverse_weight_map: mapping of resolved source file path -> list of tensor names to load from that file. Precomputed by build_inverse_weight_map() in the job-building phase. If list is empty, all tensors are pulled Example: {"/path/shard0.safetensors": ["q_proj.weight"], "/path/shard1.safetensors": ["k_proj.weight", "v_proj.weight"]} :param device: tensors will be loaded onto this device. Defaults to CPU :returns: mapping of tensor name to actual tensor loaded from safetensors file Example: {"q_proj.weight": torch.Tensor(...), "k_proj.weight: torch.Tensor(...)} """ tensors: dict[str, torch.Tensor] = {} for source_file, tensor_names in inverse_weight_map.items(): with safe_open(source_file, framework="pt") as file: keys = file.keys() # if tensor_names is empty, pull all tensors if tensor_names is None or len(tensor_names) == 0: tensor_names = keys for tensor_name in tensor_names: if tensor_name not in keys: raise ValueError( f"Expected to find tensor {tensor_name} in " f"{source_file}, but tensor was not found." ) if str(device) == "meta": _slice = file.get_slice(tensor_name) dtype = str_to_torch_dtype[_slice.get_dtype()] size = _slice.get_shape() tensors[tensor_name] = torch.empty( size=size, dtype=dtype, device="meta" ) else: tensors[tensor_name] = file.get_tensor(tensor_name).to( device=device ) return tensors def is_quantization_param(name: str) -> bool: """ Checks is a parameter name is associated with a quantization parameter :param name: parameter name to check :return: True if parameter name is a quantization parameter, else False """ if name.endswith("_scale"): return True if name.endswith("zero_point"): return True if name.endswith("g_idx"): return True return False def _fetch_and_save_prefix_tensors( source_model: str, prefix: str, dest_dir: str, shard_name: str ) -> dict: """ Extracts all tensors whose keys start with `prefix` from `source_model` and saves them as a new shard in `dest_dir`. This is useful when saving MTP layers from the original checkpoint into the quantized checkpoint, since MTP layers are not included in the quantized model and thus must be copied over as-is. :param source_model: local path or HuggingFace stub of the source model :param prefix: tensor key prefix to filter on :param dest_dir: destination directory to write the shard into :param shard_name: filename for the new shard :return: dict mapping tensor key to tensor """ source_dir = get_safetensors_folder(source_model) weight_mappings = get_weight_mappings(source_dir) tensors = {} for key, filepath in weight_mappings.items(): if key.startswith(prefix): with safe_open(filepath, framework="pt", device="cpu") as f: tensors[key] = f.get_tensor(key) if len(tensors) <= 0: raise ValueError(f"No tensors with prefix '{prefix}' found in {source_model}") save_file(tensors, os.path.join(dest_dir, shard_name)) return tensors vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/utils/semi_structured_conversions.py000066400000000000000000000312621521257237700345640ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # # Modified by Roberto Lopez Castro (roberto.lopez.castro@udc.es). # Pulled from nm-vllm/vllm/model_executor/layers/quantization/utils/format_24.py # # flake8: noqa # isort: skip_file import torch __all__ = [ "sparse_semi_structured_from_dense_cutlass", "sparse_semi_structured_to_dense_cutlass", "mask_creator", ] # This is PyTorch implementation of main part of reorder_meta() # function, from tools/util/include/cutlass/util/host_reorder.h file # of CUTLASS source tree. Furthermore, CUTLASS template for sparse # GEMM decides upon layout of this matrix, and at the moment for the # sparse GEMM executed on tensor cores, this is layout described by # ColumnMajorInterleaved<2> data structure, in # include/cutlass/layout/matrix.h of CUTLASS source tree. The # reordering of meta matrix into meta_reordered matrix calculated # according to these segments of CUTLASS code is re-implemented here. # Note that this calculation produces offsets for scattering metadata # matrix elements into reordered metadata matrix elements (or, # equivalently, for gathering reordered metadata matrix element back # into metadata matrix elements). def _calculate_meta_reordering_scatter_offsets(m, meta_ncols, meta_dtype, device): dst_rows = torch.arange(0, m, device=device)[:, None].repeat(1, meta_ncols) dst_cols = torch.arange(0, meta_ncols, device=device).repeat(m, 1) # Reorder the rows, then swizzle the 2x2 blocks. group_x = 64 group_y = 32 if meta_dtype.itemsize == 2 else 16 dst_rows = ( dst_rows // group_x * group_x + (dst_rows % 2) * 2 + (dst_rows % 8) // 4 + ((dst_rows % group_y) % 4) // 2 * 32 + ((dst_rows % group_x) // 8) * 4 ) topright = ((dst_rows % 2 == 0) & (dst_cols % 2 == 1)).to(torch.int8) bottomleft = ((dst_rows % 2 == 1) & (dst_cols % 2 == 0)).to(torch.int8) dst_rows += topright - bottomleft dst_cols -= topright - bottomleft # Assumed that meta tensor is to be stored in CUTLASS # InterleavedColumnMajor layout, and reverse engineered # corresponding code to store values into this tensor. interleave = 2 cols_maj = dst_cols // interleave cols_min = dst_cols % interleave return (cols_maj * m * interleave + dst_rows * interleave + cols_min).view(-1) # This function converts dense matrix into sparse semi-structured # representation, producing "compressed" matrix, in the layout used by # CUTLASS backend, and corresponding metadata matrix. def sparse_semi_structured_from_dense_cutlass(dense): if dense.dim() != 2: raise RuntimeError( f"Expected 2-dimensional dense tensor, got {dense.dim()}-dimensional tensor" # noqa: E501 ) m, k = dense.shape device = dense.device meta_dtype = torch.int8 if dense.dtype == torch.int8: meta_dtype = torch.int32 elif dense.dtype in [torch.half, torch.bfloat16, torch.float, torch.int32]: meta_dtype = torch.int16 else: raise RuntimeError(f"Invalid datatype {dense.dtype} of dense matrix") quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4 if quadbits_per_meta_elem not in (4, 8): raise RuntimeError("Invalid number of elements per meta element calculated") if meta_dtype == torch.int32: if m % 16 != 0: raise RuntimeError( f"Number of rows of dense matrix {m} must be divisible by 16" ) else: if m % 32 != 0: raise RuntimeError( f"Number of rows of dense matrix {m} must be divisible by 32" ) if k % (4 * quadbits_per_meta_elem) != 0: raise RuntimeError( f"Number of columns of dense matrix {k} must be divisible by {4 * quadbits_per_meta_elem}" # noqa: E501 ) if dense.dtype != torch.float: ksparse = 4 dense_4 = dense.view(-1, k // ksparse, ksparse) m0, m1, m2, m3 = (dense_4 != 0).unbind(-1) else: ksparse = 2 dense_2 = dense.view(-1, k // ksparse, ksparse) m0, m2 = m1, m3 = (dense_2 != 0).unbind(-1) meta_ncols = k // (ksparse * quadbits_per_meta_elem) # Encoding quadruples of True/False values as follows: # [True, True, False, False] -> 0b0100 # [True, False, True, False] -> 0b1000 # [False, True, True, False] -> 0b1001 # [True, False, False, True ] -> 0b1100 # [False, True, False, True ] -> 0b1101 # [False, False, True, True ] -> 0b1110 # Thus, lower two bits in the encoding are index of the True value # at the lowest index in the quadruple, and the higher two bits in # the encoding are index of the other True value in the quadruple. # In case there are less than two True values, than False value or # values at some index or indices are considered True for the # encoding. In case there are more than two True values, then the # excess True value(s) at some indices are considered False for # the encoding. The exact encodings used for these cases are as # follows: # [False, False, False, False] -> 0b1110 # [False, False, False, True ] -> 0b1110 # [False, False, True, False] -> 0b1110 # [False, True, False, False] -> 0b1001 # [False, True, True, True ] -> 0b1101 # [True, False, False, False] -> 0b1000 # [True, False, True, True ] -> 0b1100 # [True, True, False, True ] -> 0b0100 # [True, True, True, False] -> 0b0100 # [True, True, True, True ] -> 0b0100 # These particular encodings are chosen, with the help of Espresso # logic minimizer software, for the purpose of minimization of # corresponding Boolean functions, that translate non-zero flags # into encoding bits. Note also possible choices for the first # and last of these encodings were limited only to (0b0100, # 0b1110), in order to produce valid encodings for 1:2 sparsity # case. expr0 = m0 & m1 expr1 = ~m0 & m1 expr2 = ~m0 & ~m1 bit0 = expr1 bit1 = expr2 bit2 = expr0 | expr2 | m3 bit3 = expr1 | ~m1 idxs0 = bit0 | (bit1.to(torch.int64) << 1) idxs1 = bit2 | (bit3.to(torch.int64) << 1) if dense.dtype != torch.float: sparse0 = dense_4.gather( -1, idxs0.unsqueeze(-1) ) # type: ignore[possibly-undefined] sparse1 = dense_4.gather(-1, idxs1.unsqueeze(-1)) sparse = torch.stack((sparse0, sparse1), dim=-1).view(m, k // 2) else: sparse = dense_2.gather(-1, idxs0.unsqueeze(-1) // 2).view( m, k // 2 ) # type: ignore[possibly-undefined] meta_4 = idxs0 | (idxs1 << 2) meta_n = meta_4.view((-1, meta_ncols, quadbits_per_meta_elem)).to(meta_dtype) if quadbits_per_meta_elem == 4: meta = ( meta_n[:, :, 0] | (meta_n[:, :, 1] << 4) | (meta_n[:, :, 2] << 8) | (meta_n[:, :, 3] << 12) ) elif quadbits_per_meta_elem == 8: meta = ( meta_n[:, :, 0] | (meta_n[:, :, 1] << 4) | (meta_n[:, :, 2] << 8) | (meta_n[:, :, 3] << 12) | (meta_n[:, :, 4] << 16) | (meta_n[:, :, 5] << 20) | (meta_n[:, :, 6] << 24) | (meta_n[:, :, 7] << 28) ) # Reorder meta tensor elements. meta_reordered = meta.new_empty( (m * meta_ncols,) ) # type: ignore[possibly-undefined] meta_offsets = _calculate_meta_reordering_scatter_offsets( m, meta_ncols, meta_dtype, device ) meta_reordered.scatter_(0, meta_offsets, meta.view(-1)) return (sparse, meta_reordered.view(m, meta_ncols)) # This function performs reverse of the function above - it # reconstructs dense matrix from a pair of "compressed" matrix, given # in the layout used by CUTLASS backend, and accompanying metadata # matrix. def sparse_semi_structured_to_dense_cutlass(sparse, meta_reordered): if sparse.dim() != 2: raise RuntimeError( f"Expected 2-dimensional sparse tensor, got {sparse.dim()}-dimensional tensor" # noqa: E501 ) m, k = sparse.shape device = sparse.device if meta_reordered.dim() != 2: raise RuntimeError( f"Expected 2-dimensional meta tensor, got {meta_reordered.dim()}-dimensional tensor" # noqa: E501 ) if meta_reordered.device != device: raise RuntimeError( f"Expected meta matrix to be on {device} device, got matrix on {meta_reordered.device} device" # noqa: E501 ) meta_dtype = meta_reordered.dtype if meta_dtype not in (torch.int16, torch.int32): raise RuntimeError(f"Invalid datatype {meta_dtype} of meta matrix") quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4 ksparse = 4 if sparse.dtype != torch.float else 2 meta_nrows, meta_ncols = meta_reordered.shape if meta_nrows != m: raise RuntimeError( f"Number of rows of meta matrix {meta_nrows} must be equal to number of columns of spase matrix {m}" # noqa: E501 ) if meta_ncols * ksparse * quadbits_per_meta_elem != 2 * k: raise RuntimeError( f"Number of columns of sparse matrix {k} different from the {meta_ncols * ksparse * quadbits_per_meta_elem // 2}, " # noqa: E501 "expected according to the number of columns of meta matrix" ) # Undo meta tensor elements reordering. meta_offsets = _calculate_meta_reordering_scatter_offsets( m, meta_ncols, meta_dtype, device ) meta = torch.gather(meta_reordered.view(-1), 0, meta_offsets).view(m, meta_ncols) # Unpack sparse tensor back to original dense tensor, using # information provided by meta tensor. Note that torch.float # datatype is handled pretty much the same as # torch.half/torch.bfloat16, as metadata for a pair of torch.float # value is encoded as if underlying 8 bytes contain four # torch.half/torch.bfloat16 values, where either first two or last # two are zeros. meta_2 = torch.empty( (m, meta_ncols, 2 * quadbits_per_meta_elem), dtype=meta_dtype, device=device, ) if quadbits_per_meta_elem == 4: meta_2[:, :, 0] = meta & 0b11 meta_2[:, :, 1] = (meta >> 2) & 0b11 meta_2[:, :, 2] = (meta >> 4) & 0b11 meta_2[:, :, 3] = (meta >> 6) & 0b11 meta_2[:, :, 4] = (meta >> 8) & 0b11 meta_2[:, :, 5] = (meta >> 10) & 0b11 meta_2[:, :, 6] = (meta >> 12) & 0b11 meta_2[:, :, 7] = (meta >> 14) & 0b11 elif quadbits_per_meta_elem == 8: meta_2[:, :, 0] = meta & 0b11 meta_2[:, :, 1] = (meta >> 2) & 0b11 meta_2[:, :, 2] = (meta >> 4) & 0b11 meta_2[:, :, 3] = (meta >> 6) & 0b11 meta_2[:, :, 4] = (meta >> 8) & 0b11 meta_2[:, :, 5] = (meta >> 10) & 0b11 meta_2[:, :, 6] = (meta >> 12) & 0b11 meta_2[:, :, 7] = (meta >> 14) & 0b11 meta_2[:, :, 8] = (meta >> 16) & 0b11 meta_2[:, :, 9] = (meta >> 18) & 0b11 meta_2[:, :, 10] = (meta >> 20) & 0b11 meta_2[:, :, 11] = (meta >> 22) & 0b11 meta_2[:, :, 12] = (meta >> 24) & 0b11 meta_2[:, :, 13] = (meta >> 26) & 0b11 meta_2[:, :, 14] = (meta >> 28) & 0b11 meta_2[:, :, 15] = (meta >> 30) & 0b11 dense_offsets = meta_2.view(-1) + ( torch.arange(0, 2 * m * k // ksparse, device=device) * 4 ).view(-1, 1).repeat(1, 2).view(-1) dense = torch.zeros((m * 2 * k,), dtype=sparse.dtype, device=device) if sparse.dtype != torch.float: # dense.scatter_(0, dense_offsets, sparse.view(-1)) dense.scatter_(0, dense_offsets, sparse.reshape(-1)) else: dense.view(torch.half).scatter_( 0, dense_offsets, sparse.view(torch.half).view(-1) ) return dense.view(m, 2 * k) def mask_creator(tensor): """ Class for creating N:M sparsity masks. Masks will be created using the N:M ratio, where for every block of M weights, N will be pruned based on ranked weight value. Each mask will correspond to the given tensor. :param N: The number of weights in a group to keep :param M: The size of a weight group """ N = 2 M = 4 mask = None # for i, tensor in enumerate(tensors): if tensor.numel() % M != 0: raise ValueError( f"Tensor of size {tensor.shape} can't be evenly divided into " f"{M} groups" ) num_groups = tensor.numel() // M # N:M sparsity for linear layers tensor_temp = tensor.detach().abs().reshape(num_groups, M) index = torch.argsort(tensor_temp, dim=1)[:, : int(M - N)] w_b = torch.ones(tensor_temp.shape, device=tensor_temp.device) mask = w_b.scatter_(dim=1, index=index, value=0).reshape(tensor.shape) return mask vllm-project-compressed-tensors-c18a0fa/src/compressed_tensors/utils/type.py000066400000000000000000000041001521257237700276630ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Annotated, Any import torch from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler from pydantic.json_schema import JsonSchemaValue from pydantic_core import core_schema __all__ = ["TorchDtype", "TensorStateDict"] class _TorchDtypeAnnotation: @classmethod def __get_pydantic_core_schema__( cls, _source_type: Any, _handler: GetCoreSchemaHandler, ) -> core_schema.CoreSchema: # support strings of the form `torch.xxx` or `xxx` def validate_from_str(name: str) -> torch.dtype: name = name.removeprefix("torch.") try: value = getattr(torch, name) assert isinstance(value, torch.dtype) except Exception: raise ValueError(f"No such torch dtype `torch.{name}`") return value # package validation into a schema (which also validates str type) from_str_schema = core_schema.chain_schema( [ core_schema.str_schema(), core_schema.no_info_plain_validator_function(validate_from_str), ] ) return core_schema.json_or_python_schema( json_schema=from_str_schema, python_schema=core_schema.union_schema( [ # support both torch.dtype or strings core_schema.is_instance_schema(torch.dtype), from_str_schema, ] ), # serialize as `torch.xxx` serialization=core_schema.plain_serializer_function_ser_schema( lambda instance: str(instance) ), ) @classmethod def __get_pydantic_json_schema__( cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler ) -> JsonSchemaValue: return handler(core_schema.str_schema()) TorchDtype = Annotated[torch.dtype, _TorchDtypeAnnotation] TensorStateDict = dict[str, torch.Tensor | None] vllm-project-compressed-tensors-c18a0fa/tests/000077500000000000000000000000001521257237700216275ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/__init__.py000066400000000000000000000001531521257237700237370ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project vllm-project-compressed-tensors-c18a0fa/tests/conftest.py000066400000000000000000000214451521257237700240340ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from math import ceil import pytest import torch from compressed_tensors.offload import update_offload_parameter from compressed_tensors.quantization.utils import calculate_qparams # --------------------------------------------------------------------------- # XPU emulation tests (part 2): TorchFunctionMode device emulation # --------------------------------------------------------------------------- def pytest_addoption(parser): parser.addoption( "--emulate-xpu", action="store_true", default=False, help="Emulate XPU device identity on CUDA hardware via TorchFunctionMode", ) class _FakeDeviceType(str): """A string subclass that acts like a device type but has a .type attribute. Inheriting from str allows it to be passed to torch.device() directly, where DeviceRemapMode will remap it to the real device type. """ def __new__(cls, fake_type: str, real_type: str = None): # Create the string with the fake type value instance = super().__new__(cls, fake_type) instance.type = fake_type instance._fake_type = fake_type instance._real_type = real_type # Store for assert_device_equal instance.index = None # Device index (shadows str.index method) return instance def __repr__(self): return f"device(type='{self.type}')" def pytest_configure(config): """Activate device emulation before test collection (before module imports). Three layers of patching: 1. DeviceRemapMode — intercepts torch.* functions, remaps "xpu" -> "cuda" 2. Accelerator mock — torch.accelerator.current_accelerator() reports "xpu" 3. is_accelerator_type patch — accepts both "xpu" and "cuda" """ if not config.getoption("--emulate-xpu"): return from tests.emulate_device import DeviceRemapMode real_type = torch.accelerator.current_accelerator().type # "cuda" fake_type = "xpu" # Save originals for cleanup config._emulate_orig_current_accelerator = torch.accelerator.current_accelerator config._emulate_orig_device_count = torch.accelerator.device_count config._emulate_orig_is_available = torch.accelerator.is_available config._emulate_orig_current_device_index = torch.accelerator.current_device_index # Snapshot real values before mocking real_device_count = torch.accelerator.device_count() real_is_available = torch.accelerator.is_available() real_current_device_index = torch.accelerator.current_device_index # Layer 1: DeviceRemapMode mode = DeviceRemapMode(fake_type=fake_type, real_type=real_type) mode.__enter__() config._emulate_device_remap_mode = mode # Layer 2: Mock accelerator identity # Return a device-like object that has .type="xpu" and can be stringified # to "xpu" for torch.device() calls (which will then be remapped by DeviceRemapMode) fake_accel = _FakeDeviceType(fake_type, real_type) torch.accelerator.current_accelerator = lambda: fake_accel torch.accelerator.device_count = lambda: real_device_count torch.accelerator.is_available = lambda: real_is_available # Patch current_device_index() to use the real device # instead of trying to initialize the fake XPU backend torch.accelerator.current_device_index = real_current_device_index # Layer 3: Patch is_accelerator_type to accept both types import compressed_tensors.utils as _utils config._emulate_orig_is_accelerator_type = _utils.is_accelerator_type def patched_is_accelerator_type(device_type: str) -> bool: return device_type in (fake_type, real_type) _utils.is_accelerator_type = patched_is_accelerator_type # Also patch base.py's binding since it imported is_accelerator_type directly # and captured the original function before pytest_configure ran import compressed_tensors.offload.cache.base as _base config._emulate_orig_base_is_accelerator_type = _base.is_accelerator_type _base.is_accelerator_type = patched_is_accelerator_type def pytest_unconfigure(config): """Tear down device emulation — restore all patched objects.""" mode = getattr(config, "_emulate_device_remap_mode", None) if mode is not None: mode.__exit__(None, None, None) orig_accel = getattr(config, "_emulate_orig_current_accelerator", None) if orig_accel is not None: torch.accelerator.current_accelerator = orig_accel torch.accelerator.device_count = config._emulate_orig_device_count torch.accelerator.is_available = config._emulate_orig_is_available torch.accelerator.current_device_index = ( config._emulate_orig_current_device_index ) orig_is_accel = getattr(config, "_emulate_orig_is_accelerator_type", None) if orig_is_accel is not None: import compressed_tensors.utils as _utils _utils.is_accelerator_type = orig_is_accel orig_base_is_accel = getattr(config, "_emulate_orig_base_is_accelerator_type", None) if orig_base_is_accel is not None: import compressed_tensors.offload.cache.base as _base _base.is_accelerator_type = orig_base_is_accel # --------------------------------------------------------------------------- # Calibration fixtures # --------------------------------------------------------------------------- def _get_dim(dim: int, value: torch.Tensor): if isinstance(dim, int): dim = [dim] dim = set(dim) reduce_dims = tuple(idx for idx in range(value.ndim) if idx not in dim) return reduce_dims @pytest.fixture def mock_per_group_calibration(): def update_scale_zp( module: torch.nn.Module, base_name: str, value: torch.Tensor, group_size: int ): quantization_scheme = getattr(module, "quantization_scheme", None) if not quantization_scheme: # no quantization scheme nothing to do return arg_name = "weights" if base_name == "weight" else f"{base_name}_activations" args = getattr(quantization_scheme, arg_name, None) rows = value.shape[0] columns = value.shape[1] num_groups = int(ceil(columns / group_size)) scale = torch.zeros((rows, num_groups), dtype=value.dtype, device=value.device) zp_dtype = args.pytorch_dtype() zp = torch.zeros((rows, num_groups), dtype=zp_dtype, device=value.device) group_sizes = torch.full((num_groups,), group_size, dtype=torch.int) end = 0 for group_index, group_count in enumerate(group_sizes): start = end end = start + group_count dim = _get_dim( 0, value[:, start:end], ) min_val = torch.amin(value, dim=dim, keepdims=True) max_val = torch.amax(value, dim=dim, keepdims=True) scale_out, zp_out = calculate_qparams(min_val, max_val, args) scale[:, group_index] = scale_out.squeeze(1) zp[:, group_index] = zp_out.squeeze(1) update_offload_parameter(module, f"{base_name}_scale", scale) update_offload_parameter(module, f"{base_name}_zero_point", zp) return update_scale_zp @pytest.fixture def mock_per_channel_calibration(): def update_scale_zp(module: torch.nn.Module, base_name: str, value: torch.Tensor): quantization_scheme = getattr(module, "quantization_scheme", None) if not quantization_scheme: # no quantization scheme nothing to do return arg_name = "weights" if base_name == "weight" else f"{base_name}_activations" args = getattr(quantization_scheme, arg_name, None) dim = _get_dim(0, value) min_val = torch.amin(value, dim=dim, keepdims=True) max_val = torch.amax(value, dim=dim, keepdims=True) scale, zp = calculate_qparams(min_val, max_val, args) update_offload_parameter(module, f"{base_name}_scale", scale) update_offload_parameter(module, f"{base_name}_zero_point", zp) return update_scale_zp @pytest.fixture def mock_per_tensor_calibration(): def update_scale_zp(module: torch.nn.Module, base_name: str, value: torch.Tensor): quantization_scheme = getattr(module, "quantization_scheme", None) if not quantization_scheme: # no quantization scheme nothing to do return arg_name = "weights" if base_name == "weight" else f"{base_name}_activations" args = getattr(quantization_scheme, arg_name, None) # per tensor quantization just calls calculate_qparams directly min_val, max_val = torch.aminmax(value) scale, zp = calculate_qparams(min_val, max_val, args) update_offload_parameter(module, f"{base_name}_scale", scale) update_offload_parameter(module, f"{base_name}_zero_point", zp) return update_scale_zp vllm-project-compressed-tensors-c18a0fa/tests/emulate_device.py000066400000000000000000000042351521257237700251600ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ Device emulation utilities for ``--emulate-xpu`` testing. Provides :class:`DeviceRemapMode`, a :class:`~torch.overrides.TorchFunctionMode` subclass that transparently remaps device type strings in all ``torch.*`` function calls. This allows the full test suite to run on CUDA hardware while the code paths use XPU (or any other) device strings. """ import re import torch from torch.overrides import TorchFunctionMode class DeviceRemapMode(TorchFunctionMode): """Transparently remap device type strings in all torch operations. When activated, any torch function receiving a device argument with ``fake_type`` will have it silently replaced with ``real_type`` before the call reaches the C++ backend. Uses a strict regex pattern to avoid false-positive replacements on strings that happen to contain the fake device type as a substring. """ def __init__(self, fake_type: str, real_type: str): self.fake_type = fake_type self.real_type = real_type self._device_pat = re.compile(rf"^{re.escape(fake_type)}(?::\d+)?$") def _remap(self, arg): if isinstance(arg, torch.device): if arg.type == self.fake_type: return torch.device(self.real_type, arg.index) elif isinstance(arg, str) and self._device_pat.match(arg): return arg.replace(self.fake_type, self.real_type, 1) elif hasattr(arg, "type") and hasattr(arg, "_fake_type"): # Handle _FakeDeviceType from conftest - convert to string and remap return str(arg).replace(self.fake_type, self.real_type, 1) return arg def __torch_function__(self, func, types, args=(), kwargs=None): kwargs = kwargs or {} # Don't remap torch.device() calls - let them create fake-type devices # Only remap actual tensor operations if func is torch.device: return func(*args, **kwargs) new_args = tuple(self._remap(a) for a in args) new_kwargs = {k: self._remap(v) for k, v in kwargs.items()} return func(*new_args, **new_kwargs) vllm-project-compressed-tensors-c18a0fa/tests/mock_observer.py000066400000000000000000000141241521257237700250430ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Tuple from weakref import ref import torch from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy from compressed_tensors.quantization.utils import ( calculate_qparams, generate_gparam, strategy_cdiv, ) class MockMinMaxObserver(torch.nn.Module): def __init__(self, base_name: str, args: QuantizationArgs, module: torch.nn.Module): super().__init__() self.parent = ref(module) self.base_name = base_name self.args = args # used for testing self.min_vals = None self.max_vals = None def get_min_max(self, observed: torch.Tensor): min_vals = torch.amin(observed, dim=(0, -1)) max_vals = torch.amax(observed, dim=(0, -1)) return min_vals, max_vals def forward(self, observed: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: observed = flatten_for_quantization(observed, self.base_name, self.args) self.min_vals, self.max_vals = self.get_min_max(observed) scales, zero_points = calculate_qparams( min_vals=self.min_vals, max_vals=self.max_vals, quantization_args=self.args, global_scale=getattr(self.parent(), f"{self.base_name}_global_scale", None), ) return scales, zero_points def get_global_scale(self, observed: torch.Tensor): observed = observed.reshape((1, 1, -1)) # per tensor reshape min_vals, max_vals = self.get_min_max(observed) global_scale = generate_gparam(min_vals, max_vals) return global_scale def flatten_for_quantization( value: torch.Tensor, base_name: str, args: QuantizationArgs ) -> torch.Tensor: if base_name == "weight": return flatten_weight_for_quantization(value, args) elif base_name in ("input", "output"): return flatten_activation_for_quantization(value, args) elif base_name in ("q", "k", "v"): return flatten_attention_for_quantization(value, args) else: raise ValueError(f"Unknown quantization base name: {base_name}") def flatten_weight_for_quantization(value: torch.Tensor, args: QuantizationArgs): # value.shape = (num_rows, num_cols) if args.strategy == QuantizationStrategy.TENSOR: # (1, 1, num_weight_elems) return value.reshape((1, 1, -1)) if args.strategy == QuantizationStrategy.TOKEN: raise ValueError("Token quantization cannot be applied to weights") if args.strategy == QuantizationStrategy.CHANNEL: # (1, num_rows, 1, num_cols) return value.unsqueeze(-2).unsqueeze(0) if args.strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP): # (1, num_rows, num_groups, group_size) return value.unflatten(-1, (-1, args.group_size)).unsqueeze(0) if args.strategy == QuantizationStrategy.BLOCK: # (1, num_block_rows, num_block_cols, block_width * block_height) block_height, block_width = args.block_structure num_rows, num_cols = value.shape num_block_rows = strategy_cdiv(num_rows, block_height, args.strategy) num_block_cols = strategy_cdiv(num_cols, block_width, args.strategy) return ( value.reshape( num_block_rows, block_height, num_block_cols, block_width, ) .transpose(1, 2) .flatten(-2, -1) .unsqueeze(0) ) if args.strategy == QuantizationStrategy.ATTN_HEAD: raise ValueError("attention head quantization cannot be applied to weights") assert False, f"Unknown strategy {args.strategy}" def flatten_activation_for_quantization(value: torch.Tensor, args: QuantizationArgs): # value.shape = (batch_size, seq_len, hidden_dim) if args.strategy == QuantizationStrategy.TENSOR: # (batch_size * seq_len, 1, hidden_dim) return value.reshape((-1, 1, value.size(-1))) if args.strategy == QuantizationStrategy.TOKEN: # (batch_size, seq_len, hidden_dim) # warning: token quantization uses `compute_dynamic_scales_and_zp` return value.flatten(2, -1) if args.strategy == QuantizationStrategy.CHANNEL: raise ValueError("Channel quantization cannot be applied to activations") if args.strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP): # (batch_size * seq_len, num_groups, group_size) # warning: group activation quantization uses compute_dynamic_scales_and_zp return value.flatten(0, 1).unflatten(-1, (-1, args.group_size)) if args.strategy == QuantizationStrategy.BLOCK: raise ValueError("Block quantization cannot be applied to activations") if args.strategy == QuantizationStrategy.ATTN_HEAD: raise ValueError("attention head quantization cannot be applied to linear acts") assert False, f"Unknown strategy {args.strategy}" def flatten_attention_for_quantization(value: torch.Tensor, args: QuantizationArgs): # value.shape = (batch_size, num_heads, seq_len, head_dim) if args.strategy == QuantizationStrategy.TENSOR: # (batch_size * seq_len, 1, num_heads * head_dim) return value.transpose(1, 2).flatten(0, 1).flatten(-2, -1).unsqueeze(-2) if args.strategy == QuantizationStrategy.TOKEN: raise ValueError("Token quantization cannot be applied to attention") if args.strategy == QuantizationStrategy.CHANNEL: raise ValueError("Channel quantization cannot be applied to attention") if args.strategy in (QuantizationStrategy.GROUP, QuantizationStrategy.TENSOR_GROUP): raise ValueError("Group quantization cannot be applied to attention") if args.strategy == QuantizationStrategy.BLOCK: raise ValueError("Block quantization cannot be applied to attention") if args.strategy == QuantizationStrategy.ATTN_HEAD: # (batch_size * seq_len, num_heads, 1, 1, head_dim) return value.transpose(1, 2).flatten(0, 1).unsqueeze(-2).unsqueeze(-2) assert False, f"Unknown strategy {args.strategy}" vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/000077500000000000000000000000001521257237700252455ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/distributed/000077500000000000000000000000001521257237700275675ustar00rootroot00000000000000test_distributed_compression.py000066400000000000000000000257311521257237700360740ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/distributed# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Integration tests for distributed model compression.""" import pytest import torch import torch.distributed as dist import torch.nn as nn from compressed_tensors.compressors.model_compressors import ModelCompressor from compressed_tensors.offload import offload_module from compressed_tensors.quantization import ( QuantizationArgs, QuantizationConfig, QuantizationScheme, QuantizationStatus, ) from tests.test_offload.conftest import torchrun from tests.testing_utils import requires_gpu class TwoLayerModel(nn.Module): """Simple model for testing distributed compression.""" def __init__(self): super().__init__() self.layer1 = nn.Linear(10, 10, bias=False) self.layer2 = nn.Linear(10, 10, bias=False) def forward(self, x): x = self.layer1(x) x = self.layer2(x) return x def create_quantization_config(bits=4, format="pack-quantized"): """Helper to create a QuantizationConfig for testing.""" config_dict = { "format": format, "global_compression_ratio": 1.0, "quant_method": "compressed-tensors", "config_groups": { "group_0": { "targets": ["Linear"], "weights": { "num_bits": bits, "strategy": "channel", "symmetric": True, "type": "int", }, } }, } return QuantizationConfig.model_validate(config_dict) def setup_quantized_model(model: nn.Module, bits: int = 4) -> nn.Module: """Set up a model with quantization schemes and parameters.""" scheme = QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( num_bits=bits, strategy="channel", symmetric=True, type="int", ), ) for name, module in model.named_modules(): if isinstance(module, nn.Linear): module.quantization_scheme = scheme module.quantization_status = QuantizationStatus.FROZEN module.weight_scale = nn.Parameter( torch.ones(module.weight.shape[0], 1) * 0.01 ) module.weight_zero_point = nn.Parameter( torch.zeros(module.weight.shape[0], 1, dtype=torch.int32), requires_grad=False, ) return model @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_model_compression(): """Test end-to-end distributed model compression.""" model = TwoLayerModel() setup_quantized_model(model) q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) # Compress the model compressor.compress_model(model) # Verify compression happened assert hasattr(model.layer1, "weight_packed") assert hasattr(model.layer2, "weight_packed") assert model.layer1.weight_packed.dtype == torch.int32 assert model.layer2.weight_packed.dtype == torch.int32 # Verify compression status is updated assert ( compressor.quantization_config.quantization_status == QuantizationStatus.COMPRESSED ) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_compression_consistency(): """Test that all ranks have consistent state after distributed compression.""" model = TwoLayerModel() setup_quantized_model(model) q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) # Compress the model compressor.compress_model(model) dist.barrier() # Compute checksums for each layer layer1_checksum = model.layer1.weight_packed.sum().item() layer2_checksum = model.layer2.weight_packed.sum().item() # Gather checksums from all ranks if dist.get_rank() == 0: gathered_layer1 = [None] * dist.get_world_size() gathered_layer2 = [None] * dist.get_world_size() dist.gather_object(layer1_checksum, gathered_layer1, dst=0) dist.gather_object(layer2_checksum, gathered_layer2, dst=0) # All ranks should have identical checksums for i in range(1, dist.get_world_size()): assert gathered_layer1[i] == gathered_layer1[0], ( f"Layer1 mismatch: rank {i} has {gathered_layer1[i]}, " f"rank 0 has {gathered_layer1[0]}" ) assert gathered_layer2[i] == gathered_layer2[0], ( f"Layer2 mismatch: rank {i} has {gathered_layer2[i]}, " f"rank 0 has {gathered_layer2[0]}" ) else: dist.gather_object(layer1_checksum, None, dst=0) dist.gather_object(layer2_checksum, None, dst=0) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_compression_with_offload(): """Test distributed compression with offloaded modules.""" model = TwoLayerModel() setup_quantized_model(model) # Offload model to CPU offload_module(model.layer1, onload_device="cuda", offload_device="cpu") offload_module(model.layer2, onload_device="cuda", offload_device="cpu") q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) # Compress the model compressor.compress_model(model) # Verify compression happened even with offloading assert hasattr(model.layer1, "weight_packed") assert hasattr(model.layer2, "weight_packed") @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_compression_decompress_roundtrip(): """Test that distributed compression + decompression preserves values.""" model = TwoLayerModel() setup_quantized_model(model) # Store original weights original_layer1 = model.layer1.weight.data.clone() original_layer2 = model.layer2.weight.data.clone() q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) # Compress and decompress compressor.compress_model(model) dist.barrier() compressor.decompress_model(model) dist.barrier() # Weights should be back to float assert model.layer1.weight.dtype == torch.float32 assert model.layer2.weight.dtype == torch.float32 # Values should be close (within quantization error) diff1 = torch.abs(original_layer1 - model.layer1.weight.data) diff2 = torch.abs(original_layer2 - model.layer2.weight.data) assert torch.max(diff1) < 1.0 assert torch.max(diff2) < 1.0 @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_compression_many_layers(): """Test distributed compression with many layers to ensure load balancing.""" class ManyLayerModel(nn.Module): def __init__(self, num_layers=10): super().__init__() self.layers = nn.ModuleList( [nn.Linear(10, 10, bias=False) for _ in range(num_layers)] ) def forward(self, x): for layer in self.layers: x = layer(x) return x model = ManyLayerModel(num_layers=10) setup_quantized_model(model) q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) # Compress the model compressor.compress_model(model) dist.barrier() # All layers should be compressed for layer in model.layers: assert hasattr(layer, "weight_packed") assert layer.weight_packed.dtype == torch.int32 @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_compression_skips_non_quantized(): """Test that non-quantized layers are skipped in distributed compression.""" class MixedModel(nn.Module): def __init__(self): super().__init__() self.quantized = nn.Linear(10, 10, bias=False) self.non_quantized = nn.Linear(10, 10, bias=False) def forward(self, x): x = self.quantized(x) x = self.non_quantized(x) return x model = MixedModel() # Only quantize one layer scheme = QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( num_bits=4, strategy="channel", symmetric=True, type="int", ), ) model.quantized.quantization_scheme = scheme model.quantized.quantization_status = QuantizationStatus.FROZEN model.quantized.weight_scale = nn.Parameter( torch.ones(model.quantized.weight.shape[0], 1) * 0.01 ) model.quantized.weight_zero_point = nn.Parameter( torch.zeros(model.quantized.weight.shape[0], 1, dtype=torch.int32), requires_grad=False, ) q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) # Store original non-quantized weight dtype original_dtype = model.non_quantized.weight.dtype # Compress the model compressor.compress_model(model) # Quantized layer should be compressed assert hasattr(model.quantized, "weight_packed") # Non-quantized layer should remain unchanged assert model.non_quantized.weight.dtype == original_dtype assert not hasattr(model.non_quantized, "weight_packed") @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_compression_empty_model(): """Test distributed compression with an empty model.""" model = nn.Sequential() q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) # Should not raise an error compressor.compress_model(model) dist.barrier() @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_compression_single_layer(): """Test distributed compression with a single layer.""" class SingleLayerModel(nn.Module): def __init__(self): super().__init__() self.layer = nn.Linear(10, 10, bias=False) model = SingleLayerModel() setup_quantized_model(model) q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) # Compress the model compressor.compress_model(model) dist.barrier() # Layer should be compressed assert hasattr(model.layer, "weight_packed") # All ranks should have the same state checksum = model.layer.weight_packed.sum().item() if dist.get_rank() == 0: gathered = [None] * dist.get_world_size() dist.gather_object(checksum, gathered, dst=0) for i in range(1, dist.get_world_size()): assert gathered[i] == gathered[0] else: dist.gather_object(checksum, None, dst=0) vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/distributed/test_module_parallel.py000066400000000000000000000306501521257237700343450ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch import torch.distributed as dist import torch.nn as nn from compressed_tensors.distributed import replace_module_parallel from compressed_tensors.offload import offload_module from compressed_tensors.offload.utils import module_size, to_meta from tests.test_offload.conftest import torchrun from tests.testing_utils import requires_gpu class SimpleLinear(nn.Module): """Simple linear module for testing.""" def __init__(self, in_features=10, out_features=10): super().__init__() self.weight = nn.Parameter(torch.randn(out_features, in_features)) self.bias = nn.Parameter(torch.randn(out_features)) class TwoLayerModel(nn.Module): """Model with two linear layers for testing.""" def __init__(self): super().__init__() self.layer1 = SimpleLinear(10, 10) self.layer2 = SimpleLinear(10, 10) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_to_meta(): """Test that to_meta correctly moves module tensors to meta device.""" module = SimpleLinear(5, 5) original_weight = module.weight.data.clone() original_bias = module.bias.data.clone() # Move to meta to_meta(module) # Check that tensors are on meta device assert module.weight.device.type == "meta" assert module.bias.device.type == "meta" # Check that shapes are preserved assert module.weight.shape == original_weight.shape assert module.bias.shape == original_bias.shape @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_basic(): """Test basic replace_module_parallel functionality.""" modules = [SimpleLinear(10, 10) for _ in range(4)] # Track which modules were processed processed_modules = [] def apply_fn(module): processed_modules.append(id(module)) # Simple modification: set weight to rank value module.weight.data.fill_(float(dist.get_rank())) replace_module_parallel(modules, apply_fn, module_size) # All modules should be processed exactly once assert len(processed_modules) == len(modules) assert len(set(processed_modules)) == len(modules) # All modules should have valid weights (not meta) for module in modules: assert module.weight.device.type != "meta" # Each rank should have processed some modules # We can't predict which ones due to bin packing, # but weights should be 0.0 or 1.0 assert module.weight.mean().item() in [0.0, 1.0] @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_with_offload(): """Test replace_module_parallel with offloaded modules.""" modules = [SimpleLinear(10, 10) for _ in range(4)] # Offload modules to CPU for module in modules: offload_module(module, onload_device="cuda", offload_device="cpu") # Track processing processed_count = [0] def apply_fn(module): processed_count[0] += 1 # Verify we can access the module's parameters assert module.weight is not None module.weight.data.fill_(1.0) replace_module_parallel(modules, apply_fn, module_size) # All modules should be processed assert processed_count[0] == len(modules) # All modules should have updated weights for module in modules: # Access weight and check it was updated weight = module.weight if weight.device.type == "cpu": assert torch.allclose(weight, torch.ones_like(weight)) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_state_broadcast(): """Test that state is correctly broadcast across ranks.""" modules = [SimpleLinear(5, 5) for _ in range(2)] # Each rank processes different modules and sets unique values def apply_fn(module): # Set weight to a unique pattern based on rank module.weight.data.fill_(float(dist.get_rank() * 100)) replace_module_parallel(modules, apply_fn, module_size) dist.barrier() # After replace_module_parallel, all ranks should have the same state # Check that weights are consistent across ranks for i, module in enumerate(modules): # Weight should be either 0 or 100 (depending on which rank processed it) mean_val = module.weight.mean().item() assert mean_val in [0.0, 100.0], f"Module {i} has unexpected weight: {mean_val}" @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_non_processing_ranks_use_meta(): """Test that non-processing ranks temporarily use meta device.""" modules = [SimpleLinear(10, 10) for _ in range(2)] # Track device usage during processing devices_seen = [] def apply_fn(module): # Record the device when this is called devices_seen.append(module.weight.device.type) replace_module_parallel(modules, apply_fn, module_size) # At least one call should see meta on each rank (non-processing modules) assert "meta" in devices_seen # Final state should still be materialized for module in modules: assert module.weight.device.type != "meta" @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_preserves_module_structure(): """Test that module structure is preserved after parallel processing.""" module = SimpleLinear(5, 5) original_weight_shape = module.weight.shape original_bias_shape = module.bias.shape def apply_fn(m): # Modify the module but preserve structure m.weight.data.fill_(1.0) m.bias.data.fill_(0.5) replace_module_parallel([module], apply_fn, module_size) # Check structure is preserved assert module.weight.shape == original_weight_shape assert module.bias.shape == original_bias_shape assert hasattr(module, "weight") assert hasattr(module, "bias") @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_empty_list(): """Test replace_module_parallel with empty module list.""" modules = [] call_count = [0] def apply_fn(module): call_count[0] += 1 # Should not raise an error replace_module_parallel(modules, apply_fn, module_size) # Function should never be called assert call_count[0] == 0 @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_single_module(): """Test replace_module_parallel with a single module.""" module = SimpleLinear(10, 10) def apply_fn(m): m.weight.data.fill_(42.0) replace_module_parallel([module], apply_fn, module_size) # Module should be processed assert module.weight.mean().item() == 42.0 @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_many_modules(): """Test replace_module_parallel with many modules.""" modules = [SimpleLinear(10, 10) for _ in range(20)] processed = set() def apply_fn(m): processed.add(id(m)) m.weight.data.fill_(1.0) replace_module_parallel(modules, apply_fn, module_size) # All modules should be processed assert len(processed) == len(modules) # All modules should have updated weights for module in modules: assert module.weight.mean().item() == 1.0 @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_custom_weight_function(): """Test replace_module_parallel with custom weight function.""" modules = [SimpleLinear(10, 10) for _ in range(4)] # Custom weight function that returns constant weights def constant_weight_fn(m): return 1.0 # All modules have equal weight processed = [] def apply_fn(m): processed.append(id(m)) replace_module_parallel(modules, apply_fn, constant_weight_fn) # All modules should still be processed assert len(processed) == len(modules) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_exception_handling(): """Test that exceptions in apply_fn are properly propagated.""" module = SimpleLinear(10, 10) def failing_apply_fn(m): raise ValueError("Test exception") # The exception should be raised with pytest.raises(ValueError, match="Test exception"): replace_module_parallel([module], failing_apply_fn, module_size) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_parameter_replacement(): """Test replace_module_parallel when parameters are replaced.""" module = SimpleLinear(10, 10) original_weight_shape = module.weight.shape def replace_weight_fn(m): # Replace weight with a new tensor of different values new_weight = torch.ones_like(m.weight) * 99.0 m.weight = nn.Parameter(new_weight) replace_module_parallel([module], replace_weight_fn, module_size) # Check that the weight was replaced assert module.weight.shape == original_weight_shape assert module.weight.mean().item() == 99.0 @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_adds_new_parameters(): """Test replace_module_parallel when new parameters are added.""" module = SimpleLinear(10, 10) def add_parameter_fn(m): # Add a new parameter m.new_param = nn.Parameter(torch.ones(5, 5)) replace_module_parallel([module], add_parameter_fn, module_size) # Check that the new parameter exists on all ranks assert hasattr(module, "new_param") assert module.new_param.shape == (5, 5) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_removes_parameters(): """Test replace_module_parallel when parameters are removed.""" module = SimpleLinear(10, 10) def remove_bias_fn(m): # Remove the bias parameter delattr(m, "bias") replace_module_parallel([module], remove_bias_fn, module_size) # Check that bias is removed on all ranks assert not hasattr(module, "bias") assert hasattr(module, "weight") @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_with_buffers(): """Test replace_module_parallel with modules that have buffers.""" class ModuleWithBuffer(nn.Module): def __init__(self): super().__init__() self.weight = nn.Parameter(torch.randn(5, 5)) self.register_buffer("running_mean", torch.zeros(5)) module = ModuleWithBuffer() def update_buffer_fn(m): m.running_mean.fill_(1.0) m.weight.data.fill_(2.0) replace_module_parallel([module], update_buffer_fn, module_size) # Check that both parameter and buffer are updated assert module.weight.mean().item() == 2.0 assert module.running_mean.mean().item() == 1.0 @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_to_meta_preserves_parameter_properties(): """Test that to_meta preserves parameter properties like requires_grad.""" module = SimpleLinear(5, 5) module.weight.requires_grad = False to_meta(module) # Properties should be preserved assert module.weight.requires_grad is False assert isinstance(module.weight, nn.Parameter) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replace_module_parallel_rank_consistency(): """Test that all ranks see the same final state.""" modules = [SimpleLinear(5, 5) for _ in range(4)] def apply_fn(m): # Fill with rank-specific value m.weight.data.fill_(float(dist.get_rank())) replace_module_parallel(modules, apply_fn, module_size) dist.barrier() # Collect checksums from all ranks checksums = [] for module in modules: checksum = module.weight.sum().item() checksums.append(checksum) # Gather all checksums to rank 0 if dist.get_rank() == 0: gathered = [None] * dist.get_world_size() dist.gather_object(checksums, gathered, dst=0) # All ranks should have the same checksums for rank_checksums in gathered[1:]: assert rank_checksums == gathered[0], "Ranks have different states!" else: dist.gather_object(checksums, None, dst=0) vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/model_compressors/000077500000000000000000000000001521257237700310045ustar00rootroot00000000000000test_model_compressor.py000066400000000000000000000346031521257237700357200ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/model_compressors# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import tempfile from pathlib import Path import torch import torch.nn as nn from compressed_tensors import ModelCompressor from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import ( QuantizationArgs, QuantizationConfig, QuantizationScheme, QuantizationStatus, ) from compressed_tensors.transform import TransformConfig class DummyLinear(nn.Module): """Simple linear module for testing.""" def __init__(self, in_features=10, out_features=10): super().__init__() self.linear = nn.Linear(in_features, out_features, bias=False) def forward(self, x): return self.linear(x) class TwoLayerModel(nn.Module): """Model with two linear layers for testing.""" def __init__(self): super().__init__() self.layer1 = nn.Linear(10, 10, bias=False) self.layer2 = nn.Linear(10, 10, bias=False) def forward(self, x): x = self.layer1(x) x = self.layer2(x) return x def create_quantization_config( bits=8, type="int", strategy="tensor", format="int-quantized" ): """Helper to create a QuantizationConfig for testing.""" config_dict = { "format": format, "global_compression_ratio": 1.0, "quant_method": "compressed-tensors", "config_groups": { "group_0": { "targets": ["Linear"], "weights": { "num_bits": bits, "strategy": strategy, "symmetric": True, "type": type, }, } }, } return QuantizationConfig.model_validate(config_dict) def create_quantization_scheme( bits=8, type="int", strategy="tensor", format=None, with_input_activations=False ): """Helper to create a QuantizationScheme for testing.""" input_activations = None if with_input_activations: input_activations = QuantizationArgs( num_bits=bits, strategy=strategy, symmetric=True, type=type, ) return QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( num_bits=bits, strategy=strategy, symmetric=True, type=type, ), input_activations=input_activations, format=format, ) class TestModelCompressorCompression: """Test compression and decompression functionality.""" def test_compress_model_basic(self): """Test basic model compression.""" model = DummyLinear() # Use pack-quantized format (requires 4 or 8 bits, int type, # no input activations) scheme = create_quantization_scheme(bits=4, type="int", strategy="channel") model.linear.quantization_scheme = scheme model.linear.quantization_status = QuantizationStatus.FROZEN # Set up quantization parameters model.linear.weight_scale = nn.Parameter( torch.ones(model.linear.weight.shape[0], 1) * 0.01 ) model.linear.weight_zero_point = nn.Parameter( torch.zeros(model.linear.weight.shape[0], 1, dtype=torch.int32), requires_grad=False, ) q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) compressor.compress_model(model) # Weight should be packed into int32 assert hasattr(model.linear, "weight_packed") assert model.linear.weight_packed.dtype == torch.int32 assert hasattr(model, "ct_decompress_hook") def test_compress_model_skips_non_quantized_modules(self): """Test that modules without quantization_scheme are skipped.""" model = TwoLayerModel() # Only quantize layer1 scheme = create_quantization_scheme(bits=4, type="int", strategy="channel") model.layer1.quantization_scheme = scheme model.layer1.quantization_status = QuantizationStatus.FROZEN model.layer1.weight_scale = nn.Parameter( torch.ones(model.layer1.weight.shape[0], 1) * 0.01 ) # Use .int() instead of dtype=torch.int32 to avoid Parameter issues model.layer1.weight_zero_point = nn.Parameter( torch.zeros(model.layer1.weight.shape[0], 1).int(), requires_grad=False ) q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) layer2_original_dtype = model.layer2.weight.dtype compressor.compress_model(model) # layer1 should be compressed assert hasattr(model.layer1, "weight_packed") # layer2 should remain unchanged assert model.layer2.weight.dtype == layer2_original_dtype def test_compress_decompress_roundtrip(self): """Test that compress then decompress recovers original weights.""" model = DummyLinear() scheme = create_quantization_scheme(bits=4, type="int", strategy="channel") model.linear.quantization_scheme = scheme model.linear.quantization_status = QuantizationStatus.FROZEN model.linear.weight_scale = nn.Parameter( torch.ones(model.linear.weight.shape[0], 1) * 0.01 ) model.linear.weight_zero_point = nn.Parameter( torch.zeros(model.linear.weight.shape[0], 1, dtype=torch.int32), requires_grad=False, ) q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) original_weight = model.linear.weight.data.clone() compressor.compress_model(model) compressor.decompress_model(model) # After decompression, weight should be back to float assert model.linear.weight.dtype == torch.float32 # Values should be close (within quantization error for 4-bit) diff = torch.abs(original_weight - model.linear.weight.data) assert torch.max(diff) < 1.0 # Reasonable threshold for 4-bit quantization def test_decompress_hook_triggers_on_forward(self): """Test that the decompress hook is triggered on forward pass.""" model = DummyLinear() scheme = create_quantization_scheme(bits=4, type="int", strategy="channel") model.linear.quantization_scheme = scheme model.linear.quantization_status = QuantizationStatus.FROZEN model.linear.weight_scale = nn.Parameter( torch.ones(model.linear.weight.shape[0], 1) * 0.01 ) model.linear.weight_zero_point = nn.Parameter( torch.zeros(model.linear.weight.shape[0], 1, dtype=torch.int32), requires_grad=False, ) q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) compressor.compress_model(model) assert hasattr(model.linear, "weight_packed") assert hasattr(model, "ct_decompress_hook") # Forward pass should trigger decompression x = torch.randn(2, 10) _ = model(x) # After forward, weight should be decompressed assert model.linear.weight.dtype == torch.float32 assert not hasattr(model.linear, "weight_packed") def test_compress_model_updates_format_in_config(self): """Test that from_pretrained_model infers and sets the format correctly.""" model = DummyLinear() scheme = create_quantization_scheme(bits=4, type="int", strategy="channel") model.linear.quantization_scheme = scheme model.linear.quantization_status = QuantizationStatus.FROZEN model.linear.weight_scale = nn.Parameter( torch.ones(model.linear.weight.shape[0], 1) * 0.01 ) model.linear.weight_zero_point = nn.Parameter( torch.zeros(model.linear.weight.shape[0], 1, dtype=torch.int32), requires_grad=False, ) # Use from_pretrained_model which infers the format compressor = ModelCompressor.from_pretrained_model(model) # Format should be inferred as pack-quantized assert compressor.quantization_config.format == CompressionFormat.pack_quantized class TestModelCompressorConfigUpdate: """Test config file update functionality.""" def test_update_config_creates_file(self): """Test that update_config creates config.json if it doesn't exist.""" with tempfile.TemporaryDirectory() as tmpdir: q_config = create_quantization_config() compressor = ModelCompressor(quantization_config=q_config) compressor.update_config(tmpdir) config_path = Path(tmpdir) / "config.json" assert config_path.exists() with open(config_path, "r") as f: config_data = json.load(f) assert "quantization_config" in config_data assert ( config_data["quantization_config"]["quant_method"] == "compressed-tensors" ) def test_update_config_preserves_existing_data(self): """Test that update_config preserves existing config data.""" with tempfile.TemporaryDirectory() as tmpdir: config_path = Path(tmpdir) / "config.json" existing_data = {"model_type": "test", "hidden_size": 768} with open(config_path, "w") as f: json.dump(existing_data, f) q_config = create_quantization_config() compressor = ModelCompressor(quantization_config=q_config) compressor.update_config(tmpdir) with open(config_path, "r") as f: config_data = json.load(f) # Original data should be preserved assert config_data["model_type"] == "test" assert config_data["hidden_size"] == 768 # New data should be added assert "quantization_config" in config_data def test_update_config_with_transform_config(self): """Test that update_config includes transform_config.""" from compressed_tensors.transform import TransformArgs, TransformScheme with tempfile.TemporaryDirectory() as tmpdir: q_config = create_quantization_config() t_config = TransformConfig( config_groups={ "group_0": TransformScheme( type="hadamard", apply=[ TransformArgs(targets=["Linear"], location="weight_input") ], ) } ) compressor = ModelCompressor( quantization_config=q_config, transform_config=t_config ) compressor.update_config(tmpdir) config_path = Path(tmpdir) / "config.json" with open(config_path, "r") as f: config_data = json.load(f) assert "quantization_config" in config_data assert "transform_config" in config_data["quantization_config"] def test_update_config_no_configs(self): """Test that update_config does nothing when no configs are present.""" with tempfile.TemporaryDirectory() as tmpdir: compressor = ModelCompressor() compressor.update_config(tmpdir) config_path = Path(tmpdir) / "config.json" # Should not create file if no configs assert not config_path.exists() def test_update_config_includes_version(self): """Test that update_config includes compression version.""" import compressed_tensors with tempfile.TemporaryDirectory() as tmpdir: q_config = create_quantization_config() compressor = ModelCompressor(quantization_config=q_config) compressor.update_config(tmpdir) config_path = Path(tmpdir) / "config.json" with open(config_path, "r") as f: config_data = json.load(f) assert "quantization_config" in config_data # The version field is added by update_config qc = config_data["quantization_config"] assert "version" in qc, f"Keys in qc: {qc.keys()}" assert qc["version"] == compressed_tensors.__version__ class TestModelCompressorEdgeCases: """Test edge cases and error handling.""" def test_compress_model_infers_format(self): """Test that compression infers format when not set on scheme.""" model = DummyLinear() scheme = QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( num_bits=4, type="int", strategy="channel", symmetric=True ), ) # No format set - will be inferred model.linear.quantization_scheme = scheme model.linear.quantization_status = QuantizationStatus.FROZEN model.linear.weight_scale = nn.Parameter( torch.ones(model.linear.weight.shape[0], 1) * 0.01 ) model.linear.weight_zero_point = nn.Parameter( torch.zeros(model.linear.weight.shape[0], 1, dtype=torch.int32), requires_grad=False, ) q_config = create_quantization_config(bits=4) compressor = ModelCompressor(quantization_config=q_config) compressor.compress_model(model) # Format should be inferred and set assert model.linear.quantization_scheme.format is not None assert hasattr(model.linear, "weight_packed") def test_empty_model(self): """Test compression of a model with no quantized modules.""" model = nn.Sequential() q_config = create_quantization_config() compressor = ModelCompressor(quantization_config=q_config) # Should not raise, just do nothing compressor.compress_model(model) compressor.decompress_model(model) def test_model_with_no_quantization_scheme(self): """Test that modules without quantization_scheme are skipped.""" model = TwoLayerModel() # Don't add quantization_scheme to either layer q_config = create_quantization_config() compressor = ModelCompressor(quantization_config=q_config) original_dtype = model.layer1.weight.dtype compressor.compress_model(model) # Should skip all modules assert model.layer1.weight.dtype == original_dtype assert model.layer2.weight.dtype == original_dtype test_model_compressor_distributed.py000066400000000000000000000256071521257237700403260ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/model_compressors# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for distributed compression in ModelCompressor.""" import pytest import torch import torch.distributed as dist import torch.nn as nn from compressed_tensors import ModelCompressor from compressed_tensors.quantization import ( QuantizationArgs, QuantizationConfig, QuantizationScheme, QuantizationStatus, ) from tests.test_offload.conftest import torchrun from tests.testing_utils import requires_gpu def create_quantization_config(bits=4, format="pack-quantized"): """Helper to create a QuantizationConfig for testing.""" config_dict = { "format": format, "global_compression_ratio": 1.0, "quant_method": "compressed-tensors", "config_groups": { "group_0": { "targets": ["Linear"], "weights": { "num_bits": bits, "strategy": "channel", "symmetric": True, "type": "int", }, } }, } return QuantizationConfig.model_validate(config_dict) def setup_quantized_module(module: nn.Linear, bits: int = 4): """Set up a linear module with quantization scheme and parameters.""" scheme = QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( num_bits=bits, strategy="channel", symmetric=True, type="int", ), ) module.quantization_scheme = scheme module.quantization_status = QuantizationStatus.FROZEN module.weight_scale = nn.Parameter(torch.ones(module.weight.shape[0], 1) * 0.01) module.weight_zero_point = nn.Parameter( torch.zeros(module.weight.shape[0], 1, dtype=torch.int32), requires_grad=False, ) class TwoLayerModel(nn.Module): """Simple model for testing.""" def __init__(self): super().__init__() self.layer1 = nn.Linear(10, 10, bias=False) self.layer2 = nn.Linear(10, 10, bias=False) def forward(self, x): x = self.layer1(x) x = self.layer2(x) return x @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_compress_model_distributed_basic(): """Test basic distributed compression via ModelCompressor.""" model = TwoLayerModel() setup_quantized_module(model.layer1) setup_quantized_module(model.layer2) q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) # Compress the model compressor.compress_model(model) # Verify compression happened assert hasattr(model.layer1, "weight_packed") assert hasattr(model.layer2, "weight_packed") assert model.layer1.weight_packed.dtype == torch.int32 assert model.layer2.weight_packed.dtype == torch.int32 # Verify status updated assert compressor.quantization_config.quantization_status == ( QuantizationStatus.COMPRESSED ) # Verify hook was added assert hasattr(model, "ct_decompress_hook") @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_compress_model_distributed_consistency(): """Test that all ranks have consistent state after distributed compression.""" model = TwoLayerModel() setup_quantized_module(model.layer1) setup_quantized_module(model.layer2) q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) # Compress the model compressor.compress_model(model) dist.barrier() # Compute checksums layer1_sum = model.layer1.weight_packed.sum().item() layer2_sum = model.layer2.weight_packed.sum().item() # Gather checksums from all ranks if dist.get_rank() == 0: gathered_layer1 = [None] * dist.get_world_size() gathered_layer2 = [None] * dist.get_world_size() dist.gather_object(layer1_sum, gathered_layer1, dst=0) dist.gather_object(layer2_sum, gathered_layer2, dst=0) # All ranks should have identical checksums for i in range(1, dist.get_world_size()): assert ( gathered_layer1[i] == gathered_layer1[0] ), f"Layer1 mismatch between rank {i} and rank 0" assert ( gathered_layer2[i] == gathered_layer2[0] ), f"Layer2 mismatch between rank {i} and rank 0" else: dist.gather_object(layer1_sum, None, dst=0) dist.gather_object(layer2_sum, None, dst=0) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_compress_model_distributed_no_quantized_modules(): """Test distributed compression with no quantized modules.""" model = TwoLayerModel() # Don't set up quantization schemes q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) # Should not raise compressor.compress_model(model) # Should not have compressed weights assert not hasattr(model.layer1, "weight_packed") assert not hasattr(model.layer2, "weight_packed") @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_compress_model_distributed_partial_quantization(): """Test distributed compression with only some modules quantized.""" model = TwoLayerModel() setup_quantized_module(model.layer1) # Don't quantize layer2 q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) compressor.compress_model(model) # Only layer1 should be compressed assert hasattr(model.layer1, "weight_packed") assert not hasattr(model.layer2, "weight_packed") @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_compress_decompress_distributed_roundtrip(): """Test compress then decompress in distributed mode.""" model = TwoLayerModel() setup_quantized_module(model.layer1) setup_quantized_module(model.layer2) # Store original weights original_layer1 = model.layer1.weight.data.clone() original_layer2 = model.layer2.weight.data.clone() q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) # Compress and decompress compressor.compress_model(model) dist.barrier() compressor.decompress_model(model) dist.barrier() # Weights should be back to float assert model.layer1.weight.dtype == torch.float32 assert model.layer2.weight.dtype == torch.float32 # Hook should be removed assert not hasattr(model, "ct_decompress_hook") # Status should be updated assert compressor.quantization_config.quantization_status == ( QuantizationStatus.DECOMPRESSED ) # Values should be close (within quantization error) diff1 = torch.abs(original_layer1 - model.layer1.weight.data) diff2 = torch.abs(original_layer2 - model.layer2.weight.data) assert torch.max(diff1) < 1.0 assert torch.max(diff2) < 1.0 @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_compress_model_distributed_many_layers(): """Test distributed compression with many layers for load balancing.""" class ManyLayerModel(nn.Module): def __init__(self, num_layers=10): super().__init__() self.layers = nn.ModuleList( [nn.Linear(10, 10, bias=False) for _ in range(num_layers)] ) model = ManyLayerModel(num_layers=10) for layer in model.layers: setup_quantized_module(layer) q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) compressor.compress_model(model) dist.barrier() # All layers should be compressed for layer in model.layers: assert hasattr(layer, "weight_packed") assert layer.weight_packed.dtype == torch.int32 # Check consistency across ranks checksums = [layer.weight_packed.sum().item() for layer in model.layers] if dist.get_rank() == 0: gathered = [None] * dist.get_world_size() dist.gather_object(checksums, gathered, dst=0) # All ranks should have identical checksums for i in range(1, dist.get_world_size()): for layer_idx, (c1, c2) in enumerate(zip(gathered[0], gathered[i])): assert c1 == c2, f"Layer {layer_idx} mismatch between ranks 0 and {i}" else: dist.gather_object(checksums, None, dst=0) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_compress_model_distributed_force_format(): """Test that force_compression_format works in distributed mode.""" model = TwoLayerModel() setup_quantized_module(model.layer1) setup_quantized_module(model.layer2) q_config = create_quantization_config(bits=4, format="pack-quantized") # Force a specific format compressor = ModelCompressor( quantization_config=q_config, force_compression_format="pack-quantized" ) compressor.compress_model(model) # Verify compression with forced format assert hasattr(model.layer1, "weight_packed") assert hasattr(model.layer2, "weight_packed") @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_compress_model_distributed_from_pretrained(): """Test from_pretrained_model entrypoint works with distributed compression.""" model = TwoLayerModel() setup_quantized_module(model.layer1) setup_quantized_module(model.layer2) # Use from_pretrained_model compressor = ModelCompressor.from_pretrained_model(model) # Should infer format assert compressor.quantization_config is not None # Compress in distributed mode compressor.compress_model(model) dist.barrier() # Verify compression assert hasattr(model.layer1, "weight_packed") assert hasattr(model.layer2, "weight_packed") @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_compress_model_distributed_hook_triggers(): """Test that decompression hook triggers correctly in distributed mode.""" model = TwoLayerModel() setup_quantized_module(model.layer1) setup_quantized_module(model.layer2) q_config = create_quantization_config(bits=4, format="pack-quantized") compressor = ModelCompressor(quantization_config=q_config) compressor.compress_model(model) dist.barrier() # Verify compressed assert hasattr(model.layer1, "weight_packed") assert hasattr(model, "ct_decompress_hook") # Forward pass should trigger decompression x = torch.randn(2, 10) _ = model(x) # After forward, should be decompressed assert model.layer1.weight.dtype == torch.float32 assert not hasattr(model.layer1, "weight_packed") assert not hasattr(model, "ct_decompress_hook") test_transformers_integration.py000066400000000000000000000135011521257237700374660ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/model_compressors# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.compressors import ModelCompressor from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme from compressed_tensors.quantization.lifecycle.initialize import ( initialize_module_for_quantization, ) from compressed_tensors.quantization.quant_args import FP8_E4M3_DATA from tests.testing_utils import requires_gpu from transformers import AutoModelForCausalLM, AutoTokenizer @pytest.mark.parametrize( "frozen_stub,q_format,compressed_stub", [ ( "nm-testing/llama2.c-stories42M-gsm8k-quantized-only-uncompressed", "float-quantized", "nm-testing/llama2.c-stories42M-gsm8k-quantized-only-compressed", ), ( "nm-testing/llama2.c-stories15M-ultrachat-mixed-uncompressed", "pack-quantized", "nm-testing/llama2.c-stories15M-ultrachat-mixed-compressed", ), ], ) def test_compress_model(frozen_stub, q_format, compressed_stub): """Check that compression generates the expected compressed model""" model = AutoModelForCausalLM.from_pretrained(frozen_stub, torch_dtype=torch.float32) compressor = ModelCompressor.from_pretrained_model(model, None, q_format) true_compressed_model = AutoModelForCausalLM.from_pretrained( compressed_stub, torch_dtype=torch.float32 ) compressor.compress_model(model) compressed = dict(model.state_dict()) true_compressed = dict(true_compressed_model.state_dict()) assert compressed.keys() == true_compressed.keys() for key in compressed.keys(): assert compressed[key].dtype == true_compressed[key].dtype, key assert torch.equal(compressed[key], true_compressed[key]) @pytest.mark.filterwarnings("ignore::UserWarning") @pytest.mark.parametrize( "model_stub,q_format,compressed_stub", [ ( "nm-testing/llama2.c-stories42M-gsm8k-quantized-only-uncompressed", "float-quantized", "nm-testing/llama2.c-stories42M-gsm8k-quantized-only-compressed", ), ( "nm-testing/llama2.c-stories15M-ultrachat-mixed-uncompressed", "pack-quantized", "nm-testing/llama2.c-stories15M-ultrachat-mixed-compressed", ), ], ) def test_decompress_model(model_stub, q_format, compressed_stub): from transformers.utils.quantization_config import CompressedTensorsConfig model = AutoModelForCausalLM.from_pretrained(model_stub, torch_dtype=torch.float32) compressor = ModelCompressor.from_pretrained_model(model, None, q_format) true_decompressed_model = AutoModelForCausalLM.from_pretrained( compressed_stub, quantization_config=CompressedTensorsConfig(run_compressed=False), torch_dtype=torch.float32, ) compressor.compress_model(model) compressor.decompress_model(model) true_decompressed = dict(true_decompressed_model.state_dict()) decompressed = dict(model.state_dict()) assert decompressed.keys() == true_decompressed.keys() for key in decompressed.keys(): assert decompressed[key].dtype == true_decompressed[key].dtype, key assert torch.equal(decompressed[key], true_decompressed[key]) def test_multiple_quant_compressors(): model = torch.nn.Sequential(torch.nn.Linear(1, 2), torch.nn.Linear(2, 3)) input_activations = QuantizationArgs(num_bits=8, type="float") weights = QuantizationArgs(num_bits=8, type="float") scheme_fp8 = QuantizationScheme( targets=["Linear"], weights=weights, input_activations=input_activations, format=CompressionFormat.float_quantized.value, ) input_activations = QuantizationArgs( num_bits=4, type="float", scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=FP8_E4M3_DATA.dtype, ) weights = QuantizationArgs( num_bits=4, type="float", scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=FP8_E4M3_DATA.dtype, ) scheme_nvfp4 = QuantizationScheme( targets=["Linear"], weights=weights, input_activations=input_activations, format=CompressionFormat.nvfp4_pack_quantized.value, ) model[0].quantization_scheme = scheme_fp8 initialize_module_for_quantization(model[0]) model[0].quantization_status = "frozen" model[1].quantization_scheme = scheme_nvfp4 initialize_module_for_quantization(model[1]) model[1].quantization_status = "frozen" compressor = ModelCompressor.from_pretrained_model(model, None) compressor.compress_model(model) assert compressor.quantization_config.format == CompressionFormat.mixed_precision assert model[0].quantization_scheme.format == scheme_fp8.format assert model[1].quantization_scheme.format == scheme_nvfp4.format @requires_gpu def test_compressed_model_inference_with_hook(): model_stub = "nm-testing/llama2.c-stories42M-gsm8k-quantized-only-compressed" # Load compressed model model = AutoModelForCausalLM.from_pretrained( model_stub, dtype="auto", device_map="auto" ) tokenizer = AutoTokenizer.from_pretrained(model_stub) # Model should have the decompression hook attached assert hasattr(model, "ct_decompress_hook") # Run a forward pass to trigger the hook prompt = "The quick brown fox jumps over the lazy dog" inputs = tokenizer(prompt, return_tensors="pt") input_ids = inputs["input_ids"].to(device=model.device) with torch.no_grad(): outputs = model(input_ids=input_ids, labels=input_ids) # After forward pass, hook should have triggered and been removed assert not hasattr(model, "ct_decompress_hook") # Check perplexity is reasonable assert torch.exp(outputs.loss) <= 500.0 vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/test_compress_decompress_module.py000066400000000000000000000153641521257237700343130ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch import torch.nn as nn from compressed_tensors.compressors import ModelCompressor from compressed_tensors.compressors.base import compress_module, decompress_module from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import ( ActivationOrdering, QuantizationArgs, QuantizationConfig, QuantizationScheme, apply_quantization_config, initialize_module_for_quantization, preset_name_to_scheme, ) from compressed_tensors.quantization.utils import is_module_quantized from compressed_tensors.utils import get_direct_state_dict def _run_compress_decompress( scheme_name, expected_format, actorder, device, module=None, targets=("Linear",) ): if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA not available") # 1. Initialize module for quantization using a preset scheme. # Use 256x256 to avoid degenerate scale shapes (e.g. (N,1) for FP8_BLOCK) # that trip up block-vs-channel strategy inference in dequantize. # Start in bfloat16 so NVFP4 decompression (which returns bfloat16) preserves dtype. if module is None: module = nn.Linear(256, 256, bias=False) module = module.to(dtype=torch.bfloat16, device=device) scheme = preset_name_to_scheme(scheme_name, list(targets)) if actorder is not None: scheme.weights.actorder = actorder initialize_module_for_quantization(module, scheme) with torch.no_grad(): for name, param in list(module.named_parameters()): param.fill_(1) # Record pre-compression state dict shapes and dtypes. # Filter out None entries (e.g. bias=None when bias=False). pre_state = { name: (tensor.shape, tensor.dtype) for name, tensor in get_direct_state_dict(module).items() if tensor is not None } # 2. Compress the module and verify the inferred quantization format. compress_module(module) assert module.quantization_scheme.format == expected_format # 3. Decompress the module and verify shapes and dtypes are restored. decompress_module(module) post_state_dict = get_direct_state_dict(module) for name, tensor in post_state_dict.items(): if name in pre_state: pre_shape, pre_dtype = pre_state[name] assert tensor.shape == pre_shape assert tensor.dtype == pre_dtype @pytest.mark.parametrize( "scheme_name,expected_format,actorder", [ ("UNQUANTIZED", CompressionFormat.dense, None), ("W8A16", CompressionFormat.pack_quantized, None), ("W4A16", CompressionFormat.pack_quantized, None), ("W4A16", CompressionFormat.pack_quantized, ActivationOrdering.GROUP), ("W4A16_ASYM", CompressionFormat.pack_quantized, None), ("W4A16_ASYM", CompressionFormat.pack_quantized, ActivationOrdering.GROUP), ("W8A8", CompressionFormat.int_quantized, None), ("W4A8", CompressionFormat.int_quantized, None), ("W4AFP8", CompressionFormat.int_quantized, None), ("FP8", CompressionFormat.float_quantized, None), ("FP8_DYNAMIC", CompressionFormat.float_quantized, None), ("FP8_BLOCK", CompressionFormat.float_quantized, None), ("NVFP4A16", CompressionFormat.nvfp4_pack_quantized, None), ("NVFP4", CompressionFormat.nvfp4_pack_quantized, None), ], ) @pytest.mark.parametrize("device", ["cpu", "meta", "cuda"]) def test_compress_decompress_module(scheme_name, expected_format, actorder, device): _run_compress_decompress(scheme_name, expected_format, actorder, device) @pytest.mark.parametrize( "scheme_name,expected_format", [ ("MXFP4A16", CompressionFormat.mxfp4_pack_quantized), ("MXFP4", CompressionFormat.mxfp4_pack_quantized), ], ) @pytest.mark.parametrize("device", ["cpu", "meta", "cuda"]) def test_compress_decompress_module_mxfp4(scheme_name, expected_format, device): _run_compress_decompress(scheme_name, expected_format, None, device) @pytest.mark.parametrize( "scheme_name,expected_format,actorder", [ ("UNQUANTIZED", CompressionFormat.dense, None), ("W8A16", CompressionFormat.pack_quantized, None), ("W4A16", CompressionFormat.pack_quantized, None), ("W4A16", CompressionFormat.pack_quantized, ActivationOrdering.GROUP), ("W4A16_ASYM", CompressionFormat.pack_quantized, None), ("NVFP4A16", CompressionFormat.nvfp4_pack_quantized, None), ("MXFP4A16", CompressionFormat.mxfp4_pack_quantized, None), ], ) @pytest.mark.parametrize("device", ["cpu", "meta", "cuda"]) def test_compress_decompress_embedding(scheme_name, expected_format, actorder, device): # Embeddings are compressed the same way as Linear weights: weight-only # (weight-and-activation schemes don't apply since embeddings consume indices). module = nn.Embedding(256, 256) _run_compress_decompress( scheme_name, expected_format, actorder, device, module=module, targets=("Embedding",), ) def test_linear_only_config_leaves_embedding_untouched(): # Embedding compression is opt-in: a module is only compressed if it has a # quantization_scheme attached, which apply_quantization_config does only for # matched targets. A Linear-only config must leave embeddings fully untouched. class TinyModel(nn.Module): def __init__(self): super().__init__() self.embed = nn.Embedding(64, 128) self.proj = nn.Linear(128, 64, bias=False) model = TinyModel() embed_weight_before = model.embed.weight.detach().clone() config = QuantizationConfig( config_groups={ "group_0": QuantizationScheme( targets=["Linear"], weights=QuantizationArgs(num_bits=4, symmetric=True), ) } ) apply_quantization_config(model, config) # Only the Linear is targeted; the Embedding gets no scheme. assert is_module_quantized(model.proj) assert not is_module_quantized(model.embed) assert not hasattr(model.embed, "quantization_scheme") ModelCompressor.from_pretrained_model(model).compress_model(model) # Linear is compressed (weight replaced by packed params)... proj_keys = set(get_direct_state_dict(model.proj).keys()) assert "weight_packed" in proj_keys assert "weight" not in proj_keys # ...but the Embedding is byte-for-byte unchanged: no packed params, no # status attribute, original weight preserved. embed_keys = set(get_direct_state_dict(model.embed).keys()) assert embed_keys == {"weight"} assert not hasattr(model.embed, "quantization_status") assert torch.equal(model.embed.weight, embed_weight_before) vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/test_fp4_quant.py000066400000000000000000000046201521257237700305610ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.compressors.nvfp4.base import NVFP4PackedCompressor from compressed_tensors.compressors.nvfp4.helpers import ( pack_fp4_to_uint8, unpack_fp4_from_uint8, ) from compressed_tensors.quantization import QuantizationArgs, QuantizationType def test_pack_unpack(): x = torch.Tensor( [ [-0.5000, -6.0000, -0.5000, -1.5000, -1.0000, 6.0000, 0.0000, -0.0000], [-1.0000, -6.0000, -0.5000, -0.0000, 0.5000, 0.5000, -0.0000, 0.0000], [-3.0000, -6.0000, -0.5000, -2.0000, -0.5000, -1.5000, -0.0000, -0.0000], [1.5000, 6.0000, -0.0000, -0.5000, 1.0000, 1.0000, -0.0000, 0.0000], ] ) dense_dtype = torch.bfloat16 x = x.to(dense_dtype) m, n = x.shape packed = pack_fp4_to_uint8(x) assert packed.dtype == torch.uint8 unpacked = unpack_fp4_from_uint8(packed, m, n, dtype=dense_dtype) assert unpacked.dtype == dense_dtype assert torch.equal(unpacked, x) # misleading as -0 and 0 are considered equal sign_bitx = torch.signbit(x) sign_bitout = torch.signbit(unpacked) assert torch.equal(sign_bitout, sign_bitx) def test_pack_unpack_odd_dims(): x = torch.Tensor( [ [-0.5000, -6.0000, -0.5000, -1.5000, -1.0000, 6.0000, 0.0000], [-1.0000, -6.0000, -0.5000, -0.0000, 0.5000, 0.5000, -0.0000], [1.5000, 6.0000, -0.0000, -0.5000, 1.0000, 1.0000, -0.0000], ] ) with pytest.raises((ValueError, torch._dynamo.exc.Unsupported)): _ = pack_fp4_to_uint8(x) def test_compress_scale_without_scale_dtype(): """ Test that NVFP4 compressor handles missing scale_dtype. (backward compatibility) """ # Create a scale tensor scale = torch.randn(10, dtype=torch.bfloat16) # Create QuantizationArgs without scale_dtype (as in older models) quant_args = QuantizationArgs( num_bits=4, type=QuantizationType.FLOAT, symmetric=True, group_size=16, # scale_dtype is not set (defaults to None) ) # This should not raise an error and should default to float8_e4m3fn compressed_scale = NVFP4PackedCompressor._compress_scale(scale, quant_args) # Verify the output dtype is float8_e4m3fn assert compressed_scale.dtype == torch.float8_e4m3fn vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/test_fp8_quant.py000066400000000000000000000122541521257237700305670ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math from collections import OrderedDict import pytest import torch from compressed_tensors import FloatQuantizationCompressor from compressed_tensors.quantization import ( QuantizationArgs, QuantizationConfig, QuantizationScheme, QuantizationStatus, QuantizationStrategy, apply_quantization_config, ) from compressed_tensors.quantization.lifecycle.forward import fake_quantize from torch.nn.modules import Linear, Sequential def make_dummy_g_idx(columns: int, group_size: int) -> torch.Tensor: perm = torch.randperm(columns) return torch.tensor([index // group_size for index in range(columns)])[perm] @pytest.mark.parametrize( "strategy,group_size,sc,zp", [ [QuantizationStrategy.TENSOR, None, torch.tensor(0.01), torch.tensor(0)], [ QuantizationStrategy.GROUP, 128, torch.rand((512, 8)) * 0.01, torch.zeros((512, 8), dtype=torch.int8), ], [ QuantizationStrategy.CHANNEL, None, torch.rand((512, 1)) * 0.01, torch.zeros((512, 1), dtype=torch.int8), ], ], ) def test_quant_format(strategy, group_size, sc, zp): module_sd = { "weight": torch.rand((512, 1024)), "weight_scale": sc.to(torch.float32), "weight_zero_point": zp.to(torch.float32), } if group_size is not None: module_sd["weight_g_idx"] = make_dummy_g_idx(1024, group_size) scheme = QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( strategy=strategy, type="float", group_size=group_size ), ) compressed = FloatQuantizationCompressor.compress(module_sd, scheme=scheme) # symmetric → zero_point is dropped assert "weight_zero_point" not in compressed assert compressed["weight_scale"].dtype == torch.float32 assert torch.equal(compressed["weight_scale"], module_sd["weight_scale"]) if group_size is not None: assert torch.equal(compressed["weight_g_idx"], module_sd["weight_g_idx"]) @pytest.mark.parametrize( "strategy,group_size", [ [QuantizationStrategy.TENSOR, None], [QuantizationStrategy.CHANNEL, None], ], ) def test_compress_decompress_match( mock_per_group_calibration, mock_per_channel_calibration, strategy, group_size, ): model = Sequential(OrderedDict([("dummy", Linear(512, 1024, bias=None))])) quant_config = QuantizationConfig( config_groups={ "group_1": QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( strategy=strategy, type="float", group_size=group_size ), ) }, ignore=["lm_head"], ) apply_quantization_config(model, quant_config) model.dummy.quantization_status = QuantizationStatus.CALIBRATION if strategy == QuantizationStrategy.GROUP: mock_per_group_calibration( model.dummy, base_name="weight", value=model.dummy.weight, group_size=128 ) if strategy == QuantizationStrategy.CHANNEL: mock_per_channel_calibration( model.dummy, base_name="weight", value=model.dummy.weight ) scheme = quant_config.config_groups["group_1"] # Build per-module state dict module_sd = { name: param.data.clone() for name, param in model.dummy.named_parameters() } compressed = FloatQuantizationCompressor.compress(module_sd, scheme=scheme) decompressed = FloatQuantizationCompressor.decompress(compressed, scheme=scheme) fake_quant_dummy = fake_quantize( model.dummy.weight, scale=model.dummy.weight_scale, zero_point=model.dummy.weight_zero_point, args=scheme.weights, ) assert torch.equal(fake_quant_dummy, decompressed["weight"]) @pytest.mark.parametrize( "rows,cols,block_height,block_width", [ (10944, 2048, 128, 128), (2048, 10944, 128, 128), (256, 256, 128, 128), (300, 400, 128, 128), (256, 300, 128, 128), (300, 256, 128, 128), ], ) def test_block_quant_compression_padding(rows, cols, block_height, block_width): """ Block quantization compresses weights with non-divisible dimensions without changing the original shape. """ block_structure = [block_height, block_width] num_rb = math.ceil(rows / block_height) num_cb = math.ceil(cols / block_width) module_sd = { "weight": torch.rand((rows, cols)), "weight_scale": torch.rand((num_rb, num_cb)) * 0.01 + 0.001, "weight_zero_point": torch.zeros((num_rb, num_cb)), } scheme = QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( strategy=QuantizationStrategy.BLOCK, type="float", block_structure=block_structure, ), ) compressed = FloatQuantizationCompressor.compress(module_sd, scheme=scheme) # Compressed weight should retain the original shape assert compressed["weight"].shape == ( rows, cols, ), "Compressed weight shape should match original" vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/test_int_quant.py000066400000000000000000000062221521257237700306620ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors import IntQuantizationCompressor from compressed_tensors.quantization import ( QuantizationArgs, QuantizationScheme, QuantizationStrategy, ) from compressed_tensors.quantization.lifecycle.forward import fake_quantize def make_quant_scheme(strategy, group_size=None, symmetric=True): return QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( strategy=strategy, group_size=group_size, symmetric=symmetric ), ) @pytest.mark.parametrize( "strategy,symmetric,group_size,sc,zp", [ [QuantizationStrategy.TENSOR, True, None, torch.tensor(0.01), torch.tensor(0)], [ QuantizationStrategy.GROUP, True, 128, torch.rand((512, 8)) * 0.01, torch.zeros((512, 8), dtype=torch.int8), ], [ QuantizationStrategy.CHANNEL, False, None, torch.rand((512, 1)) * 0.01, ((torch.rand((512, 1)) - 0.5) * 127).to(torch.int8), ], ], ) def test_quant_format(strategy, symmetric, group_size, sc, zp): module_sd = { "weight": torch.rand((512, 1024)), "weight_scale": sc.to(torch.float32), "weight_zero_point": zp.to(torch.int32), } scheme = make_quant_scheme( strategy=strategy, group_size=group_size, symmetric=symmetric ) compressed = IntQuantizationCompressor.compress(module_sd, scheme=scheme) # zero_point is dropped for symmetric quantization if symmetric: assert "weight_zero_point" not in compressed else: assert compressed["weight_zero_point"].dtype == torch.int32 # weight should be compressed to int8 assert compressed["weight"].dtype == torch.int8 assert compressed["weight_scale"].dtype == torch.float32 @pytest.mark.parametrize( "strategy,group_size,sc,zp", [ [QuantizationStrategy.TENSOR, None, torch.tensor(0.01), torch.tensor(0)], [ QuantizationStrategy.GROUP, 128, torch.rand((300, 8)) * 0.01, torch.zeros((300, 8), dtype=torch.int8), ], [ QuantizationStrategy.CHANNEL, None, torch.rand((300, 1)) * 0.01, torch.zeros((300, 1), dtype=torch.int8), ], ], ) def test_compress_decompress_match(strategy, group_size, sc, zp): module_sd = { "weight": torch.rand((300, 1024)), "weight_scale": sc.to(torch.float32), "weight_zero_point": zp.to(torch.int32), } scheme = make_quant_scheme(strategy=strategy, group_size=group_size) compressed = IntQuantizationCompressor.compress(module_sd, scheme=scheme) decompressed = IntQuantizationCompressor.decompress(compressed, scheme=scheme) fake_quant = fake_quantize( module_sd["weight"], scale=module_sd["weight_scale"], zero_point=module_sd["weight_zero_point"], args=scheme.weights, ) assert torch.equal(fake_quant, decompressed["weight"].to(torch.float32)) vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/test_mxfp4_quant.py000066400000000000000000000052641521257237700311330ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.compressors.mxfp4.base import MXFP4PackedCompressor from compressed_tensors.compressors.nvfp4.helpers import pack_fp4_to_uint8 from compressed_tensors.quantization import ( QuantizationArgs, QuantizationScheme, QuantizationType, ) def test_compress_scale_without_scale_dtype(): """ Test that MXFP4 compressor handles missing scale_dtype. (backward compatibility) """ # Create a scale tensor scale = torch.randn(10, dtype=torch.bfloat16).abs() + 1e-6 # Ensure positive values # Create QuantizationArgs without scale_dtype (as in older models) quant_args = QuantizationArgs( num_bits=4, type=QuantizationType.FLOAT, symmetric=True, group_size=32, # scale_dtype is not set (defaults to None) ) # This should not raise an error and should default to uint8 compressed_scale = MXFP4PackedCompressor._compress_scale(scale, quant_args) # Verify the output dtype is uint8 assert compressed_scale.dtype == torch.uint8 def test_compress_scale_with_scale_dtype(): """Test that MXFP4 compressor respects explicit scale_dtype""" # Create a scale tensor scale = torch.randn(10, dtype=torch.bfloat16).abs() + 1e-6 # Ensure positive values # Create QuantizationArgs with explicit scale_dtype quant_args = QuantizationArgs( num_bits=4, type=QuantizationType.FLOAT, symmetric=True, group_size=32, scale_dtype=torch.uint8, ) # Compress the scale compressed_scale = MXFP4PackedCompressor._compress_scale(scale, quant_args) # Verify the output dtype matches the specified scale_dtype assert compressed_scale.dtype == torch.uint8 def test_decompress_decodes_mx_scales_and_restores_weight(): quant_args = QuantizationArgs( num_bits=4, type=QuantizationType.FLOAT, symmetric=True, group_size=32, scale_dtype=torch.uint8, ) scale = torch.tensor([[0.25, 0.5]], dtype=torch.bfloat16) packed = pack_fp4_to_uint8( torch.tensor([[0.5, 1.0, 1.5, 2.0]], dtype=torch.bfloat16) ) decompressed = MXFP4PackedCompressor.decompress( { "weight_packed": packed, "weight_scale": MXFP4PackedCompressor._compress_scale(scale, quant_args), }, QuantizationScheme(targets=["Linear"], weights=quant_args), ) expected_weight = torch.tensor([[0.125, 0.25, 0.75, 1.0]], dtype=torch.bfloat16) assert torch.equal(decompressed["weight_scale"], scale) assert torch.equal(decompressed["weight"], expected_weight) vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/test_mxfp8_quant.py000066400000000000000000000130161521257237700311310ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.compressors.mxfp8 import MXFP8QuantizationCompressor from compressed_tensors.quantization import ( QuantizationArgs, QuantizationScheme, QuantizationStrategy, ) from compressed_tensors.quantization.utils.helpers import calculate_qparams def test_mxfp8_compress_decompress(): """ Test MXFP8 compress/decompress round-trip with group strategy and group_size=32. Verifies weights survive the cycle (lossy but close). """ rows, cols = 512, 1024 group_size = 32 num_groups = cols // group_size quant_args = QuantizationArgs( num_bits=8, type="float", strategy=QuantizationStrategy.GROUP, group_size=group_size, scale_dtype=torch.uint8, zp_dtype=torch.uint8, symmetric=True, ) weight = torch.randn((rows, cols)) # Compute scales using calculate_qparams (which generates MX scales) reshaped = weight.reshape(rows, num_groups, group_size) min_vals = reshaped.amin(dim=-1) max_vals = reshaped.amax(dim=-1) scale, zp = calculate_qparams(min_vals, max_vals, quant_args) scheme = QuantizationScheme( targets=["Linear"], weights=quant_args, ) # Build per-module state dict (local names, no module prefix) module_sd = { "weight": weight, "weight_scale": scale, "weight_zero_point": zp, } # Compress compressed = MXFP8QuantizationCompressor.compress(module_sd, scheme=scheme) # Check compressed weight is FP8 assert compressed["weight"].dtype == torch.float8_e4m3fn # Check scale is stored as uint8 (E8M0 exponent format) assert compressed["weight_scale"].dtype == torch.uint8 # Decompress decompressed = MXFP8QuantizationCompressor.decompress(compressed, scheme=scheme) # Check shapes match assert decompressed["weight"].shape == weight.shape # FP8 quantization is lossy, but should be reasonably close assert torch.allclose(decompressed["weight"].float(), weight, atol=0.1, rtol=0.1) def test_mxfp8_scale_roundtrip(): """ Test that E8M0 scale encoding/decoding is lossless for power-of-2 scales. """ rows, cols = 128, 256 group_size = 32 num_groups = cols // group_size quant_args = QuantizationArgs( num_bits=8, type="float", strategy=QuantizationStrategy.GROUP, group_size=group_size, scale_dtype=torch.uint8, zp_dtype=torch.uint8, symmetric=True, ) weight = torch.randn((rows, cols)) reshaped = weight.reshape(rows, num_groups, group_size) min_vals = reshaped.amin(dim=-1) max_vals = reshaped.amax(dim=-1) scale, zp = calculate_qparams(min_vals, max_vals, quant_args) scheme = QuantizationScheme( targets=["Linear"], weights=quant_args, ) module_sd = { "weight": weight, "weight_scale": scale, "weight_zero_point": zp, } compressed = MXFP8QuantizationCompressor.compress(module_sd, scheme=scheme) # E8M0 encoded scale e8m0_scale = compressed["weight_scale"] assert e8m0_scale.dtype == torch.uint8 # Decode: 2^(exp - 127) scale_exp = e8m0_scale.to(torch.int32) - 127 decoded_scale = 2.0 ** scale_exp.to(torch.float32) # The original scale after floor(log2) should round-trip exactly expected_scale = 2.0 ** torch.floor(torch.log2(scale)).to(torch.float32) assert torch.allclose(decoded_scale, expected_scale) def test_mxfp8_can_compress(): """Test that can_compress matches MXFP8 signature correctly.""" import torch.nn as nn mxfp8_scheme = QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( num_bits=8, type="float", strategy=QuantizationStrategy.GROUP, group_size=32, scale_dtype=torch.uint8, ), ) assert MXFP8QuantizationCompressor.can_compress(nn.Linear, mxfp8_scheme) is True # Non-MXFP8: group_size != 32 non_mxfp8_scheme = QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( num_bits=8, type="float", strategy=QuantizationStrategy.GROUP, group_size=128, ), ) assert ( MXFP8QuantizationCompressor.can_compress(nn.Linear, non_mxfp8_scheme) is False ) # Non-MXFP8: int type int_scheme = QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( num_bits=8, type="int", strategy=QuantizationStrategy.GROUP, group_size=32, scale_dtype=torch.uint8, ), ) assert MXFP8QuantizationCompressor.can_compress(nn.Linear, int_scheme) is False def test_compress_scale_without_scale_dtype(): """ Test that MXFP8 compressor handles missing scale_dtype. (backward compatibility) """ # Create a scale tensor scale = torch.randn(10, dtype=torch.bfloat16).abs() + 1e-6 # Ensure positive values # Create QuantizationArgs without scale_dtype (as in older models) quant_args = QuantizationArgs( num_bits=8, type="float", symmetric=True, group_size=32, # scale_dtype is not set (defaults to None) ) # This should not raise an error and should default to uint8 compressed_scale = MXFP8QuantizationCompressor._compress_scale(scale, quant_args) # Verify the output dtype is uint8 assert compressed_scale.dtype == torch.uint8 vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/test_pack_quant.py000066400000000000000000000377051521257237700310200ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math from collections import OrderedDict import pytest import torch from compressed_tensors import PackedQuantizationCompressor from compressed_tensors.compressors.pack_quantized.helpers import ( pack_to_int32, unpack_from_int32, ) from compressed_tensors.quantization import ( QuantizationArgs, QuantizationConfig, QuantizationScheme, QuantizationStatus, QuantizationStrategy, apply_quantization_config, ) from compressed_tensors.quantization.lifecycle.forward import fake_quantize from compressed_tensors.quantization.quant_args import ActivationOrdering from torch.nn.modules import Linear, Sequential def get_dummy_quant_config( num_bits=4, strategy=None, group_size=None, actorder=None, symmetric=True ) -> QuantizationConfig: config_groups = { "group_1": QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( num_bits=num_bits, strategy=strategy, group_size=group_size, actorder=actorder, symmetric=symmetric, ), ), } return QuantizationConfig(config_groups=config_groups) def make_dummy_g_idx(columns: int, group_size: int) -> torch.Tensor: perm = torch.randperm(columns) return torch.nn.Parameter( (torch.arange(columns, dtype=torch.int) // group_size)[perm], requires_grad=False, ) @pytest.mark.parametrize( "shape", [ (512, 1024), (830, 545), (342, 512), (256, 700), ], ) def test_quant_format(shape): module_sd = { "weight": torch.rand(shape), "weight_scale": torch.tensor(0.01, dtype=torch.float32), "weight_zero_point": torch.tensor(0, dtype=torch.int8), } scheme = QuantizationScheme( targets=["Linear"], weights=QuantizationArgs(num_bits=4, symmetric=True) ) compressed = PackedQuantizationCompressor.compress(module_sd, scheme=scheme) # 'weight' replaced by 'weight_packed' + 'weight_shape'; zp dropped (symmetric) assert "weight" not in compressed assert "weight_packed" in compressed assert "weight_shape" in compressed assert "weight_zero_point" not in compressed assert compressed["weight_packed"].dtype == torch.int32 expected_rows = shape[0] expected_columns = math.ceil(shape[1] / 8) assert compressed["weight_packed"].shape == (expected_rows, expected_columns) assert torch.equal(compressed["weight_shape"], torch.tensor(shape)) assert compressed["weight_scale"].dtype == torch.float32 @pytest.mark.parametrize( "value", [ torch.tensor([[1, 2], [3, 4]]), torch.tensor([[1, 2, 3, 4, 5, 6, 7, 0], [-1, -2, -3, -4, -5, -6, -7, -8]]), (torch.rand((32, 100)) * 16 - 8), ], ) def test_repack_4bit(value): value = value.to(torch.int8) shape = value.shape assert not torch.any(value > 7).item() assert not torch.any(value < -8).item() packed = pack_to_int32(value, 4) unpacked = unpack_from_int32(packed, 4, shape) assert torch.equal(value, unpacked) @pytest.mark.parametrize( "value", [ torch.tensor([[30, 40], [50, 60]]), torch.tensor( [[10, 15, 20, 25, 30, 35, 40, 45], [-10, -20, -30, -40, -50, -60, -70, -80]] ), (torch.rand((32, 100)) * 256 - 128), ], ) def test_repack_8bit(value): value = value.to(torch.int8) shape = value.shape assert not torch.any(value > 127).item() assert not torch.any(value < -128).item() packed = pack_to_int32(value, 8) unpacked = unpack_from_int32(packed, 8, shape) assert torch.equal(value, unpacked) @pytest.mark.parametrize("num_bits", [1, 2, 3, 4, 5, 6, 7, 8]) @pytest.mark.parametrize("shape", [(256, 1024), (512, 100), (128, 33)]) def test_pack_unpack_roundtrip(num_bits, shape): """Pack/unpack roundtrip preserves values for all supported bit widths.""" lo, hi = -(1 << (num_bits - 1)), (1 << (num_bits - 1)) - 1 value = torch.randint(lo, hi + 1, shape, dtype=torch.int8) packed = pack_to_int32(value, num_bits) assert packed.dtype == torch.int32 pack_factor = 32 // num_bits assert packed.shape == (shape[0], math.ceil(shape[1] / pack_factor)) unpacked = unpack_from_int32(packed, num_bits, torch.Size(shape)) assert torch.equal(unpacked, value) @pytest.mark.parametrize("num_bits", [1, 2, 3, 4, 5, 6, 7, 8]) def test_compress_decompress_match(num_bits): """Round-trip compress → decompress in memory.""" module_sd = { "weight": torch.rand((511, 350)), "weight_scale": torch.tensor(0.01, dtype=torch.float32), "weight_zero_point": torch.tensor(0, dtype=torch.int8), } scheme = QuantizationScheme( targets=["Linear"], weights=QuantizationArgs(num_bits=num_bits, symmetric=False), ) compressed = PackedQuantizationCompressor.compress(module_sd.copy(), scheme=scheme) decompressed = PackedQuantizationCompressor.decompress(compressed, scheme=scheme) fake_quant = fake_quantize( module_sd["weight"], scale=module_sd["weight_scale"], zero_point=module_sd["weight_zero_point"], args=scheme.weights, ) assert torch.equal(fake_quant, decompressed["weight"].to(torch.float32)) @pytest.mark.parametrize( "strategy", {QuantizationStrategy.GROUP, QuantizationStrategy.CHANNEL}, ) def test_asymmetric_packed_support(strategy): shape = (1024, 1024) group_size = None if strategy == QuantizationStrategy.GROUP: group_size = 128 if strategy == QuantizationStrategy.CHANNEL: expected_shape = (shape[0], 1) elif strategy == QuantizationStrategy.GROUP: num_groups = shape[1] // group_size expected_shape = (shape[0], max(num_groups, 1)) module_sd = { "weight": torch.rand(shape), "weight_scale": torch.rand(expected_shape).to(torch.float32), "weight_zero_point": torch.rand(expected_shape).to(torch.int8), } scheme = QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( num_bits=4, strategy=strategy.value, symmetric=False, group_size=group_size ), ) compressed = PackedQuantizationCompressor.compress(module_sd, scheme=scheme) # weight + shape entry + packed zp assert "weight_packed" in compressed assert "weight_shape" in compressed assert "weight_zero_point" in compressed assert compressed["weight_packed"].dtype == torch.int32 assert compressed["weight_zero_point"].dtype == torch.int32 assert compressed["weight_scale"].dtype == torch.float32 expected_rows = shape[0] expected_columns = math.ceil(shape[1] / 8) assert compressed["weight_packed"].shape == (expected_rows, expected_columns) assert torch.equal(compressed["weight_shape"], torch.tensor(shape)) packed_size_zp = math.ceil(shape[0] / 8) zp_factor = group_size if strategy == QuantizationStrategy.GROUP else shape[-1] assert compressed["weight_zero_point"].shape == ( packed_size_zp, shape[-1] // zp_factor, ) @pytest.mark.parametrize( "actorder", [ ActivationOrdering.GROUP, ActivationOrdering.WEIGHT, None, ], ) def test_actorder_compress_decompress_match(actorder, mock_per_group_calibration): model = Sequential(OrderedDict([("dummy", Linear(512, 1024, bias=None))])) group_size = 128 quant_config = get_dummy_quant_config( strategy="group", group_size=group_size, actorder=actorder ) apply_quantization_config(model, quant_config) model.quantization_status = QuantizationStatus.CALIBRATION mock_per_group_calibration( model.dummy, base_name="weight", value=model.dummy.weight, group_size=group_size ) if actorder == ActivationOrdering.GROUP: init_g_idx = make_dummy_g_idx(512, group_size) model.dummy.register_parameter("weight_g_idx", init_g_idx) scheme = quant_config.config_groups["group_1"] module_sd = { name: param.data.clone() for name, param in model.dummy.named_parameters() } compressed = PackedQuantizationCompressor.compress(module_sd, scheme=scheme) decompressed = PackedQuantizationCompressor.decompress(compressed, scheme=scheme) fake_quant = fake_quantize( model.dummy.weight, scale=model.dummy.weight_scale, zero_point=model.dummy.weight_zero_point, g_idx=getattr(model.dummy, "weight_g_idx", None), args=scheme.weights, ) assert torch.equal(fake_quant, decompressed["weight"]) @pytest.mark.parametrize( "num_bits,values,expected_values", [ ( 4, torch.tensor([[1]]), torch.tensor([[9]], dtype=torch.int32), ), ( 8, torch.tensor([[1]]), torch.tensor([[129]], dtype=torch.int32), ), (4, torch.tensor([[1, 2, 3, 4]]), torch.tensor([[52137]], dtype=torch.int32)), ( 4, torch.tensor([[-8, -7, -6, -5, -4, -3, -2, -1]]), torch.tensor([[1985229328]], dtype=torch.int32), ), ( 8, torch.tensor([[1, 2, 3, 4]]), torch.tensor([[-2071756159]], dtype=torch.int32), ), ( 8, torch.tensor([[-128, -127, -126, -125]]), torch.tensor([[50462976]], dtype=torch.int32), ), ( 4, torch.tensor([[-8, -7, -6, -5, -4, -3, -2, -1, 1, 2, 3, 4]]), torch.tensor([[1985229328, 52137]], dtype=torch.int32), ), ( 4, torch.tensor( [ [-8, -7, -6, -5, -4, -3, -2, -1, 1, 2, 3, 4, -8, -8, -8, -8], [1, 2, 3, 4, -8, -8, -8, -8, -8, -7, -6, -5, -4, -3, -2, -1], ] ), torch.tensor([[1985229328, 52137], [52137, 1985229328]], dtype=torch.int32), ), ( 8, torch.tensor([[1, 2, 3, 4], [-128, -127, -126, -125]]), torch.tensor([[-2071756159], [50462976]], dtype=torch.int32), ), ( 8, torch.tensor( [ [1, 2, 3, 4, -128, -127, -126, -125], [-128, -127, -126, -125, 1, 2, 3, 4], ] ), torch.tensor( [[-2071756159, 50462976], [50462976, -2071756159]], dtype=torch.int32 ), ), ], ) def test_pack_to_int32(num_bits, values, expected_values): values = values.to(torch.int8) packed_values = pack_to_int32(values, num_bits) assert torch.equal(packed_values, expected_values) assert packed_values.dtype == expected_values.dtype @pytest.mark.parametrize( "num_bits,values,expected_tensor", [ ( 4, torch.tensor([[9]], dtype=torch.int32), torch.tensor([[1]], dtype=torch.int8), ), ( 8, torch.tensor([[129]], dtype=torch.int32), torch.tensor([[1]], dtype=torch.int8), ), ( 4, torch.tensor([[52137]], dtype=torch.int32), torch.tensor([[1, 2, 3, 4]], dtype=torch.int8), ), ( 4, torch.tensor([[1985229328]], dtype=torch.int32), torch.tensor([[-8, -7, -6, -5, -4, -3, -2, -1]], dtype=torch.int8), ), ( 8, torch.tensor([[-2071756159]], dtype=torch.int32), torch.tensor([[1, 2, 3, 4]], dtype=torch.int8), ), ( 8, torch.tensor([[50462976]], dtype=torch.int32), torch.tensor([[-128, -127, -126, -125]], dtype=torch.int8), ), ( 4, torch.tensor([[1985229328, 52137]], dtype=torch.int32), torch.tensor( [[-8, -7, -6, -5, -4, -3, -2, -1, 1, 2, 3, 4]], dtype=torch.int8 ), ), ( 4, torch.tensor([[1985229328, 52137], [52137, 1985229328]], dtype=torch.int32), torch.tensor( [ [-8, -7, -6, -5, -4, -3, -2, -1, 1, 2, 3, 4, -8, -8, -8, -8], [1, 2, 3, 4, -8, -8, -8, -8, -8, -7, -6, -5, -4, -3, -2, -1], ], dtype=torch.int8, ), ), ( 8, torch.tensor([[-2071756159], [50462976]], dtype=torch.int32), torch.tensor([[1, 2, 3, 4], [-128, -127, -126, -125]], dtype=torch.int8), ), ( 8, torch.tensor( [[-2071756159, 50462976], [50462976, -2071756159]], dtype=torch.int32 ), torch.tensor( [ [1, 2, 3, 4, -128, -127, -126, -125], [-128, -127, -126, -125, 1, 2, 3, 4], ], dtype=torch.int8, ), ), ], ) def test_unpack_from_int32(num_bits, values, expected_tensor): unpacked_tensor = unpack_from_int32(values, num_bits, expected_tensor.shape) assert torch.equal(unpacked_tensor, expected_tensor) assert unpacked_tensor.dtype == expected_tensor.dtype @pytest.mark.parametrize( "strategy,group_size", [ (QuantizationStrategy.GROUP, 128), (QuantizationStrategy.CHANNEL, None), ], ) def test_asymmetric_zero_point_decompression(strategy, group_size): shape = (512, 1024) if strategy == QuantizationStrategy.CHANNEL: expected_zp_shape = (shape[0], 1) elif strategy == QuantizationStrategy.GROUP: num_groups = shape[1] // group_size expected_zp_shape = (shape[0], max(num_groups, 1)) module_sd = { "weight": torch.randn(shape), "weight_scale": torch.rand(expected_zp_shape).to(torch.float32), "weight_zero_point": torch.randint(-8, 8, expected_zp_shape).to(torch.int8), } scheme = QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( num_bits=4, strategy=strategy.value, symmetric=False, group_size=group_size ), ) compressed = PackedQuantizationCompressor.compress(module_sd.copy(), scheme=scheme) assert "weight_zero_point" in compressed assert compressed["weight_zero_point"].dtype == torch.int32 decompressed = PackedQuantizationCompressor.decompress(compressed, scheme=scheme) assert "weight" in decompressed assert decompressed["weight"].shape == shape @pytest.mark.parametrize( "num_bits,strategy", [ (4, QuantizationStrategy.GROUP), (4, QuantizationStrategy.CHANNEL), (8, QuantizationStrategy.GROUP), (8, QuantizationStrategy.CHANNEL), ], ) def test_zero_point_pack_unpack_consistency(num_bits, strategy): if strategy == QuantizationStrategy.GROUP: shape = (512, 8) else: shape = (512, 1) max_val = (1 << (num_bits - 1)) - 1 min_val = -(1 << (num_bits - 1)) original_zp = torch.randint(min_val, max_val + 1, shape).to(torch.int8) packed_zp = pack_to_int32(original_zp, num_bits, packed_dim=0) unpacked_zp = unpack_from_int32(packed_zp, num_bits, shape, packed_dim=0) assert torch.equal(original_zp, unpacked_zp) assert unpacked_zp.dtype == torch.int8 def test_pack_unpack_3d_round_trip(): """3D tensors (e.g. MoE expert weights) should pack/unpack correctly.""" num_bits = 4 shape = (4, 8, 32) # (num_experts, rows, cols) value = torch.randint(-8, 7, shape, dtype=torch.int8) packed = pack_to_int32(value, num_bits) unpacked = unpack_from_int32(packed, num_bits, torch.Size(shape)) assert torch.equal(value, unpacked) def test_pack_unpack_3d_matches_stacked_2d(): """3D pack/unpack should match stacking individual 2D results.""" num_bits = 4 shape = (4, 8, 32) value = torch.randint(-8, 7, shape, dtype=torch.int8) packed_3d = pack_to_int32(value, num_bits) packed_2d = torch.stack( [pack_to_int32(value[i], num_bits) for i in range(value.shape[0])] ) assert torch.equal(packed_3d, packed_2d) vllm-project-compressed-tensors-c18a0fa/tests/test_compressors/test_packed_asym_decompression.py000066400000000000000000000125371521257237700341000ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ End-to-end tests for asymmetric quantization with zero-point decompression. """ import pytest import torch from compressed_tensors.compressors.model_compressors.model_compressor import ( ModelCompressor, ) from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import ( QuantizationArgs, QuantizationConfig, QuantizationScheme, QuantizationStrategy, apply_quantization_config, ) from torch.nn import Linear, Module class SimpleModel(Module): """Simple model for testing""" def __init__(self, input_dim=512, hidden_dim=256, output_dim=128): super().__init__() self.layer1 = Linear(input_dim, hidden_dim, bias=False) self.layer2 = Linear(hidden_dim, output_dim, bias=False) def forward(self, x): x = self.layer1(x) x = torch.relu(x) x = self.layer2(x) return x def create_asymmetric_quant_config( num_bits=4, strategy=QuantizationStrategy.GROUP, group_size=128 ) -> QuantizationConfig: """Create an asymmetric quantization config""" config_groups = { "group_1": QuantizationScheme( targets=["Linear"], weights=QuantizationArgs( num_bits=num_bits, strategy=strategy.value, group_size=( group_size if strategy == QuantizationStrategy.GROUP else None ), symmetric=False, ), ), } return QuantizationConfig(config_groups=config_groups) @pytest.mark.parametrize( "strategy,group_size", [ (QuantizationStrategy.GROUP, 128), (QuantizationStrategy.CHANNEL, None), ], ) def test_end_to_end_asymmetric_quantization( strategy, group_size, mock_per_group_calibration, mock_per_channel_calibration, ): """ Test end-to-end workflow: quantize -> compress -> decompress in memory """ model = SimpleModel() original_weights = { "layer1": model.layer1.weight.detach().clone(), "layer2": model.layer2.weight.detach().clone(), } quant_config = create_asymmetric_quant_config( num_bits=4, strategy=strategy, group_size=group_size ) # Set pack-quantized format for ModelCompressor usage quant_config.format = CompressionFormat.pack_quantized.value apply_quantization_config(model, quant_config) if strategy == QuantizationStrategy.GROUP: mock_per_group_calibration( model.layer1, "weight", model.layer1.weight, group_size ) mock_per_group_calibration( model.layer2, "weight", model.layer2.weight, group_size ) else: mock_per_channel_calibration(model.layer1, "weight", model.layer1.weight) mock_per_channel_calibration(model.layer2, "weight", model.layer2.weight) # Compress and decompress in memory using ModelCompressor mc = ModelCompressor(quantization_config=quant_config) mc.compress_model(model) # Verify compression created zero-point parameters assert hasattr(model.layer1, "weight_zero_point") assert hasattr(model.layer2, "weight_zero_point") assert model.layer1.weight_zero_point.dtype == torch.int32 assert model.layer2.weight_zero_point.dtype == torch.int32 # Decompress in memory mc.decompress_model(model) # Verify decompression restored weights correctly assert model.layer1.weight.shape == original_weights["layer1"].shape assert model.layer2.weight.shape == original_weights["layer2"].shape assert model.layer1.weight.dtype.is_floating_point assert model.layer2.weight.dtype.is_floating_point assert not torch.isnan(model.layer1.weight).any() assert not torch.isnan(model.layer2.weight).any() assert not torch.isinf(model.layer1.weight).any() assert not torch.isinf(model.layer2.weight).any() @pytest.mark.parametrize("num_bits", [4, 8]) def test_asymmetric_quantization_accuracy(num_bits, mock_per_group_calibration): """ Test that asymmetric quantization with zero-point preserves accuracy better than symmetric quantization for biased weight distributions. """ shape = (256, 512) biased_weights = torch.randn(shape) + 2.0 quant_config = create_asymmetric_quant_config( num_bits=num_bits, strategy=QuantizationStrategy.GROUP, group_size=128, ) quant_config.format = CompressionFormat.pack_quantized.value class SingleLayer(Module): def __init__(self): super().__init__() self.layer = Linear(shape[1], shape[0], bias=False) model = SingleLayer() apply_quantization_config(model, quant_config) with torch.no_grad(): model.layer.weight.copy_(biased_weights) mock_per_group_calibration(model.layer, "weight", model.layer.weight, 128) # Compress and decompress in memory using ModelCompressor mc = ModelCompressor(quantization_config=quant_config) mc.compress_model(model) mc.decompress_model(model) decompressed_weights = model.layer.weight assert decompressed_weights.shape == shape assert not torch.isnan(decompressed_weights).any() assert not torch.isinf(decompressed_weights).any() threshold = torch.std(torch.rand(shape) - torch.rand(shape)) assert torch.std(biased_weights - decompressed_weights) < threshold vllm-project-compressed-tensors-c18a0fa/tests/test_configs/000077500000000000000000000000001521257237700243165ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_configs/__init__.py000066400000000000000000000001531521257237700264260ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project vllm-project-compressed-tensors-c18a0fa/tests/test_configs/test_base.py000066400000000000000000000033171521257237700266450ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest from compressed_tensors.config import SparsityStructure def test_sparsity_structure_valid_cases(): assert ( SparsityStructure("2:4") == SparsityStructure.TWO_FOUR ), "Failed to match '2:4' with TWO_FOUR" assert ( SparsityStructure("unstructured") == SparsityStructure.UNSTRUCTURED ), "Failed to match 'unstructured' with UNSTRUCTURED" assert ( SparsityStructure("UNSTRUCTURED") == SparsityStructure.UNSTRUCTURED ), "Failed to match 'UNSTRUCTURED' with UNSTRUCTURED" assert ( SparsityStructure(None) == SparsityStructure.UNSTRUCTURED ), "Failed to match None with UNSTRUCTURED" def test_sparsity_structure_invalid_case(): with pytest.raises(ValueError, match="invalid is not a valid SparsityStructure"): SparsityStructure("invalid") def test_sparsity_structure_case_insensitivity(): assert ( SparsityStructure("2:4") == SparsityStructure.TWO_FOUR ), "Failed to match '2:4' with TWO_FOUR" assert ( SparsityStructure("2:4".upper()) == SparsityStructure.TWO_FOUR ), "Failed to match '2:4'.upper() with TWO_FOUR" assert ( SparsityStructure("unstructured".upper()) == SparsityStructure.UNSTRUCTURED ), "Failed to match 'unstructured'.upper() with UNSTRUCTURED" assert ( SparsityStructure("UNSTRUCTURED".lower()) == SparsityStructure.UNSTRUCTURED ), "Failed to match 'UNSTRUCTURED'.lower() with UNSTRUCTURED" def test_sparsity_structure_default_case(): assert ( SparsityStructure(None) == SparsityStructure.UNSTRUCTURED ), "Failed to match None with UNSTRUCTURED" vllm-project-compressed-tensors-c18a0fa/tests/test_configs/test_infer_quant.py000066400000000000000000000027371521257237700302530ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import OrderedDict import pytest import torch from compressed_tensors.compressors.format import infer_model_format from compressed_tensors.quantization import preset_name_to_scheme @pytest.mark.parametrize( "preset,expected_format", [ ["W8A8", "int-quantized"], ["W8A16", "pack-quantized"], ["W4A16", "pack-quantized"], ["FP8", "float-quantized"], ], ) def test_infer_quant_format(preset, expected_format): quant_scheme = preset_name_to_scheme(preset, targets=["Linear"]) dummy_model = torch.nn.Sequential( OrderedDict( [ ("fc1", torch.nn.Linear(8, 16, bias=True)), ("fc2", torch.nn.Linear(16, 32, bias=True)), ( "block1", torch.nn.Sequential( OrderedDict( [ ("fc1", torch.nn.Linear(32, 16, bias=True)), ("fc2", torch.nn.Linear(16, 8, bias=True)), ] ) ), ), ] ) ) for _, module in dummy_model.named_modules(): if isinstance(module, torch.nn.Linear): module.quantization_scheme = quant_scheme assert infer_model_format(dummy_model).value == expected_format vllm-project-compressed-tensors-c18a0fa/tests/test_entrypoints/000077500000000000000000000000001521257237700252645ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_entrypoints/convert/000077500000000000000000000000001521257237700267445ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_entrypoints/convert/converters/000077500000000000000000000000001521257237700311365ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_entrypoints/convert/converters/test_autoawq.py000066400000000000000000000130411521257237700342270ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.config import CompressionFormat from compressed_tensors.entrypoints.convert import AutoAWQConverter from compressed_tensors.quantization import QuantizationStatus from transformers import AutoConfig def _pack_int4(values: torch.Tensor) -> torch.Tensor: values = values.to(torch.int32) packed = torch.zeros(values.shape[0], values.shape[1] // 8, dtype=torch.int32) for offset in range(8): packed |= values[:, offset::8] << (offset * 4) return packed @pytest.mark.unit def test_unpack_awq_and_reverse_order(): packed_values = torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7]], dtype=torch.int8) qweight = _pack_int4(packed_values) unpacked, _ = AutoAWQConverter.unpack_awq(qweight, None, bits=4) reordered, _ = AutoAWQConverter.reverse_awq_order(unpacked, None, bits=4) assert torch.equal(torch.bitwise_and(unpacked, 15), packed_values) assert torch.equal( torch.bitwise_and(reordered, 15), torch.tensor([[0, 4, 1, 5, 2, 6, 3, 7]], dtype=torch.int8), ) @pytest.mark.unit @pytest.mark.parametrize("zero_point", [True, False]) def test_autoawq_converter_processes_gemm_tensors(zero_point): converter = AutoAWQConverter( group_size=2, targets=[r"re:.*proj$"], zero_point=zero_point, ) qweight_values = torch.tensor( [ [8, 9, 10, 11, 12, 13, 14, 15], [0, 1, 2, 3, 4, 5, 6, 7], ], dtype=torch.int8, ) qzeros_values = torch.tensor([[8, 8, 8, 8, 8, 8, 8, 8]], dtype=torch.int8) scales = torch.ones(1, 8, dtype=torch.float16) tensors = { "model.layers.0.mlp.up_proj.qweight": _pack_int4(qweight_values), "model.layers.0.mlp.up_proj.qzeros": _pack_int4(qzeros_values), "model.layers.0.mlp.up_proj.scales": scales, "model.embed_tokens.weight": torch.ones(4, 4), } if not zero_point: del tensors["model.layers.0.mlp.up_proj.qzeros"] converter.validate(tensors) converter.process(tensors) assert "model.layers.0.mlp.up_proj.qweight" not in tensors assert "model.layers.0.mlp.up_proj.qzeros" not in tensors assert "model.layers.0.mlp.up_proj.scales" not in tensors assert "model.layers.0.mlp.up_proj.weight" not in tensors assert tensors["model.layers.0.mlp.up_proj.weight_packed"].shape == (8, 1) assert torch.equal( tensors["model.layers.0.mlp.up_proj.weight_shape"], torch.tensor([8, 2]) ) assert tensors["model.layers.0.mlp.up_proj.weight_scale"].shape == (8, 1) assert tensors["model.layers.0.mlp.up_proj.weight_scale"].is_contiguous() if zero_point: assert tensors["model.layers.0.mlp.up_proj.weight_zero_point"].shape == (1, 1) else: assert "model.layers.0.mlp.up_proj.weight_zero_point" not in tensors @pytest.mark.unit def test_autoawq_converter_config_from_autoawq_config(): converter = AutoAWQConverter.from_autoawq_config( { "bits": 4, "group_size": 64, "zero_point": True, "version": "gemm", "modules_to_not_convert": ["vision_tower"], }, ) config = converter.create_config() scheme = config.config_groups["config_group_0"] assert config.format == CompressionFormat.pack_quantized.value assert config.quantization_status == QuantizationStatus.COMPRESSED assert config.ignore == ["lm_head", "re:.*vision_tower.*"] assert scheme.format == CompressionFormat.pack_quantized.value assert scheme.weights.num_bits == 4 assert scheme.weights.group_size == 64 assert scheme.weights.symmetric is False @pytest.mark.unit def test_autoawq_converter_from_pretrained(monkeypatch): class MockConfig: quantization_config = { "quant_method": "awq", "bits": 4, "group_size": 32, "zero_point": True, "version": "gemm", "modules_to_not_convert": ["visual"], } def mock_from_pretrained(model_name_or_path, trust_remote_code=False): assert model_name_or_path == "test/model" assert trust_remote_code is True return MockConfig() monkeypatch.setattr(AutoConfig, "from_pretrained", mock_from_pretrained) converter = AutoAWQConverter.from_pretrained( "test/model", trust_remote_code=True, ) assert converter.bits == 4 assert converter.group_size == 32 assert converter.zero_point is True assert converter.version == "gemm" assert converter.ignore == ["lm_head", "re:.*visual.*"] @pytest.mark.unit def test_autoawq_converter_dependencies(): converter = AutoAWQConverter(targets=[r"re:.*down_proj$"]) assert converter.get_dependencies("model.layers.0.mlp.down_proj.qweight") == { "model.layers.0.mlp.down_proj.qzeros", "model.layers.0.mlp.down_proj.scales", } assert converter.get_dependencies("model.layers.0.mlp.up_proj.qweight") == set() symmetric_converter = AutoAWQConverter( targets=[r"re:.*down_proj$"], zero_point=False, ) assert symmetric_converter.get_dependencies( "model.layers.0.mlp.down_proj.qweight" ) == {"model.layers.0.mlp.down_proj.scales"} @pytest.mark.unit def test_autoawq_converter_validate_requires_dependencies(): converter = AutoAWQConverter() with pytest.raises(ValueError, match="without corresponding"): converter.validate( {"model.layers.0.mlp.down_proj.qweight": torch.zeros(1, 1, device="meta")} ) test_build_inverse_weight_maps.py000066400000000000000000000101521521257237700377100ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_entrypoints/convert/converters# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import pytest import torch from compressed_tensors.entrypoints.convert import ( FP8BlockDequantizer, build_inverse_weight_maps, ) from compressed_tensors.utils.safetensors_load import get_checkpoint_files from safetensors.torch import save_file @pytest.mark.unit def test_build_inverse_weight_maps(tmp_path): """ Test that reindex_checkpoint correctly moves tensors across files so that weight and weight_scale_inv end up in the same file. """ # Create dummy checkpoint with weights split across files model_dir = tmp_path / "model" model_dir.mkdir() # File 1: has layer0.weight but NOT layer0.weight_scale_inv file1_tensors = { "embed_tokens.weight": torch.randn(128, 128, dtype=torch.float32), "layer0.weight": torch.randn(128, 128, dtype=torch.float32).to( torch.float8_e4m3fn ), "layer1.weight_scale_inv": torch.randn(1, 1, dtype=torch.float32), } file1_path = model_dir / "model-00001-of-00002.safetensors" save_file(file1_tensors, str(file1_path)) # File 2: has layer0.weight_scale_inv and layer1.weight_scale_inv file2_tensors = { "layer0.weight_scale_inv": torch.randn(1, 1, dtype=torch.float32), "layer1.weight": torch.randn(128, 128, dtype=torch.float32).to( torch.float8_e4m3fn ), "layer2.weight": torch.randn(128, 128, dtype=torch.float32).to( torch.float8_e4m3fn ), "layer2.weight_scale_inv": torch.randn(1, 1, dtype=torch.float32), "lm_head.weight": torch.randn(128, 128, dtype=torch.float32), } file2_path = model_dir / "model-00002-of-00002.safetensors" save_file(file2_tensors, str(file2_path)) # Create index file weight_map = { "embed_tokens.weight": "model-00001-of-00002.safetensors", "layer0.weight": "model-00001-of-00002.safetensors", "layer1.weight": "model-00002-of-00002.safetensors", "layer0.weight_scale_inv": "model-00002-of-00002.safetensors", "layer1.weight_scale_inv": "model-00001-of-00002.safetensors", "layer2.weight": "model-00002-of-00002.safetensors", "layer2.weight_scale_inv": "model-00002-of-00002.safetensors", "lm_head.weight": "model-00002-of-00002.safetensors", } index_data = { "metadata": { "total_size": sum( t.numel() * t.element_size() for tensors in [file1_tensors, file2_tensors] for t in tensors.values() ) }, "weight_map": weight_map, } index_path = model_dir / "model.safetensors.index.json" with open(index_path, "w") as f: json.dump(index_data, f) # Create config.json (required by get_checkpoint_files) config_path = model_dir / "config.json" with open(config_path, "w") as f: json.dump({"model_type": "test"}, f) converter = FP8BlockDequantizer(targets=[r"re:.*layer\d.*"]) inverse_weight_maps = build_inverse_weight_maps( weight_map=weight_map, model_files=get_checkpoint_files(model_dir), converters=[converter], ) for file_name in ( "model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors", ): assert ( file_name in inverse_weight_maps ), f"File {file_name} missing in inverse_weight_maps" seen_weight_names = set() for inverse_weight_map in inverse_weight_maps.values(): for weight_names in inverse_weight_map.values(): for weight_name in weight_names: assert ( weight_name not in seen_weight_names ), f"duplicate weight {weight_name} found" seen_weight_names.add(weight_name) all_weight_names = set(weight_map.keys()) assert ( seen_weight_names >= all_weight_names ), f"Some weights are missing, {all_weight_names - seen_weight_names}" assert ( all_weight_names >= seen_weight_names ), f"Extraneous weights added, {seen_weight_names - all_weight_names}" test_ct_dequantizer.py000066400000000000000000000120731521257237700355140ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_entrypoints/convert/converters# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.config import CompressionFormat from compressed_tensors.entrypoints.convert.converters.ct_dequantizer import ( CompressedTensorsDequantizer, ) from compressed_tensors.quantization import QuantizationConfig, QuantizationScheme from compressed_tensors.quantization.quant_args import ( QuantizationArgs, QuantizationStrategy, QuantizationType, ) def _create_dequantizer(ignore=None): dequantizer = object.__new__(CompressedTensorsDequantizer) dequantizer.dtype = torch.bfloat16 scheme = QuantizationScheme( targets=["re:.*mlp.*"], weights=QuantizationArgs( num_bits=8, type=QuantizationType.INT, strategy=QuantizationStrategy.CHANNEL, symmetric=True, dynamic=False, ), format=CompressionFormat.naive_quantized, ) dequantizer.quant_config = QuantizationConfig( config_groups={"group_0": scheme}, ignore=ignore or [], ) return dequantizer def _create_dummy_tensors(device: torch.device = torch.device("cpu")): with device: return { "model.layers.0.mlp.up_proj.weight": torch.randint( -128, 127, (64, 64), dtype=torch.int8 ), "model.layers.0.mlp.up_proj.weight_scale": torch.rand( 64, 1, dtype=torch.float32 ), "model.layers.0.mlp.down_proj.weight": torch.randint( -128, 127, (64, 64), dtype=torch.int8 ), "model.language_model.layers.0.input_layernorm.weight": torch.randn( 64, 1, dtype=torch.bfloat16 ), "model.language_model.layers.0.pre_feedforward_layernorm.weight": torch.randn( # noqa: E501 64, 1, dtype=torch.bfloat16 ), "model.language_model.layers.0.post_feedforward_layernorm.weight": torch.randn( # noqa: E501 64, 1, dtype=torch.bfloat16 ), "model.layers.0.mlp.down_proj.weight_scale": torch.rand( 64, 1, dtype=torch.float32 ), "model.layers.0.self_attn.q_proj.weight": torch.randn( 128, 64, dtype=torch.bfloat16 ), "model.embed_tokens.weight": torch.randn(128, 64, dtype=torch.bfloat16), } @pytest.mark.unit def test_process_dequantizes_targeted_layers(): dequantizer = _create_dequantizer(ignore=["model.embed_tokens"]) tensors = _create_dummy_tensors() qproj_weight = tensors["model.layers.0.self_attn.q_proj.weight"].clone() embed_tokens_weight = tensors["model.embed_tokens.weight"].clone() result = dequantizer.process(tensors) assert "model.layers.0.mlp.up_proj.weight" in result assert "model.layers.0.mlp.down_proj.weight" in result assert result["model.layers.0.mlp.up_proj.weight"].dtype == torch.bfloat16 assert result["model.layers.0.mlp.down_proj.weight"].dtype == torch.bfloat16 assert "model.layers.0.mlp.up_proj.weight_scale" not in result assert "model.layers.0.mlp.down_proj.weight_scale" not in result assert torch.equal(result["model.layers.0.self_attn.q_proj.weight"], qproj_weight) assert torch.equal(result["model.embed_tokens.weight"], embed_tokens_weight) @pytest.mark.unit def test_validate_passes_with_valid_tensors(): dequantizer = _create_dequantizer(ignore=["model.embed_tokens"]) tensors = _create_dummy_tensors(device=torch.device("meta")) dequantizer.validate(tensors) @pytest.mark.unit def test_validate_raises_on_missing_scale(): dequantizer = _create_dequantizer(ignore=["model.embed_tokens"]) tensors = _create_dummy_tensors(device=torch.device("meta")) del tensors["model.layers.0.mlp.up_proj.weight_scale"] with pytest.raises(ValueError, match="Expected key"): dequantizer.validate(tensors) @pytest.mark.unit def test_validate_raises_on_unconsumed_key(): dequantizer = _create_dequantizer(ignore=["model.embed_tokens"]) tensors = _create_dummy_tensors(device=torch.device("meta")) tensors["model.layers.0.mlp.up_proj.extra_param"] = torch.rand(64) with pytest.raises(ValueError, match="unconsumed keys"): dequantizer.validate(tensors) @pytest.mark.unit def test_get_dependencies_returns_scale_for_targeted_weight(): dequantizer = _create_dequantizer(ignore=["model.embed_tokens"]) deps = dequantizer.get_dependencies("model.layers.0.mlp.up_proj.weight") assert deps == {"model.layers.0.mlp.up_proj.weight_scale"} @pytest.mark.unit def test_get_dependencies_returns_empty_for_non_root_param(): dequantizer = _create_dequantizer(ignore=["model.embed_tokens"]) deps = dequantizer.get_dependencies("model.layers.0.mlp.up_proj.weight_scale") assert deps == set() @pytest.mark.unit def test_get_dependencies_returns_empty_for_ignored_module(): dequantizer = _create_dequantizer(ignore=["model.embed_tokens"]) deps = dequantizer.get_dependencies("model.embed_tokens.weight") assert deps == set() test_fp8block_dequantizer.py000066400000000000000000000177061521257237700366260ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_entrypoints/convert/converters# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.entrypoints.convert import FP8BlockDequantizer @pytest.mark.unit def test_fp8_block_to_bfloat16_conversion(): """ Test that _create_bfloat16_weight correctly converts FP8 block-quantized weights to bfloat16 by multiplying by the scale_inv per block. """ converter = FP8BlockDequantizer(weight_block_size=(128, 128)) # Create a weight tensor divisible by block size (256x256 = 2x2 blocks of 128x128) original_weight = torch.randn(256, 256, dtype=torch.bfloat16) # Simulate block quantization: divide into blocks and create per-block scales num_row_blocks = 2 num_col_blocks = 2 # Create per-block scale_inv (2x2 for 2x2 blocks) weight_scale_inv = torch.randn(num_row_blocks, num_col_blocks, dtype=torch.float32) # Convert original to fp8 (simulate quantization by just converting dtype) weight_fp8 = original_weight.to(torch.float32).to(torch.float8_e4m3fn) # Test conversion result = converter._create_dequantized_weight(weight_fp8, weight_scale_inv) # Verify using helper _verify_block_conversion(result, weight_fp8, weight_scale_inv, (128, 128)) @pytest.mark.unit def test_fp8_block_to_bfloat16_conversion_with_padding(): """ Test that _create_bfloat16_weight correctly handles tensors that need padding (dimensions not evenly divisible by block size). """ converter = FP8BlockDequantizer(weight_block_size=(128, 128)) # Create a weight tensor NOT divisible by block size (200x300) # Should be padded to 256x384 (2x3 blocks) weight_fp8 = torch.randn(200, 300, dtype=torch.float32).to(torch.float8_e4m3fn) # Scale_inv for padded size: 2 row blocks x 3 col blocks num_row_blocks = 2 # ceil(200/128) = 2 num_col_blocks = 3 # ceil(300/128) = 3 weight_scale_inv = torch.ones(num_row_blocks, num_col_blocks, dtype=torch.float32) # Test conversion result = converter._create_dequantized_weight(weight_fp8, weight_scale_inv) # Verify output shape matches original (not padded) assert result.shape == (200, 300), "Output shape should match original, not padded" assert result.dtype == torch.bfloat16, "Output dtype should be bfloat16" @pytest.mark.unit def test_fp8_block_converter_process(): """ Test that the converter's process method correctly converts FP8 block-quantized tensors in a dict to bfloat16, removing weight_scale_inv tensors. """ converter = FP8BlockDequantizer( targets=[r"re:.*layer\d+\.mlp\..*proj$"], weight_block_size=(128, 128) ) # Create mock tensors dict with FP8 weights and scale_inv tensors num_row_blocks = 2 num_col_blocks = 2 # Non-targeted tensor (should not be modified) non_targeted_weight = torch.randn(128, 128, dtype=torch.bfloat16) tensors = { "model.layer0.mlp.up_proj.weight": torch.randn( 256, 256, dtype=torch.float32 ).to(torch.float8_e4m3fn), "model.layer0.mlp.up_proj.weight_scale_inv": torch.randn( num_row_blocks, num_col_blocks, dtype=torch.float32 ), "model.layer1.mlp.down_proj.weight": torch.randn( 256, 256, dtype=torch.float32 ).to(torch.float8_e4m3fn), "model.layer1.mlp.down_proj.weight_scale_inv": torch.randn( num_row_blocks, num_col_blocks, dtype=torch.float32 ), "model.embed_tokens.weight": non_targeted_weight, } # Save references to original tensors before processing weight_fp8_layer0 = tensors["model.layer0.mlp.up_proj.weight"].clone() scale_inv_layer0 = tensors["model.layer0.mlp.up_proj.weight_scale_inv"].clone() weight_fp8_layer1 = tensors["model.layer1.mlp.down_proj.weight"].clone() scale_inv_layer1 = tensors["model.layer1.mlp.down_proj.weight_scale_inv"].clone() # Process the tensors converter.process(tensors) # Verify that weight_scale_inv tensors were removed assert ( "model.layer0.mlp.up_proj.weight_scale_inv" not in tensors ), "weight_scale_inv should be removed" assert ( "model.layer1.mlp.down_proj.weight_scale_inv" not in tensors ), "weight_scale_inv should be removed" # Verify that weights were converted to bfloat16 assert "model.layer0.mlp.up_proj.weight" in tensors, "weight should still exist" assert "model.layer1.mlp.down_proj.weight" in tensors, "weight should still exist" # Verify the conversion is correct using helper _verify_block_conversion( tensors["model.layer0.mlp.up_proj.weight"], weight_fp8_layer0, scale_inv_layer0, (128, 128), ) _verify_block_conversion( tensors["model.layer1.mlp.down_proj.weight"], weight_fp8_layer1, scale_inv_layer1, (128, 128), ) # Verify non-targeted tensor was not modified assert torch.equal( tensors["model.embed_tokens.weight"], non_targeted_weight ), "Non-targeted tensor should not be modified" @pytest.mark.unit def test_fp8_block_converter_validate_with_meta_tensors(): """ Test that the converter's validate method works correctly with meta tensors. """ converter = FP8BlockDequantizer( targets=[r"re:.*layer\d+\.mlp\..*proj$"], weight_block_size=(128, 128) ) # Create mock tensors dict with FP8 weights and scale_inv tensors on meta device num_row_blocks = 2 num_col_blocks = 2 with torch.device("meta"): tensors = { "model.layer0.mlp.up_proj.weight": torch.empty( 256, 256, dtype=torch.float8_e4m3fn ), "model.layer0.mlp.up_proj.weight_scale_inv": torch.empty( num_row_blocks, num_col_blocks, dtype=torch.float32 ), "model.layer1.mlp.down_proj.weight": torch.empty( 256, 256, dtype=torch.float8_e4m3fn ), "model.layer1.mlp.down_proj.weight_scale_inv": torch.empty( num_row_blocks, num_col_blocks, dtype=torch.float32 ), "model.embed_tokens.weight": torch.empty(128, 128, dtype=torch.bfloat16), } # Should not raise any errors converter.validate(tensors) def _verify_block_conversion( result: torch.Tensor, weight_fp8: torch.Tensor, weight_scale_inv: torch.Tensor, block_size: tuple[int, int], ): """ Helper method to verify that FP8 block conversion to bfloat16 is correct. Checks that each block is correctly scaled by its corresponding scale_inv value. """ block_height, block_width = block_size num_row_blocks = weight_scale_inv.shape[0] num_col_blocks = weight_scale_inv.shape[1] # Verify output properties assert result.shape == weight_fp8.shape, "Output shape should match input shape" assert result.dtype == torch.bfloat16, "Output dtype should be bfloat16" # Verify the conversion logic: each block should be multiplied by its scale_inv for row_block in range(num_row_blocks): for col_block in range(num_col_blocks): row_start = row_block * block_height row_end = min((row_block + 1) * block_height, result.shape[0]) col_start = col_block * block_width col_end = min((col_block + 1) * block_width, result.shape[1]) # Get the block from result result_block = result[row_start:row_end, col_start:col_end] # Get expected: weight_fp8 block * scale_inv expected_block = ( weight_fp8[row_start:row_end, col_start:col_end].to(torch.float32) * weight_scale_inv[row_block, col_block].to(torch.float32) ).to(torch.bfloat16) # They should be equal (within floating point precision) assert torch.allclose( result_block.to(torch.float32), expected_block.to(torch.float32), rtol=1e-2, atol=1e-3, ), f"Block ({row_block}, {col_block}) conversion mismatch" test_modelopt_nvfp4.py000066400000000000000000000113671521257237700354400ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_entrypoints/convert/converters# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.entrypoints.convert import ModelOptNvfp4Converter @pytest.mark.unit def test_modelopt_nvfp4_converter_process(): """ Test that the converter's process method correctly transforms ModelOpt NVFP4 tensors to compressed-tensors NVFP4 format. """ converter = ModelOptNvfp4Converter(targets=[r"re:.*layer\d+\.mlp\..*proj$"]) # Create mock tensors dict with ModelOpt NVFP4 format input_scale = torch.tensor([2.0], dtype=torch.float32) weight = torch.randint(0, 255, (256, 256), dtype=torch.uint8) weight_scale = torch.rand(256, 1, dtype=torch.float32).to(torch.float8_e4m3fn) weight_scale_2 = torch.tensor([4.0], dtype=torch.float32) embed_weight = torch.randn(128, 128, dtype=torch.bfloat16) tensors = { "model.layer0.mlp.up_proj.input_scale": input_scale, "model.layer0.mlp.up_proj.weight": weight, "model.layer0.mlp.up_proj.weight_scale": weight_scale, "model.layer0.mlp.up_proj.weight_scale_2": weight_scale_2, "model.embed_tokens.weight": embed_weight, } # Process the tensors result = converter.process(tensors) # Verify transformations # input_scale -> input_global_scale (inverted) assert "model.layer0.mlp.up_proj.input_scale" not in result assert "model.layer0.mlp.up_proj.input_global_scale" in result assert torch.allclose( result["model.layer0.mlp.up_proj.input_global_scale"], 1 / input_scale, ) # weight -> weight_packed (renamed) assert "model.layer0.mlp.up_proj.weight" not in result assert "model.layer0.mlp.up_proj.weight_packed" in result assert torch.equal(result["model.layer0.mlp.up_proj.weight_packed"], weight) # weight_scale stays the same assert "model.layer0.mlp.up_proj.weight_scale" in result assert ( result["model.layer0.mlp.up_proj.weight_scale"].data_ptr() == weight_scale.data_ptr() ) # weight_scale_2 -> weight_global_scale (inverted) assert "model.layer0.mlp.up_proj.weight_scale_2" not in result assert "model.layer0.mlp.up_proj.weight_global_scale" in result assert torch.allclose( result["model.layer0.mlp.up_proj.weight_global_scale"], 1 / weight_scale_2, ) # Non-targeted tensor should not be modified assert torch.equal(result["model.embed_tokens.weight"], embed_weight) @pytest.mark.unit def test_modelopt_nvfp4_converter_get_dependencies(): """ Test that get_dependencies returns the correct dependent tensors for targeted weight tensors. """ converter = ModelOptNvfp4Converter(targets=[r"re:.*down_proj$"]) # Targeted layer should have dependencies deps = converter.get_dependencies("model.layer0.mlp.down_proj.weight") assert deps == { "model.layer0.mlp.down_proj.input_scale", "model.layer0.mlp.down_proj.weight_scale", "model.layer0.mlp.down_proj.weight_scale_2", } # Non-targeted layer should have no dependencies deps = converter.get_dependencies("model.layer0.mlp.up_proj.weight") assert deps == set() # Non-weight tensor should have no dependencies deps = converter.get_dependencies("model.layer0.mlp.down_proj.weight_scale") assert deps == set() @pytest.mark.unit def test_modelopt_nvfp4_converter_validate_with_meta_tensors(): """ Test that the converter's validate method works correctly with meta tensors. """ converter = ModelOptNvfp4Converter(targets=[r"re:.*layer\d+\.mlp\..*proj$"]) # Create mock tensors dict with NVFP4 tensors on meta device with torch.device("meta"): tensors = { "model.layer0.mlp.up_proj.input_scale": torch.empty(1, dtype=torch.float32), "model.layer0.mlp.up_proj.weight": torch.empty(256, 256, dtype=torch.uint8), "model.layer0.mlp.up_proj.weight_scale": torch.empty( 256, 1, dtype=torch.float8_e4m3fn ), "model.layer0.mlp.up_proj.weight_scale_2": torch.empty( 1, dtype=torch.float32 ), "model.layer1.mlp.down_proj.input_scale": torch.empty( 1, dtype=torch.float32 ), "model.layer1.mlp.down_proj.weight": torch.empty( 256, 256, dtype=torch.uint8 ), "model.layer1.mlp.down_proj.weight_scale": torch.empty( 256, 1, dtype=torch.float8_e4m3fn ), "model.layer1.mlp.down_proj.weight_scale_2": torch.empty( 1, dtype=torch.float32 ), "model.embed_tokens.weight": torch.empty(128, 128, dtype=torch.bfloat16), } # Should not raise any errors converter.validate(tensors) vllm-project-compressed-tensors-c18a0fa/tests/test_examples/000077500000000000000000000000001521257237700245045ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_examples/test_bitmask_compression_ipynb.py000066400000000000000000000014111521257237700333660ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest nbformat = pytest.importorskip("nbformat") from nbconvert.preprocessors import ExecutePreprocessor # noqa: E402 @pytest.mark.skip( reason="GHA not setup yet to run those tests. The test should work locally" ) @pytest.mark.parametrize("notebook", ["examples/bitmask_compression.ipynb"]) def test_notebook_exec(notebook): with open(notebook) as f: nb = nbformat.read(f, as_version=4) ep = ExecutePreprocessor(timeout=600, kernel_name="python3") try: assert ep.preprocess(nb) is not None, f"Got empty notebook for {notebook}" except Exception: assert False, f"Failed executing {notebook}" vllm-project-compressed-tensors-c18a0fa/tests/test_modeling/000077500000000000000000000000001521257237700244645ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_modeling/test_attention_and_cache.py000066400000000000000000000062211521257237700320500ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.modeling import ( IMPL_ATTR, KV_CACHE_ATTR, QuantizedAttentionImpl, QuantizedKVCache, initialize_hooked_attention, initialize_hooked_kv_cache, register_key_hook, register_query_hook, register_value_hook, ) from tests.testing_utils import requires_gpu from transformers import AutoModelForCausalLM @requires_gpu def test_attention_cache(): model = AutoModelForCausalLM.from_pretrained( "nm-testing/llama2.c-stories15M", device_map="cuda" ) inputs = {key: value.to("cuda") for key, value in model.dummy_inputs.items()} true_outputs = model(**inputs) layers = model.model.layers # check if hooks work k_called = [False for _ in range(len(layers))] v_called = [False for _ in range(len(layers))] # apply kv cache quantization _apply_kv_cache(model, layers, k_called, v_called) # check kv cache quantization outputs = model(**inputs) assert torch.equal(outputs.logits, true_outputs.logits) assert all(k_called) and all(v_called) """ apply attention quantization after kv cache quantization """ # check if hooks work q_called = [False for _ in range(len(layers))] k_called = [False for _ in range(len(layers))] v_called = [False for _ in range(len(layers))] # apply attention quantization _apply_attention(model, layers, q_called, k_called, v_called) # check attention quantization outputs = model(**inputs) assert torch.equal(outputs.logits, true_outputs.logits) assert all(q_called) and all(k_called) and all(v_called) def _apply_kv_cache(model, layers, k_called, v_called): for layer_index, layer in enumerate(layers): module = layer.self_attn initialize_hooked_kv_cache(model, module) assert isinstance(getattr(module, KV_CACHE_ATTR), QuantizedKVCache) # reapply is no-op initialize_hooked_kv_cache(model, module) def k_hook(_module, _states, layer_index=layer_index): # NOTE: capture by value k_called[layer_index] = True def v_hook(_module, _states, layer_index=layer_index): my_index = layer_index v_called[my_index] = True register_key_hook(module, k_hook) register_value_hook(module, v_hook) def _apply_attention(model, layers, q_called, k_called, v_called): for layer_index, layer in enumerate(layers): module = layer.self_attn initialize_hooked_attention(model, module) assert isinstance(getattr(module, IMPL_ATTR), QuantizedAttentionImpl) # reapply is no-op initialize_hooked_attention(model, module) def q_hook(_module, _states, layer_index=layer_index): q_called[layer_index] = True def k_hook(_module, _states, layer_index=layer_index): k_called[layer_index] = True def v_hook(_module, _states, layer_index=layer_index): v_called[layer_index] = True register_query_hook(module, q_hook) register_key_hook(module, k_hook) register_value_hook(module, v_hook) vllm-project-compressed-tensors-c18a0fa/tests/test_modeling/test_deepseekv3_kvcache_quant.py000066400000000000000000000055221521257237700330330ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from compressed_tensors.modeling import ( IMPL_ATTR, KV_CACHE_ATTR, QuantizedAttentionImpl, QuantizedKVCache, register_query_hook, ) from compressed_tensors.quantization.lifecycle.apply import apply_quantization_config from compressed_tensors.quantization.quant_args import ( QuantizationArgs, QuantizationStrategy, ) from compressed_tensors.quantization.quant_config import ( QuantizationConfig, QuantizationStatus, ) from compressed_tensors.quantization.quant_scheme import QuantizationScheme from tests.testing_utils import requires_gpu from transformers import AutoModelForCausalLM @requires_gpu def test_apply_config_detects_deepseekv3_attention_and_hooks(): model = AutoModelForCausalLM.from_pretrained( "trl-internal-testing/tiny-DeepseekV3ForCausalLM", device_map="cuda" ) inputs = {key: value.to("cuda") for key, value in model.dummy_inputs.items()} # Build attention quantization scheme targeting attention modules qa = QuantizationArgs( strategy=QuantizationStrategy.TENSOR, symmetric=True, num_bits=8, dynamic=False, ) scheme = QuantizationScheme( targets=["re:.*self_attn$"], input_activations=qa, ) config = QuantizationConfig( config_groups={"group_0": scheme}, quantization_status=QuantizationStatus.INITIALIZED, kv_cache_scheme=None, ) apply_quantization_config(model, config) # Validate q/k/v qparams initialized and hooks attached q_called = [] k_called = [] v_called = [] for idx, layer in enumerate(model.model.layers): attn = layer.self_attn assert isinstance(getattr(attn, IMPL_ATTR), QuantizedAttentionImpl) assert isinstance(getattr(attn, KV_CACHE_ATTR), QuantizedKVCache) assert hasattr(attn, "q_scale") assert hasattr(attn, "k_scale") assert hasattr(attn, "v_scale") assert hasattr(attn, "q_zero_point") assert hasattr(attn, "k_zero_point") assert hasattr(attn, "v_zero_point") q_called.append(False) k_called.append(False) v_called.append(False) def q_hook(_module, _states, i=idx): q_called[i] = True register_query_hook(attn, q_hook) impl = getattr(attn, IMPL_ATTR) def _k_pre_hook(_impl, args, kwargs, i=idx): k_called[i] = True return args, kwargs def _v_pre_hook(_impl, args, kwargs, i=idx): v_called[i] = True return args, kwargs impl.register_forward_pre_hook(_k_pre_hook, with_kwargs=True) impl.register_forward_pre_hook(_v_pre_hook, with_kwargs=True) model(**inputs, use_cache=True) assert all(q_called) and all(k_called) and all(v_called) vllm-project-compressed-tensors-c18a0fa/tests/test_offload/000077500000000000000000000000001521257237700243005ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_offload/cache/000077500000000000000000000000001521257237700253435ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_offload/cache/conftest.py000066400000000000000000000011161521257237700275410ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os import pytest from compressed_tensors.offload import OffloadCache @pytest.fixture() def offload_cache(offload_device, onload_device, tmp_path): if offload_device == "disk": offload_dir = str(tmp_path / "offload_dir") os.makedirs(offload_dir) return OffloadCache.cls_from_device(offload_device)( onload_device, offload_dir=offload_dir ) else: return OffloadCache.cls_from_device(offload_device)(onload_device) vllm-project-compressed-tensors-c18a0fa/tests/test_offload/cache/helpers.py000066400000000000000000000125311521257237700273610ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import gc from weakref import ref import torch from tests.test_offload.conftest import assert_device_equal, assert_tensor_equal def _test_onloading(offload_device, onload_device, offload_cache): tensor = torch.ones(10) offload_cache["weight"] = tensor onloaded = offload_cache["weight"] assert type(onloaded) is type(tensor) assert_tensor_equal(onloaded, tensor, onload_device) def _test_garbage_collect(offload_device, onload_device, offload_cache): offload_cache["weight"] = torch.ones(10) onloaded = offload_cache["weight"] onloaded_ref = ref(onloaded) del onloaded gc.collect() assert onloaded_ref() is None def _test_offload(offload_device, onload_device, offload_cache): tensor = torch.ones(10, device=onload_device) offloaded = offload_cache.offload(tensor) assert_device_equal(offloaded.device, offload_device) assert_tensor_equal(offloaded, tensor, offload_device) def _test_onload(offload_device, onload_device, offload_cache): tensor = torch.ones(10, device=onload_device) onloaded = offload_cache.onload(offload_cache.offload(tensor)) assert_device_equal(onloaded.device, onload_device) assert_tensor_equal(onloaded, tensor, onload_device) def _test_disable_offloading(offload_device, onload_device, offload_cache): offload_cache["weight"] = torch.ones(10) outside_onloaded = offload_cache["weight"] outside_onloaded_ref = ref(outside_onloaded) assert_device_equal(outside_onloaded.device, onload_device) with offload_cache.disable_offloading(): inside_onloaded = offload_cache["weight"] inside_onloaded_ref = ref(inside_onloaded) assert_device_equal(inside_onloaded.device, onload_device) del outside_onloaded del inside_onloaded gc.collect() assert outside_onloaded_ref() is None assert inside_onloaded_ref() is not None assert outside_onloaded_ref() is None assert inside_onloaded_ref() is None def _test_disable_onloading(offload_device, onload_device, offload_cache): tensor = torch.ones(10) offload_cache.offloaded_values["weight"] = tensor with offload_cache.disable_onloading(): onloaded = offload_cache["weight"] assert onloaded is tensor assert onloaded is tensor def _test_delete(offload_device, onload_device, offload_cache): offload_cache["weight"] = torch.ones(10) onloaded = offload_cache["weight"] onloaded_ref = ref(onloaded) with offload_cache.disable_offloading(): del offload_cache["weight"] del onloaded gc.collect() assert onloaded_ref() is None assert onloaded_ref() is None def _test_shared_attributes(offload_device, onload_device, offload_cache): assert ( offload_cache.offloading_disabled is offload_cache.__class__.offloading_disabled ) assert ( offload_cache.onloading_disabled is offload_cache.__class__.onloading_disabled ) assert ( offload_cache.keep_onloaded_values is offload_cache.__class__.keep_onloaded_values ) assert not hasattr(offload_cache.__class__, "onload_device") assert not hasattr(offload_cache.__class__, "offloaded_values") def _test_tensor_subclass(offload_device, onload_device, offload_cache): tensor = torch.ones(10) param = torch.nn.Parameter(torch.ones(10), requires_grad=False) buffer = torch.nn.Buffer(torch.ones(10)) offload_cache["tensor"] = tensor offload_cache["param"] = param offload_cache["buffer"] = buffer assert_tensor_equal(offload_cache["tensor"], tensor, onload_device) assert_tensor_equal(offload_cache["param"], param, onload_device) assert_tensor_equal(offload_cache["buffer"], buffer, onload_device) with offload_cache.disable_onloading(): assert_tensor_equal(offload_cache["tensor"], tensor, offload_device) assert_tensor_equal(offload_cache["param"], param, offload_device) assert_tensor_equal(offload_cache["buffer"], buffer, offload_device) def _test_update_offload(offload_device, onload_device, offload_cache): # Create initial tensor and offload it initial_data = torch.ones(10, device=onload_device) offload_cache["weight"] = initial_data # Verify initial value onloaded = offload_cache["weight"] assert_tensor_equal(onloaded, initial_data, onload_device) # Update with new data new_data = torch.ones(10, device=onload_device) * 2.0 offload_cache["weight"] = new_data # Verify update worked updated_onloaded = offload_cache["weight"] assert_tensor_equal(updated_onloaded, new_data, onload_device) # Verify offloaded tensor was updated in place (not replaced) with offload_cache.disable_onloading(): offloaded = offload_cache["weight"] assert_tensor_equal(offloaded, new_data, offload_device) # Test update with disable_offloading context with offload_cache.disable_offloading(): offload_cache["weight"] = torch.ones(10, device=onload_device) * 3.0 cached_onloaded = offload_cache["weight"] assert_tensor_equal(cached_onloaded, torch.ones(10) * 3.0, onload_device) # Verify update persisted after context exit final_onloaded = offload_cache["weight"] assert_tensor_equal(final_onloaded, torch.ones(10) * 3.0, onload_device) vllm-project-compressed-tensors-c18a0fa/tests/test_offload/cache/test_cpu.py000066400000000000000000000075761521257237700275620ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import compressed_tensors.offload.cache.cpu as cpu_cache import pytest import torch from loguru import logger as loguru_logger from tests.test_offload.cache.helpers import ( _test_delete, _test_disable_offloading, _test_disable_onloading, _test_garbage_collect, _test_offload, _test_onload, _test_onloading, _test_shared_attributes, _test_tensor_subclass, ) from tests.testing_utils import requires_gpu @pytest.fixture() def onload_device(): return torch.device(torch.accelerator.current_accelerator().type) @pytest.fixture() def offload_device(): return torch.device("cpu") @pytest.mark.unit @requires_gpu def test_delete(offload_device, onload_device, offload_cache): _test_delete(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_disable_offloading(offload_device, onload_device, offload_cache): _test_disable_offloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_disable_onloading(offload_device, onload_device, offload_cache): _test_disable_onloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_garbage_collect(offload_device, onload_device, offload_cache): _test_garbage_collect(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_offload(offload_device, onload_device, offload_cache): _test_offload(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_onload(offload_device, onload_device, offload_cache): _test_onload(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_onloading(offload_device, onload_device, offload_cache): _test_onloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_shared_attributes(offload_device, onload_device, offload_cache): _test_shared_attributes(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_tensor_subclass(offload_device, onload_device, offload_cache): _test_tensor_subclass(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_offload_logs_memory_hint(onload_device): cache = cpu_cache.CPUCache(onload_device) original_send_tensors = cpu_cache.send_tensors def raise_memory_error(*args, **kwargs): raise RuntimeError("mmap failed: Cannot allocate memory") cpu_cache.send_tensors = raise_memory_error warnings = [] handler_id = loguru_logger.add( lambda msg: warnings.append(msg.record["message"]), level="WARNING" ) try: with pytest.raises(RuntimeError, match="Cannot allocate memory"): cache.offload(torch.zeros(1, device=onload_device)) finally: cpu_cache.send_tensors = original_send_tensors loguru_logger.remove(handler_id) assert any( "CPU offloading ran out of host RAM or mmap descriptors." in w for w in warnings ) @pytest.mark.unit @requires_gpu def test_offload_logs_memory_hint_oserror(onload_device): import errno cache = cpu_cache.CPUCache(onload_device) original_send_tensors = cpu_cache.send_tensors def raise_memory_error(*args, **kwargs): raise OSError(errno.ENOMEM, "Cannot allocate memory") cpu_cache.send_tensors = raise_memory_error warnings = [] handler_id = loguru_logger.add( lambda msg: warnings.append(msg.record["message"]), level="WARNING" ) try: with pytest.raises(OSError): cache.offload(torch.zeros(1, device=onload_device)) finally: cpu_cache.send_tensors = original_send_tensors loguru_logger.remove(handler_id) assert any( "CPU offloading ran out of host RAM or mmap descriptors." in w for w in warnings ) vllm-project-compressed-tensors-c18a0fa/tests/test_offload/cache/test_device.py000066400000000000000000000061621521257237700302200ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import gc from weakref import ref import pytest import torch from compressed_tensors.offload.cache.device import DeviceCache from tests.test_offload.cache.helpers import ( _test_delete, _test_disable_onloading, _test_offload, _test_onload, _test_onloading, _test_shared_attributes, _test_tensor_subclass, ) from tests.test_offload.conftest import assert_device_equal from tests.testing_utils import requires_gpu @pytest.fixture() def onload_device(): return torch.accelerator.current_accelerator() @pytest.fixture() def offload_device(): return torch.accelerator.current_accelerator() @pytest.mark.unit @requires_gpu def test_delete(offload_device, onload_device, offload_cache): _test_delete(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_disable_offloading(onload_device): # unlike other device caches, the onload is not garbage collected cache = DeviceCache(onload_device) cache["weight"] = torch.ones(10) outside_onloaded = cache["weight"] outside_onloaded_ref = ref(outside_onloaded) assert_device_equal(outside_onloaded.device, onload_device) with cache.disable_offloading(): inside_onloaded = cache["weight"] inside_onloaded_ref = ref(inside_onloaded) assert_device_equal(inside_onloaded.device, onload_device) del outside_onloaded del inside_onloaded gc.collect() assert outside_onloaded_ref() is not None # changed assert inside_onloaded_ref() is not None assert outside_onloaded_ref() is not None # changed assert inside_onloaded_ref() is not None # changed @pytest.mark.unit @requires_gpu def test_disable_onloading(offload_device, onload_device, offload_cache): _test_disable_onloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_garbage_collect(onload_device): # unlike other device caches, the onload is not garbage collected cache = DeviceCache(onload_device) cache["weight"] = torch.ones(10) onloaded = cache["weight"] onloaded_ref = ref(onloaded) del onloaded gc.collect() assert onloaded_ref() is not None # changed @pytest.mark.unit @requires_gpu def test_offload(offload_device, onload_device, offload_cache): _test_offload(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu @requires_gpu def test_onload(offload_device, onload_device, offload_cache): _test_onload(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_onloading(offload_device, onload_device, offload_cache): _test_onloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_shared_attributes(offload_device, onload_device, offload_cache): _test_shared_attributes(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_tensor_subclass(offload_device, onload_device, offload_cache): _test_tensor_subclass(offload_device, onload_device, offload_cache) vllm-project-compressed-tensors-c18a0fa/tests/test_offload/cache/test_disk.py000066400000000000000000000067731521257237700277230ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os import pytest import torch from compressed_tensors.offload.cache.disk import DiskCache from safetensors import safe_open from tests.test_offload.cache.helpers import ( _test_delete, _test_disable_offloading, _test_disable_onloading, _test_garbage_collect, _test_offload, _test_onload, _test_onloading, _test_shared_attributes, _test_tensor_subclass, _test_update_offload, ) from tests.test_offload.conftest import assert_tensor_equal from tests.testing_utils import requires_gpu @pytest.fixture() def onload_device(): return torch.accelerator.current_accelerator() @pytest.fixture() def offload_device(): return "disk" @pytest.mark.unit @requires_gpu def test_delete(offload_device, onload_device, offload_cache): _test_delete(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_disable_offloading(offload_device, onload_device, offload_cache): _test_disable_offloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_disable_onloading(offload_device, onload_device, offload_cache): _test_disable_onloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_garbage_collect(offload_device, onload_device, offload_cache): _test_garbage_collect(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_offload(offload_device, onload_device, offload_cache): _test_offload(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu @requires_gpu def test_onload(offload_device, onload_device, offload_cache): _test_onload(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_onloading(offload_device, onload_device, offload_cache): _test_onloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_shared_attributes(offload_device, onload_device, offload_cache): _test_shared_attributes(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_tensor_subclass(offload_device, onload_device, offload_cache): _test_tensor_subclass(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu def test_update_offload(offload_device, onload_device, offload_cache): _test_update_offload(offload_device, onload_device, offload_cache) @pytest.mark.unit def test_files(tmp_path): offload_dir = tmp_path / "offload_dir" os.mkdir(offload_dir) # initial write DiskCache.index = {} cache = DiskCache("cpu", offload_dir=str(offload_dir)) tensor = torch.zeros(10) cache["weight"] = tensor files = os.listdir(offload_dir) assert len(DiskCache.index) == 1 assert len(files) == 1 with safe_open(offload_dir / files[0], framework="pt", device="cpu") as file: read_tensor = file.get_tensor("weight") assert_tensor_equal(read_tensor, tensor) # modify tensor = torch.ones(10) cache["weight"] = tensor files = os.listdir(offload_dir) assert len(DiskCache.index) == 1 assert len(files) == 1 with safe_open(offload_dir / files[0], framework="pt", device="cpu") as file: read_tensor = file.get_tensor("weight") assert_tensor_equal(read_tensor, tensor) # delete del cache["weight"] files = os.listdir(offload_dir) assert len(DiskCache.index) == 0 assert len(files) == 0 vllm-project-compressed-tensors-c18a0fa/tests/test_offload/cache/test_dist_cpu.py000066400000000000000000000151361521257237700305740ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch import torch.distributed as dist from compressed_tensors.offload import disable_onloading from compressed_tensors.offload.cache.dist_cpu import DistributedCPUCache from loguru import logger as loguru_logger from tests.test_offload.cache.helpers import ( _test_delete, _test_disable_offloading, _test_disable_onloading, _test_garbage_collect, _test_offload, _test_onload, _test_onloading, _test_shared_attributes, _test_tensor_subclass, ) from tests.test_offload.conftest import torchrun from tests.testing_utils import requires_gpu @pytest.fixture() def onload_device(): return torch.accelerator.current_accelerator() @pytest.fixture() def offload_device(): return torch.device("cpu") @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_delete(offload_device, onload_device, offload_cache): _test_delete(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_disable_offloading(offload_device, onload_device, offload_cache): _test_disable_offloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_disable_onloading(offload_device, onload_device, offload_cache): _test_disable_onloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_garbage_collect(offload_device, onload_device, offload_cache): _test_garbage_collect(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_offload(offload_device, onload_device, offload_cache): _test_offload(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_onload(offload_device, onload_device, offload_cache): _test_onload(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_onloading(offload_device, onload_device, offload_cache): _test_onloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_shared_attributes(offload_device, onload_device, offload_cache): _test_shared_attributes(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_tensor_subclass(offload_device, onload_device, offload_cache): _test_tensor_subclass(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_offload(onload_device): cache = DistributedCPUCache(onload_device) tensor = torch.zeros((5, 2)) cache["tensor"] = tensor # check tensor construction assert torch.equal(cache["tensor"].cpu(), tensor) with disable_onloading(): assert torch.equal(cache["tensor"].cpu(), tensor) # update tensor tensor = torch.ones((5, 2)) cache["tensor"] = tensor # check tensor construction assert torch.equal(cache["tensor"].cpu(), tensor) with disable_onloading(): assert torch.equal(cache["tensor"].cpu(), tensor) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_shared_cpu_offload(onload_device): cache = DistributedCPUCache(onload_device) tensor = torch.zeros((5, 2)) cache["tensor"] = tensor # modify the offloaded cpu tensor directly tensor = torch.ones((5, 2)) if dist.get_rank() == 0: with disable_onloading(): cache["tensor"].copy_(tensor) dist.barrier() # check that the value is affected on all ranks assert torch.equal(cache["tensor"].cpu(), tensor) with disable_onloading(): assert torch.equal(cache["tensor"].cpu(), tensor) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_async_update(onload_device): """ Test that different ranks can update different tensors asynchronously, and that values are correct after a barrier. """ cache = DistributedCPUCache(onload_device) # Initialize two tensors in the cache cache["tensor_0"] = torch.zeros(10, device=onload_device) cache["tensor_1"] = torch.zeros(10, device=onload_device) # Each rank updates a different tensor rank = dist.get_rank() if rank == 0: # Rank 0 updates tensor_0 cache[f"tensor_{rank}"] = torch.ones(10, device=onload_device) * 1.0 elif rank == 1: # Rank 1 updates tensor_1 cache[f"tensor_{rank}"] = torch.ones(10, device=onload_device) * 2.0 # Synchronize to ensure all updates are complete dist.barrier() # Verify that both tensors have the correct values on all ranks tensor_0 = cache["tensor_0"] tensor_1 = cache["tensor_1"] assert torch.allclose(tensor_0.cpu(), torch.ones(10) * 1.0) assert torch.allclose(tensor_1.cpu(), torch.ones(10) * 2.0) # Verify offloaded values are also correct with disable_onloading(): offloaded_0 = cache["tensor_0"] offloaded_1 = cache["tensor_1"] assert torch.allclose(offloaded_0.cpu(), torch.ones(10) * 1.0) assert torch.allclose(offloaded_1.cpu(), torch.ones(10) * 2.0) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_offload_logs_memory_hint(onload_device): cache = DistributedCPUCache(onload_device) original_share_memory = torch.Tensor.share_memory_ def raise_memory_error(*args, **kwargs): raise RuntimeError("mmap failed: Cannot allocate memory") torch.Tensor.share_memory_ = raise_memory_error warnings = [] handler_id = loguru_logger.add( lambda msg: warnings.append(msg.record["message"]), level="WARNING" ) try: # Only Rank 0 calls offload(), which throws an error *before* the # dist.broadcast barrier. This cleanly avoids the hang since Rank 1 # safely exits without waiting. if dist.get_rank() == 0: with pytest.raises(RuntimeError, match="Cannot allocate memory"): cache.offload(torch.zeros(1, device=onload_device)) finally: torch.Tensor.share_memory_ = original_share_memory loguru_logger.remove(handler_id) if dist.get_rank() == 0: assert any( "CPU offloading ran out of host RAM or mmap descriptors." in w for w in warnings ) vllm-project-compressed-tensors-c18a0fa/tests/test_offload/cache/test_dist_device.py000066400000000000000000000130601521257237700312360ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import gc from weakref import ref import pytest import torch import torch.distributed as dist from compressed_tensors.offload import disable_onloading from compressed_tensors.offload.cache.dist_device import DistributedDeviceCache from tests.test_offload.cache.helpers import ( _test_delete, _test_disable_onloading, _test_offload, _test_onload, _test_onloading, _test_shared_attributes, _test_tensor_subclass, ) from tests.test_offload.conftest import assert_device_equal, torchrun from tests.testing_utils import requires_gpu # Note that tests only require at least 1 gpu # b/c different ranks can share the same gpu @pytest.fixture() def onload_device(): return torch.accelerator.current_accelerator() @pytest.fixture() def offload_device(): return torch.accelerator.current_accelerator() @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_delete(offload_device, onload_device, offload_cache): _test_delete(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_disable_offloading(onload_device): # unlike other device caches, the onload is not garbage collected cache = DistributedDeviceCache(onload_device) cache["weight"] = torch.ones(10) outside_onloaded = cache["weight"] outside_onloaded_ref = ref(outside_onloaded) assert_device_equal(outside_onloaded.device, onload_device) with cache.disable_offloading(): inside_onloaded = cache["weight"] inside_onloaded_ref = ref(inside_onloaded) assert_device_equal(inside_onloaded.device, onload_device) del outside_onloaded del inside_onloaded gc.collect() assert outside_onloaded_ref() is not None # changed assert inside_onloaded_ref() is not None assert outside_onloaded_ref() is not None # changed assert inside_onloaded_ref() is not None # changed @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_disable_onloading(offload_device, onload_device, offload_cache): _test_disable_onloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_garbage_collect(onload_device): # unlike other device caches, the onload is not garbage collected cache = DistributedDeviceCache(onload_device) cache["weight"] = torch.ones(10) onloaded = cache["weight"] onloaded_ref = ref(onloaded) del onloaded gc.collect() assert onloaded_ref() is not None # changed @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_offload(offload_device, onload_device, offload_cache): _test_offload(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_onload(offload_device, onload_device, offload_cache): _test_onload(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_onloading(offload_device, onload_device, offload_cache): _test_onloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_shared_attributes(offload_device, onload_device, offload_cache): _test_shared_attributes(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_tensor_subclass(offload_device, onload_device, offload_cache): _test_tensor_subclass(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_offload(onload_device): cache = DistributedDeviceCache(onload_device) tensor = torch.zeros((5, 2)) cache["tensor"] = tensor # check tensor construction assert torch.equal(cache["tensor"].cpu(), tensor) with disable_onloading(): assert torch.equal(cache["tensor"].cpu(), tensor) # update tensor tensor = torch.ones((5, 2)) cache["tensor"] = tensor # check tensor construction assert torch.equal(cache["tensor"].cpu(), tensor) with disable_onloading(): assert torch.equal(cache["tensor"].cpu(), tensor) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_offload_fp8(onload_device): """FP8 tensors should broadcast successfully and preserve their dtype""" float8_dtypes = [ torch.float8_e4m3fn, torch.float8_e5m2, torch.float8_e4m3fnuz, torch.float8_e5m2fnuz, ] for dtype in float8_dtypes: cache = DistributedDeviceCache(onload_device) tensor = torch.zeros((5, 2), dtype=dtype) cache["tensor"] = tensor result = cache["tensor"].cpu() assert result.shape == tensor.shape assert result.dtype == dtype @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_replicated_device_offload(onload_device): cache = DistributedDeviceCache(onload_device) tensor = torch.empty((5, 2)) cache["tensor"] = tensor # modify the offloaded cpu tensor directly tensor = torch.full((5, 2), dist.get_rank()) cache["tensor"].copy_(tensor) dist.barrier() # check that the value is affected on all ranks assert torch.equal(cache["tensor"].cpu(), tensor) with disable_onloading(): assert torch.equal(cache["tensor"].cpu(), tensor) vllm-project-compressed-tensors-c18a0fa/tests/test_offload/cache/test_dist_disk.py000066400000000000000000000170071521257237700307360ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os import pytest import torch import torch.distributed as dist from compressed_tensors.offload import disable_onloading from compressed_tensors.offload.cache.disk import DiskCache from compressed_tensors.offload.cache.dist_disk import DistributedDiskCache from safetensors import safe_open from tests.test_offload.cache.helpers import ( _test_delete, _test_disable_offloading, _test_disable_onloading, _test_garbage_collect, _test_offload, _test_onload, _test_onloading, _test_shared_attributes, _test_tensor_subclass, ) from tests.test_offload.conftest import assert_tensor_equal, torchrun from tests.testing_utils import requires_gpu @pytest.fixture() def onload_device(): return torch.accelerator.current_accelerator() @pytest.fixture() def offload_device(): return "disk" @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_delete(offload_device, onload_device, offload_cache): _test_delete(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_disable_offloading(offload_device, onload_device, offload_cache): _test_disable_offloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_disable_onloading(offload_device, onload_device, offload_cache): _test_disable_onloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_garbage_collect(offload_device, onload_device, offload_cache): _test_garbage_collect(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_offload(offload_device, onload_device, offload_cache): _test_offload(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_onload(offload_device, onload_device, offload_cache): _test_onload(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_onloading(offload_device, onload_device, offload_cache): _test_onloading(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_shared_attributes(offload_device, onload_device, offload_cache): _test_shared_attributes(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_tensor_subclass(offload_device, onload_device, offload_cache): _test_tensor_subclass(offload_device, onload_device, offload_cache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_offload(onload_device, tmp_path): # Broadcast directory path from rank 0 to all ranks if dist.get_rank() == 0: offload_dir = tmp_path / "offload_dir" os.mkdir(offload_dir) broadcast_obj = [str(offload_dir)] else: broadcast_obj = [None] dist.broadcast_object_list(broadcast_obj, src=0) offload_dir = broadcast_obj[0] # Ensure directory creation completes before other ranks proceed dist.barrier() cache = DistributedDiskCache(onload_device, offload_dir=offload_dir) tensor = torch.zeros((5, 2)) cache["tensor"] = tensor # check tensor construction assert torch.equal(cache["tensor"], tensor.to(onload_device)) with disable_onloading(): assert_tensor_equal(cache["tensor"], tensor.to("meta")) # update tensor tensor = torch.ones((5, 2)) cache["tensor"] = tensor # check tensor construction assert torch.equal(cache["tensor"], tensor.to(onload_device)) with disable_onloading(): assert_tensor_equal(cache["tensor"], tensor.to("meta")) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_files(tmp_path): # Broadcast directory path from rank 0 to all ranks if dist.get_rank() == 0: offload_dir = tmp_path / "offload_dir" os.mkdir(offload_dir) broadcast_obj = [str(offload_dir)] else: broadcast_obj = [None] dist.broadcast_object_list(broadcast_obj, src=0) offload_dir = broadcast_obj[0] # Ensure directory creation completes before other ranks proceed dist.barrier() # initial write, broadcasted to all ranks DiskCache.index = {} cache = DistributedDiskCache("cpu", offload_dir=offload_dir) tensor = torch.zeros(10) cache["weight"] = tensor assert len(DiskCache.index) == 1 if dist.get_rank() == 0: # only rank0 bc `tmp_path` is not shared between ranks files = os.listdir(offload_dir) assert len(files) == 1 with safe_open( os.path.join(offload_dir, files[0]), framework="pt", device="cpu" ) as file: read_tensor = file.get_tensor("weight") assert_tensor_equal(read_tensor, tensor) # modify on one rank tensor = torch.ones(10) if dist.get_rank() == 0: cache["weight"] = tensor assert len(DiskCache.index) == 1 if dist.get_rank() == 0: # only rank0 bc `tmp_path` is not shared between ranks files = os.listdir(offload_dir) assert len(files) == 1 with safe_open( os.path.join(offload_dir, files[0]), framework="pt", device="cpu" ) as file: read_tensor = file.get_tensor("weight") assert_tensor_equal(read_tensor, tensor) # delete del cache["weight"] assert len(DiskCache.index) == 0 if dist.get_rank() == 0: # only rank0 bc `tmp_path` is not shared between ranks files = os.listdir(offload_dir) assert len(files) == 0 @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_async_update(tmp_path): """ Test that different ranks can update different tensors asynchronously, and that values are correct after a barrier. """ offload_dir = tmp_path / "offload_dir" if dist.get_rank() == 0: os.mkdir(offload_dir) # Ensure directory creation completes before other ranks proceed dist.barrier() onload_device = torch.accelerator.current_accelerator() cache = DistributedDiskCache(onload_device, offload_dir=str(offload_dir)) # Initialize two tensors in the cache cache["tensor_0"] = torch.zeros(10, device=onload_device) cache["tensor_1"] = torch.zeros(10, device=onload_device) # Each rank updates a different tensor rank = dist.get_rank() if rank == 0: # Rank 0 updates tensor_0 cache[f"tensor_{rank}"] = torch.ones(10, device=onload_device) * 1.0 elif rank == 1: # Rank 1 updates tensor_1 cache[f"tensor_{rank}"] = torch.ones(10, device=onload_device) * 2.0 # Synchronize to ensure all updates are complete dist.barrier() # Verify that both tensors have the correct values on all ranks tensor_0 = cache["tensor_0"] tensor_1 = cache["tensor_1"] assert torch.allclose(tensor_0.cpu(), torch.ones(10) * 1.0) assert torch.allclose(tensor_1.cpu(), torch.ones(10) * 2.0) # Verify offloaded values are also correct with disable_onloading(): offloaded_0 = cache["tensor_0"] offloaded_1 = cache["tensor_1"] assert_tensor_equal(offloaded_0, torch.ones(10) * 1.0, "disk") assert_tensor_equal(offloaded_1, torch.ones(10) * 2.0, "disk") vllm-project-compressed-tensors-c18a0fa/tests/test_offload/conftest.py000066400000000000000000000126471521257237700265110ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os import subprocess import sys from functools import wraps from types import FunctionType from typing import Any, Callable, Literal, Optional import pytest import torch from compressed_tensors.offload.utils import send_tensors accelerator_device = torch.accelerator.current_accelerator() skip_if_mps_device = pytest.mark.skipif( accelerator_device.type == "mps", reason="[Known issue] https://github.com/pytorch/pytorch/issues/167447", ) def assert_device_equal( device_a: torch.device | Literal["disk"], device_b: torch.device | Literal["disk"], ): if device_a == "disk": device_a = torch.device("meta") if device_b == "disk": device_b = torch.device("meta") cur_index = torch.accelerator.current_device_index() a_index = cur_index if device_a.index is None else device_a.index b_index = cur_index if device_b.index is None else device_b.index # Handle device emulation: when --emulate-xpu is active, tensors created # on "xpu" actually live on the real accelerator, so their .device reports # the real type. Normalize device types: if one matches the fake type and # the other matches the real type, treat them as equal. accel = torch.accelerator.current_accelerator() fake_type = accel.type real_type = getattr(accel, "_real_type", None) a_type = device_a.type b_type = device_b.type # If emulation is active, normalize: fake_type and real_type are equivalent if real_type is not None: if a_type == real_type: a_type = fake_type if b_type == real_type: b_type = fake_type assert a_type == b_type and a_index == b_index def assert_tensor_equal( tensor_a: torch.Tensor, tensor_b: torch.Tensor, device: Optional[torch.device | str] = None, ): if device is not None: tensor_b = send_tensors(tensor_b, "meta" if device == "disk" else device) assert tensor_a.__class__ == tensor_b.__class__ assert tensor_a.requires_grad == tensor_b.requires_grad assert tensor_a.__dict__ == tensor_b.__dict__ if tensor_a.is_meta or tensor_b.is_meta: assert ( tensor_a.device == tensor_b.device and tensor_a.shape == tensor_b.shape and tensor_a.dtype == tensor_b.dtype ) else: assert torch.equal(tensor_a, tensor_b) def torchrun( world_size: int = 1, init_dist: bool = False ) -> Callable[[Callable[..., Any]], Callable[..., Any]]: """ Pytest decorator to run a test within parallel torchrun subprocesses. This decorator automatically spawns torchrun when the test is run with regular pytest. When running under torchrun (detected via TORCHELASTIC_RUN_ID env var), it optionally initializes the distributed process group before running the test. Usage: @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_distributed_feature(): # Distributed already initialized ... @torchrun(world_size=2) # init_dist=False by default def test_custom_init(): # Handle your own distributed setup from compressed_tensors.distributed import init_dist init_dist() ... :param world_size: number of ranks to spawn :param init_dist: whether to automatically call init_dist() (default: False) """ def decorator(func: FunctionType): @wraps(func) def wrapper(*args, **kwargs): # We're running in a torchrun subprocess: optionally init then run the test if "TORCHELASTIC_RUN_ID" in os.environ: if init_dist: from compressed_tensors.distributed import init_dist as _init_dist _init_dist() return func(*args, **kwargs) # First time calling in the main process: # trigger torchrun with this function as the pytest target else: module = sys.modules.get(func.__module__) if module is None: raise RuntimeError( f"Can't find module {func.__module__} for func {func.__name__}" ) file_path = module.__file__ if file_path is None: raise RuntimeError( f"Module {func.__module__} has no __file__ attribute" ) func_name = func.__name__ cmd = [ sys.executable, "-m", "torch.distributed.run", "--nproc_per_node", str(world_size), "--log-dir", "/tmp/torchrun-logs", "--tee", "3", "--role", "torchrun", "-m", "pytest", f"{file_path}::{func_name}", "-sx", ] proc = subprocess.run(cmd) assert proc.returncode == 0 return wrapper return decorator @pytest.fixture() def accel_device(): accel_type = torch.accelerator.current_accelerator().type return ( torch.device(accel_type) if "TORCHELASTIC_RUN_ID" in os.environ else torch.device(accel_type, 0) ) vllm-project-compressed-tensors-c18a0fa/tests/test_offload/convert/000077500000000000000000000000001521257237700257605ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_offload/convert/test_convert.py000066400000000000000000000047711521257237700310620ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os import pytest import torch from compressed_tensors.distributed import is_source_process from compressed_tensors.offload import disable_onloading, from_accelerate, to_accelerate from tests.test_offload.conftest import torchrun from tests.testing_utils import requires_gpu pytest.importorskip("accelerate") def get_hf_dispatched_model(accel_device, tmp_path): from accelerate.big_modeling import dispatch_model offload_dir = tmp_path / "offload_dir" os.mkdir(offload_dir) model = torch.nn.Sequential( torch.nn.Linear(5, 5), torch.nn.Linear(5, 5), torch.nn.Linear(5, 5) ) if is_source_process(): dispatch_model( model, {"0": 0, "1": "cpu", "2": "disk"}, main_device=str(accel_device), force_hooks=True, offload_dir=offload_dir, ) else: model.to("meta") return model, offload_dir @pytest.mark.unit @requires_gpu def test_conversion_lifecycle(accel_device, tmp_path): model, offload_dir = get_hf_dispatched_model(accel_device, tmp_path) exp_device_map = { "": (None, None), "0": (accel_device, accel_device), "1": (accel_device, torch.device("cpu")), "2": (accel_device, "disk"), } exp_hf_device_map = {"": "cpu", "0": str(accel_device), "1": "cpu", "2": "disk"} # 1. from_accelerate (oneshot/ load_offloaded_model) device_map, _offload_dir = from_accelerate(model) assert device_map == exp_device_map with disable_onloading(): state_dict = model.state_dict(keep_vars=True) # 2. to_accelerate (transformers save) hf_device_map = to_accelerate(model) assert hf_device_map == exp_hf_device_map # 3. from_accelerate (post-save restore) device_map, _offload_dir = from_accelerate(model) assert device_map == exp_device_map # Note that rank 0 tensor pointers remain unchanged: # - no extra gpu/cpu/disk memory is allocated # - disk index remains valid with disable_onloading(): assert model.state_dict(keep_vars=True) == state_dict # 4. Just for completeness, test final invert (not part of lifecycle) hf_device_map = to_accelerate(model) assert hf_device_map == exp_hf_device_map @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_conversion_lifecycle_dist(accel_device, tmp_path): test_conversion_lifecycle(accel_device, tmp_path) vllm-project-compressed-tensors-c18a0fa/tests/test_offload/convert/test_from_accelerate.py000066400000000000000000000164261521257237700325150ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import hashlib import os from pathlib import Path import pytest import torch import torch.distributed as dist from compressed_tensors.distributed import is_source_process from compressed_tensors.offload import ( disable_onloading, from_accelerate, load_offloaded_model, ) from compressed_tensors.offload.cache import CPUCache, DeviceCache, DiskCache from compressed_tensors.offload.convert.from_accelerate import ( remove_accelerate_from_module, ) from tests.test_offload.conftest import torchrun from tests.testing_utils import requires_gpu from transformers import AutoModelForCausalLM acclerate = pytest.importorskip("accelerate") @pytest.mark.unit @requires_gpu def test_remove_accelerate_from_module_device(accel_device): # there"s no way to force accelerate to "offload" to Torch accelerator. Instead, # it just stays on Torch accelerator with no hooks current_accelerator = torch.accelerator.current_accelerator() linear = torch.nn.Linear(5, 5, device=current_accelerator) assert remove_accelerate_from_module(linear) == (accel_device, accel_device, None) assert not hasattr(linear, "_hf_hook") # test idempotency assert remove_accelerate_from_module(linear) == (accel_device, accel_device, None) assert not hasattr(linear, "_hf_hook") @pytest.mark.unit @requires_gpu def test_remove_accelerate_from_module_cpu(accel_device): from accelerate.big_modeling import dispatch_model current_accelerator = torch.accelerator.current_accelerator() linear = torch.nn.Linear(5, 5) dispatch_model( linear, {"": "cpu"}, main_device=current_accelerator, state_dict=linear.state_dict(), force_hooks=True, ) assert remove_accelerate_from_module(linear) == ( accel_device, torch.device("cpu"), None, ) assert not hasattr(linear, "_hf_hook") @pytest.mark.unit @requires_gpu @pytest.mark.filterwarnings("ignore::UserWarning") def test_remove_accelerate_from_module_disk(accel_device, tmp_path): # `disk_offload` is a super buggy function, and not reflective of real dispatches # `dispatch_model` is also super buggy, and requires at least one cpu device from accelerate.big_modeling import dispatch_model current_accelerator = torch.accelerator.current_accelerator() offload_dir = tmp_path / "offload_dir" os.mkdir(offload_dir) linear = torch.nn.Linear(5, 5) model = torch.nn.Sequential(linear) dispatch_model( model, {"0": "disk", "fake_module": "cpu"}, main_device=current_accelerator, force_hooks=True, offload_dir=offload_dir, ) assert remove_accelerate_from_module(linear) == (accel_device, "disk", offload_dir) assert not hasattr(linear, "_hf_hook") @pytest.mark.unit @requires_gpu def test_from_accelerate(accel_device, tmp_path): from accelerate.big_modeling import dispatch_model offload_dir = tmp_path / "offload_dir" os.mkdir(offload_dir) model = torch.nn.Sequential( torch.nn.Linear(5, 5), torch.nn.Linear(5, 5), torch.nn.Linear(5, 5) ) if is_source_process(): dispatch_model( model, {"0": 0, "1": "cpu", "2": "disk"}, main_device=str(accel_device), force_hooks=True, offload_dir=offload_dir, ) else: model.to("meta") device_map, _offload_dir = from_accelerate(model) # cuda is index agnostic when distributed assert device_map == { "": (None, None), "0": (accel_device, accel_device), "1": (accel_device, torch.device("cpu")), "2": (accel_device, "disk"), } if is_source_process(): assert _offload_dir == offload_dir assert isinstance(model[0]._parameters, DeviceCache) assert isinstance(model[1]._parameters, CPUCache) assert isinstance(model[2]._parameters, DiskCache) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_from_accelerate_dist(accel_device, tmp_path): test_from_accelerate(accel_device, tmp_path) @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) @torch.no_grad() def test_dist_disk_safetensors_update(tmp_path): offload_folder = tmp_path / "offload_folder" os.makedirs(offload_folder, exist_ok=True) with load_offloaded_model(): model = AutoModelForCausalLM.from_pretrained( "Qwen/Qwen3-0.6B", dtype="auto", device_map="auto_offload", max_memory={"cpu": 6e8}, offload_folder=str(offload_folder), ) # Get the model checkpoint files and hash them if dist.get_rank() == 0: checkpoint_files = {} for file_path in Path(offload_folder).glob("*.safetensors"): if not file_path.name.startswith(DiskCache._ct_file_prefix): with open(file_path, "rb") as f: checkpoint_files[file_path.name] = hashlib.sha256( f.read() ).hexdigest() dist.barrier() # Each rank updates a different module's tensor rank_0_module = model.model.layers[-1].self_attn.q_proj rank_1_module = model.model.layers[-1].self_attn.k_proj rank = dist.get_rank() if rank == 0: rank_0_module.weight *= 0 elif rank == 1: rank_1_module.weight *= 0 rank_1_module.weight += 1 dist.barrier() # Check that onloaded values are updated across ranks assert torch.all(rank_0_module.weight == 0) assert torch.all(rank_1_module.weight == 1) # Compare model checkpoint files and make sure they're unchanged if dist.get_rank() == 0: for file_name, original_hash in checkpoint_files.items(): file_path = offload_folder / file_name with open(file_path, "rb") as f: current_hash = hashlib.sha256(f.read()).hexdigest() assert current_hash == original_hash # Check that the files exist and are not symlinks with disable_onloading(): q_file_path = DiskCache.index[rank_0_module.weight]["safetensors_file"] k_file_path = DiskCache.index[rank_1_module.weight]["safetensors_file"] if dist.get_rank() == 0: assert os.path.exists(q_file_path) assert not os.path.islink(q_file_path) if dist.get_rank() == 1: assert os.path.exists(k_file_path) assert not os.path.islink(k_file_path) # Delete the parameters dist.barrier() delattr(rank_0_module, "weight") delattr(rank_1_module, "weight") # Check that the new files were deleted dist.barrier() assert not os.path.exists(q_file_path) assert not os.path.exists(k_file_path) # Compare model checkpoint files again and make sure they're still unchanged if dist.get_rank() == 0: for file_name, original_hash in checkpoint_files.items(): file_path = offload_folder / file_name with open(file_path, "rb") as f: current_hash = hashlib.sha256(f.read()).hexdigest() assert current_hash == original_hash vllm-project-compressed-tensors-c18a0fa/tests/test_offload/convert/test_to_accelerate.py000066400000000000000000000052311521257237700321640ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os import pytest import torch from compressed_tensors.offload import dispatch_with_map, offload_module, to_accelerate from compressed_tensors.offload.convert.helpers import norm_device from compressed_tensors.offload.convert.to_accelerate import to_accelerate_module from tests.test_offload.conftest import torchrun from tests.testing_utils import requires_gpu acclerate = pytest.importorskip("accelerate") def get_offload_devices() -> list[str]: offload_devices = ["cpu", "disk"] accelerator_device = torch.accelerator.current_accelerator() if accelerator_device is None: return offload_devices offload_devices.append(accelerator_device.type) multi_gpu_supported_types = ["cuda", "xpu", "hip", "hpu"] if accelerator_device.type in multi_gpu_supported_types: offload_devices.append(f"{accelerator_device.type}:0") return offload_devices @pytest.mark.unit @requires_gpu @pytest.mark.parametrize("offload_device", get_offload_devices()) def test_to_accelerate_module(offload_device, tmp_path): accelerator_device = torch.accelerator.current_accelerator() linear = torch.nn.Linear(5, 5) if offload_device == "disk": offload_dir = tmp_path / "offload_dir" os.mkdir(offload_dir) offload_module( linear, accelerator_device, offload_device, offload_dir=str(offload_dir) ) else: offload_module(linear, accelerator_device, offload_device) _offload_device = to_accelerate_module(linear, name="", hf_disk_index={}) assert _offload_device == str(norm_device(offload_device)) @pytest.mark.unit @requires_gpu def test_to_accelerate(accel_device, tmp_path): offload_dir = tmp_path / "offload_dir" os.mkdir(offload_dir) current_accelerator = torch.accelerator.current_accelerator() model = torch.nn.Sequential( torch.nn.Linear(5, 5), torch.nn.Linear(5, 5), torch.nn.Linear(5, 5) ) device_map = { "0": (current_accelerator, torch.device("cpu")), "1": (current_accelerator, current_accelerator), "2": (current_accelerator, "disk"), } dispatch_with_map(model, device_map, offload_dir) hf_device_map = to_accelerate(model) assert hf_device_map == {"": "cpu", "0": "cpu", "1": str(accel_device), "2": "disk"} assert hasattr(model[0], "_hf_hook") assert hasattr(model[1], "_hf_hook") assert hasattr(model[2], "_hf_hook") @pytest.mark.unit @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_to_accelerate_dist(accel_device, tmp_path): test_to_accelerate(accel_device, tmp_path) vllm-project-compressed-tensors-c18a0fa/tests/test_offload/test_dispatch.py000066400000000000000000000211751521257237700275160ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from unittest.mock import patch import pytest import torch from compressed_tensors.offload.cache import CPUCache, OffloadCache from compressed_tensors.offload.dispatch import ( dispatch_model, get_device_memory, set_onload_device, ) from compressed_tensors.offload.utils import module_size from tests.test_offload.conftest import skip_if_mps_device from tests.testing_utils import requires_gpu from transformers import AutoModelForCausalLM, AutoTokenizer class Decoder(torch.nn.Module): def __init__(self): super().__init__() self.linear0 = torch.nn.Linear(5, 5) self.linear1 = torch.nn.Linear(5, 5) def forward(self, input): return self.linear1(self.linear0(input)) class Model(torch.nn.Module): _no_split_modules = ["Decoder"] def __init__(self): super().__init__() self.decoder0 = Decoder() self.decoder1 = Decoder() def forward(self, input): return self.decoder1(self.decoder0(input)) def assert_module_on_device(module: torch.nn.Module, device: torch.device | str): assert not isinstance(module._parameters, CPUCache) for name, param in module.named_parameters(): assert torch.device(param.device) == torch.device(device), name def assert_module_offloaded( module: torch.nn.Module, onload_device: torch.device | str, offload_device: torch.device | str, req_params: bool = False, ): for name, submodule in module.named_modules(): if isinstance(submodule, torch.nn.ModuleList): continue if req_params and module_size(submodule)[0] <= 0: continue assert isinstance(submodule._parameters, OffloadCache), name assert torch.device(submodule._parameters.onload_device) == torch.device( onload_device ) assert torch.device(submodule._parameters.offload_device) == torch.device( offload_device ) def has_memory_requirements(device_memory: dict[torch.device, int]): real_device_memory = get_device_memory() for key, req in device_memory.items(): if key not in real_device_memory or real_device_memory[key] < req: return False return True @pytest.mark.unit @skip_if_mps_device @requires_gpu def test_dispatch_one_device(): model = Model() device_memory = {torch.device("cuda:0"): module_size(model)} if not has_memory_requirements(device_memory): pytest.skip("Cannot perform one device dispatch test, not enough device memory") dispatch_model(model, device_memory=device_memory, extra_memory=0) assert_module_on_device(model, "cuda:0") @pytest.mark.unit @skip_if_mps_device @requires_gpu def test_dispatch_two_devices(): model = Model() device_memory = { torch.device("cuda:0"): module_size(model.decoder0), torch.device("cuda:1"): module_size(model) - module_size(model.decoder0), } if not has_memory_requirements(device_memory): pytest.skip("Cannot perform split dispatch test: not enough devices or memory") # first decoder on first device, rest on second device dispatch_model(model, device_memory=device_memory, extra_memory=0) assert_module_on_device(model.decoder0, "cuda:0") assert_module_on_device(model.decoder1, "cuda:1") @pytest.mark.unit @skip_if_mps_device @requires_gpu def test_dispatch_no_split(): model = Model() device_memory = { torch.device("cuda:0"): module_size(model.decoder0.linear0), torch.device("cuda:1"): module_size(model), } if not has_memory_requirements(device_memory): pytest.skip("Cannot perform split dispatch test: not enough devices or mem") # first device is skipped: all ends up on second device dispatch_model(model, device_memory=device_memory, extra_memory=0) assert_module_on_device(model, "cuda:1") @pytest.mark.unit @skip_if_mps_device @requires_gpu def test_dispatch_split(): model = Model() first_linear = model.decoder0.linear0 device_memory = { torch.device("cuda:0"): module_size(first_linear), torch.device("cuda:1"): module_size(model) - module_size(first_linear), } if not has_memory_requirements(device_memory): pytest.skip("Cannot perform split dispatch test: not enough devices or memory") # first linear on first device, rest on second device dispatch_model( model, device_memory=device_memory, no_split_modules=tuple(), extra_memory=0 ) assert_module_on_device(model.decoder0.linear0, "cuda:0") assert_module_on_device(model.decoder0.linear1, "cuda:1") assert_module_on_device(model.decoder1, "cuda:1") @pytest.mark.unit @skip_if_mps_device @requires_gpu def test_dispatch_offloaded(): model = Model() device_memory = { torch.device("cuda:0"): ( module_size(model.decoder0.linear0) + module_size(model.decoder1) ), } if not has_memory_requirements(device_memory): pytest.skip("Cannot perform split dispatch test: not enough devices or mem") with patch("compressed_tensors.offload.dispatch.get_module_sizes") as mock_sizes: # first two linears are disjoint, but not enough memory to fit decoder1 mock_sizes.return_value = [ (model.decoder0.linear0, module_size(model.decoder0.linear0)), (model.decoder0.linear1, module_size(model.decoder0.linear1)), (model.decoder1, module_size(model.decoder1)), ] # first linear stays onloaded # second linear is popped off to fit offloaded decoder1 dispatch_model( model, device_memory=device_memory, no_split_modules=tuple(), extra_memory=0 ) assert_module_on_device(model.decoder0.linear0, "cuda:0") assert_module_offloaded(model.decoder0.linear1, "cuda:0", "cpu") assert_module_offloaded(model.decoder1, "cuda:0", "cpu") @pytest.mark.integration @requires_gpu @pytest.mark.parametrize("model_id", ["nm-testing/tinysmokellama-3.2"]) @skip_if_mps_device @torch.inference_mode() def test_offload_and_dispatch_model(model_id): model = AutoModelForCausalLM.from_pretrained(model_id).eval() tokenizer = AutoTokenizer.from_pretrained(model_id) tied_tensors_size = model.lm_head.weight.nbytes device_memory = {torch.device("cuda:0"): module_size(model) + tied_tensors_size} if not has_memory_requirements(device_memory): pytest.skip("Cannot perform split dispatch test: not enough devices or mem") model.to("cuda:0") sample = tokenizer("Hello my name is", return_tensors="pt") sample = {k: v.to("cuda:0") for k, v in sample.items()} true_logits = model(**sample).logits # offload entire model model.to("cpu") model = set_onload_device(model, "cuda:0") offloaded_logits = model(**sample).logits for module in model.modules(): assert_module_offloaded(module, "cuda:0", torch.device("cpu")) assert torch.allclose(offloaded_logits, true_logits) # dispatch model and fits model = dispatch_model(model, device_memory=device_memory, extra_memory=0) dispatched_logits = model(**sample).logits for module in model.modules(): assert_module_on_device(module, "cuda:0") assert torch.allclose(dispatched_logits, true_logits) # dispatch model with offload device_memory[torch.device("cuda:0")] = device_memory[torch.device("cuda:0")] // 2 model = dispatch_model(model, device_memory=device_memory, extra_memory=0) dispatched_logits = model(**sample).logits assert torch.allclose(dispatched_logits, true_logits) @pytest.mark.unit @pytest.mark.skip_xpu # mocks torch.accelerator for CPU fallback path def test_get_device_memory_cpu_fallback(): with patch("compressed_tensors.offload.dispatch.torch.accelerator") as mock_accel: mock_accel.is_available.return_value = False device_memory = get_device_memory() assert len(device_memory) == 1 assert torch.device("cpu") in device_memory assert device_memory[torch.device("cpu")] > 0 @pytest.mark.unit def test_dispatch_cpu_only(): model = Model() cpu_memory = module_size(model) * 2 device_memory = {torch.device("cpu"): cpu_memory} dispatch_model(model, device_memory=device_memory, extra_memory=0) assert_module_on_device(model, "cpu") @pytest.mark.unit @pytest.mark.skip_xpu # mocks torch.accelerator for CPU fallback path def test_dispatch_cpu_only_via_fallback(): model = Model() with patch("compressed_tensors.offload.dispatch.torch.accelerator") as mock_accel: mock_accel.is_available.return_value = False dispatch_model(model, extra_memory=0) assert_module_on_device(model, "cpu") vllm-project-compressed-tensors-c18a0fa/tests/test_offload/test_interface.py000066400000000000000000000242431521257237700276560ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import tempfile from pathlib import Path import pytest import torch from compressed_tensors.offload import ( align_module_device, align_modules, disable_offloading, disable_onloading, get_cache_init_kwargs, get_cache_kwargs, get_execution_device, get_offloaded_device, update_offload_parameter, ) from compressed_tensors.offload.cache import CPUCache from compressed_tensors.offload.module import offload_module from tests.test_offload.conftest import assert_device_equal, assert_tensor_equal from tests.testing_utils import requires_gpu ONLOAD_DEVICE = torch.accelerator.current_accelerator() OFFLOAD_DEVICE = torch.device("cpu") @pytest.fixture(scope="function") def cache(): return CPUCache(ONLOAD_DEVICE) @pytest.fixture(scope="function") def linear(): return torch.nn.Linear(5, 5, bias=True, device=OFFLOAD_DEVICE) @pytest.fixture(scope="function") def offloaded_linear(linear, cache): offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) return linear @pytest.mark.unit @requires_gpu def test_disable_offloading(): cache1 = CPUCache(ONLOAD_DEVICE) cache2 = CPUCache(ONLOAD_DEVICE) cache1["weight"] = torch.tensor(0, device=OFFLOAD_DEVICE) cache2["weight"] = torch.tensor(1, device=OFFLOAD_DEVICE) with disable_offloading(): assert cache1["weight"] in cache1.keep_onloaded_values.values() assert cache2["weight"] in cache2.keep_onloaded_values.values() @pytest.mark.unit @requires_gpu def test_disable_onloading(): cache1 = CPUCache(ONLOAD_DEVICE) cache2 = CPUCache(ONLOAD_DEVICE) cache1["weight"] = torch.tensor(0, device=OFFLOAD_DEVICE) cache2["weight"] = torch.tensor(1, device=OFFLOAD_DEVICE) with disable_onloading(): assert_device_equal(cache1["weight"].device, OFFLOAD_DEVICE) assert_device_equal(cache2["weight"].device, OFFLOAD_DEVICE) @pytest.mark.unit @requires_gpu @pytest.mark.parametrize("offload", (True, False)) def test_update_offload_parameter(linear: torch.nn.Linear, cache, offload): init_data = torch.tensor(0.0, device=OFFLOAD_DEVICE) linear.weight = torch.nn.Parameter(init_data, requires_grad=False) if offload: offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) assert linear.weight == 0 update_offload_parameter(linear, "weight", torch.tensor(1)) assert linear.weight == 1 with disable_offloading(): update_offload_parameter(linear, "weight", torch.tensor(2)) assert linear.weight == 2 assert linear.weight == 2 with disable_onloading(): update_offload_parameter(linear, "weight", torch.tensor(3)) assert linear.weight == 3 assert linear.weight == 3 @pytest.mark.unit def test_update_offload_parameter_with_grad(linear: torch.nn.Linear): zeros = torch.nn.Parameter(torch.zeros(5, 5), requires_grad=True) update_offload_parameter(linear, "weight", zeros) assert_tensor_equal(linear.weight, zeros) ones = torch.nn.Parameter(torch.ones(5, 5), requires_grad=True) offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) update_offload_parameter(linear, "weight", ones) assert_tensor_equal(linear.weight, ones, ONLOAD_DEVICE) @pytest.mark.unit @requires_gpu def test_get_execution_device(linear: torch.nn.Linear, cache): assert_device_equal(get_execution_device(linear), OFFLOAD_DEVICE) linear.to(ONLOAD_DEVICE) assert_device_equal(get_execution_device(linear), ONLOAD_DEVICE) linear.to(OFFLOAD_DEVICE) offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) assert_device_equal(get_execution_device(linear), ONLOAD_DEVICE) with disable_onloading(): assert_device_equal(get_execution_device(linear), ONLOAD_DEVICE) with disable_offloading(): assert_device_equal(get_execution_device(linear), ONLOAD_DEVICE) @pytest.mark.unit @requires_gpu def test_get_offloaded_device(linear: torch.nn.Linear, cache): assert_device_equal(get_offloaded_device(linear), OFFLOAD_DEVICE) linear.to(ONLOAD_DEVICE) assert_device_equal(get_offloaded_device(linear), ONLOAD_DEVICE) linear.to(OFFLOAD_DEVICE) offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) assert_device_equal(get_offloaded_device(linear), OFFLOAD_DEVICE) with disable_onloading(): assert_device_equal(get_offloaded_device(linear), OFFLOAD_DEVICE) with disable_offloading(): assert_device_equal(get_offloaded_device(linear), OFFLOAD_DEVICE) @pytest.mark.unit @requires_gpu def test_get_cache_kwargs_cpu(): """Test get_cache_kwargs for CPUCache.""" # Non-offloaded module should return empty dict linear = torch.nn.Linear(5, 5, device=OFFLOAD_DEVICE) kwargs = get_cache_kwargs(linear) assert kwargs == {} # With default provided default = {"custom": "value"} kwargs = get_cache_kwargs(linear, default=default) assert kwargs == {"custom": "value"} # Offloaded module with CPUCache (no extra kwargs) offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) kwargs = get_cache_kwargs(linear) assert kwargs == {} @pytest.mark.unit @requires_gpu def test_get_cache_kwargs_disk(): """Test get_cache_kwargs for DiskCache extracts offload_dir.""" with tempfile.TemporaryDirectory() as tmpdir: linear = torch.nn.Linear(5, 5, device="cpu") offload_module(linear, ONLOAD_DEVICE, "disk", offload_dir=tmpdir) kwargs = get_cache_kwargs(linear) assert kwargs == {"offload_dir": Path(tmpdir).resolve()} @pytest.mark.unit @requires_gpu def test_get_cache_init_kwargs_cpu(): """ Test using get_cache_init_kwargs to copy offloading from one module to another. """ # Create and offload a reference module ref_linear = torch.nn.Linear(3, 3, device=OFFLOAD_DEVICE) offload_module(ref_linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) # Get the init kwargs from the reference module kwargs = get_cache_init_kwargs(ref_linear) # Create a new module and apply the same offload settings new_linear = torch.nn.Linear(5, 5, device=OFFLOAD_DEVICE) offload_module(new_linear, **kwargs) # Verify the new module has the same offload settings assert isinstance(new_linear._parameters, CPUCache) assert_device_equal(new_linear._parameters.onload_device, ONLOAD_DEVICE) assert_device_equal(new_linear._parameters.offload_device, OFFLOAD_DEVICE) # Verify weights work correctly assert_device_equal(new_linear.weight.device, ONLOAD_DEVICE) @pytest.mark.unit @requires_gpu def test_get_cache_init_kwargs_disk(): """Test using get_cache_init_kwargs to copy disk offload settings.""" with tempfile.TemporaryDirectory() as tmpdir: # Create and offload a reference module with disk offloading ref_linear = torch.nn.Linear(3, 3, device="cpu") offload_module(ref_linear, ONLOAD_DEVICE, "disk", offload_dir=tmpdir) # Get the init kwargs from the reference module kwargs = get_cache_init_kwargs(ref_linear) # Create a new module and apply the same offload settings new_linear = torch.nn.Linear(5, 5, device="cpu") offload_module(new_linear, **kwargs) # Verify the new module has the same offload settings including offload_dir assert_device_equal(new_linear._parameters.onload_device, ONLOAD_DEVICE) assert new_linear._parameters.offload_device == "disk" assert new_linear._parameters.offload_dir == Path(tmpdir).resolve() # Verify weights work correctly assert_device_equal(new_linear.weight.device, ONLOAD_DEVICE) @pytest.mark.unit @requires_gpu def test_register_offload_module_cpu(linear: torch.nn.Linear): from compressed_tensors.offload import register_offload_module # Non-offloaded parent - submodule should not be offloaded sub1 = torch.nn.Linear(1, 1) register_offload_module(linear, "sub1", sub1) assert linear.sub1 is sub1 assert not isinstance(sub1._parameters, CPUCache) # Offloaded parent - submodule should inherit offloading offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) sub2 = torch.nn.Linear(1, 1) register_offload_module(linear, "sub2", sub2) assert linear.sub2 is sub2 assert isinstance(sub2._parameters, CPUCache) assert_device_equal(sub2._parameters.onload_device, ONLOAD_DEVICE) assert_device_equal(sub2._parameters.offload_device, OFFLOAD_DEVICE) assert_device_equal(sub2.weight.device, ONLOAD_DEVICE) @pytest.mark.unit @requires_gpu def test_register_offload_module_disk(): """Test register_offload_module inherits offload_dir from parent with DiskCache.""" from compressed_tensors.offload import register_offload_module with tempfile.TemporaryDirectory() as tmpdir: # Create a parent module with disk offloading parent = torch.nn.Linear(5, 5, device="cpu") offload_module(parent, ONLOAD_DEVICE, "disk", offload_dir=tmpdir) # Register a submodule - it should inherit the disk offloading settings sub = torch.nn.Linear(3, 3, device="cpu") register_offload_module(parent, "sub", sub) # Verify the submodule has the same offload settings including offload_dir assert parent.sub is sub assert_device_equal(sub._parameters.onload_device, ONLOAD_DEVICE) assert sub._parameters.offload_device == "disk" assert sub._parameters.offload_dir == Path(tmpdir).resolve() # Verify weights work correctly assert_device_equal(sub.weight.device, ONLOAD_DEVICE) @pytest.mark.unit @requires_gpu def test_align_modules(offloaded_linear: torch.nn.Linear): linear = torch.nn.Linear(1, 1, device=ONLOAD_DEVICE) with align_modules((linear, offloaded_linear), OFFLOAD_DEVICE): assert_device_equal(linear.weight.device, OFFLOAD_DEVICE) assert_device_equal(offloaded_linear.weight.device, OFFLOAD_DEVICE) @pytest.mark.unit @requires_gpu @pytest.mark.parametrize("offload", (True, False)) def test_align_module_device(linear: torch.nn.Linear, cache, offload): if offload: offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) else: linear.to(ONLOAD_DEVICE) with align_module_device(linear, OFFLOAD_DEVICE): assert_device_equal(linear.weight.device, OFFLOAD_DEVICE) vllm-project-compressed-tensors-c18a0fa/tests/test_offload/test_load.py000066400000000000000000000105411521257237700266310ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from unittest.mock import MagicMock, patch import pytest import torch from compressed_tensors.offload import ( disable_onloading, from_accelerate, get_offloaded_device, ) from compressed_tensors.offload.convert import to_accelerate from compressed_tensors.offload.convert.from_accelerate import _infer_module_device from compressed_tensors.offload.load import load_offloaded_model from tests.test_offload.conftest import ( assert_device_equal, skip_if_mps_device, torchrun, ) from tests.testing_utils import requires_gpu from transformers import AutoModelForCausalLM acclerate = pytest.importorskip("accelerate") accelerator_device = torch.accelerator.current_accelerator() TEST_PARAMETERS = [ ( "auto", {0: 596049920, "cpu": 1e15}, # force cpu offload for testing accelerator_device, torch.device("cpu"), ), ( accelerator_device.type, None, accelerator_device, accelerator_device, ), ( "cpu", None, torch.device("cpu"), torch.device("cpu"), ), ( "auto_offload", {"cpu": 596049920}, # force disk offload for testing torch.device("cpu"), "disk", ), ] @pytest.mark.integration @requires_gpu @pytest.mark.parametrize("device_map,max_memory,first,second", TEST_PARAMETERS) def test_load(device_map, max_memory, first, second, tmp_path): with load_offloaded_model(AutoModelForCausalLM): model = AutoModelForCausalLM.from_pretrained( "Qwen/Qwen3-0.6B", device_map=device_map, max_memory=max_memory, dtype=torch.bfloat16, offload_folder=str(tmp_path / "disk_offload"), ) for layer_index in range(0, 8): module = model.get_submodule(f"model.layers.{layer_index}.self_attn.q_proj") assert_device_equal(get_offloaded_device(module), first) for layer_index in range(8, 28): module = model.get_submodule(f"model.layers.{layer_index}.self_attn.q_proj") assert_device_equal(get_offloaded_device(module), second) with disable_onloading(): state_dict = model.state_dict(keep_vars=True) to_accelerate(model) for layer_index in range(0, 8): module = model.get_submodule(f"model.layers.{layer_index}.self_attn.q_proj") assert_device_equal(_get_accelerate_offloaded_device(module), first) for layer_index in range(8, 28): module = model.get_submodule(f"model.layers.{layer_index}.self_attn.q_proj") assert_device_equal(_get_accelerate_offloaded_device(module), second) model.save_pretrained(tmp_path / "save_path") from_accelerate(model) # TODO: accelerate's disk onloading implementation does not keep consistent meta # tensors, :. tensor pointers change and cannot be converted back properly if second != "disk": with disable_onloading(): assert model.state_dict(keep_vars=True) == state_dict @pytest.mark.integration @requires_gpu(2) @torchrun(world_size=2, init_dist=True) def test_load_dist(tmp_path): for parameters in TEST_PARAMETERS: test_load(*parameters, tmp_path=tmp_path) def _get_accelerate_offloaded_device(module: torch.nn.Module) -> str | None: device = _infer_module_device(module) if device == torch.device("meta"): return "disk" return device @pytest.mark.unit @skip_if_mps_device @patch("compressed_tensors.offload.load.from_accelerate") def test_patch_forwards_positional_args(mock_from_accelerate): """Regression: positional args must be forwarded without rebinding to cls.""" received = {} class FakeModel: @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): received["cls"] = cls received["path"] = pretrained_model_name_or_path received["model_args"] = model_args received["kwargs"] = kwargs return MagicMock() with load_offloaded_model(FakeModel, extra_cpu_mem=0): FakeModel.from_pretrained("org/model", device_map="cpu", torch_dtype="auto") assert received["cls"] is FakeModel assert received["path"] == "org/model" assert received["kwargs"]["device_map"] == "cpu" assert received["kwargs"]["torch_dtype"] == "auto" vllm-project-compressed-tensors-c18a0fa/tests/test_offload/test_module.py000066400000000000000000000162131521257237700272010ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import gc import inspect from weakref import ref import pytest import torch from compressed_tensors.offload import disable_offloading, disable_onloading from compressed_tensors.offload.cache.cpu import CPUCache from compressed_tensors.offload.module import offload_module from tests.test_offload.conftest import assert_device_equal from tests.testing_utils import requires_gpu ONLOAD_DEVICE = torch.accelerator.current_accelerator() OFFLOAD_DEVICE = torch.device("cpu") @pytest.fixture(scope="function") def cache(): return CPUCache(ONLOAD_DEVICE) @pytest.fixture(scope="function") def linear(): return torch.nn.Linear(5, 5, bias=True, device=OFFLOAD_DEVICE) @pytest.fixture(scope="function") def offloaded_linear(linear, cache): offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) return linear @pytest.fixture(scope="function") def input(): return torch.zeros(6, device=OFFLOAD_DEVICE) @pytest.mark.unit @requires_gpu def test_onloading(linear: torch.nn.Linear, cache): weight = linear.weight bias = linear.bias offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) onloaded_weight = linear.weight onloaded_bias = linear.bias assert_device_equal(onloaded_weight.device, ONLOAD_DEVICE) assert_device_equal(onloaded_bias.device, ONLOAD_DEVICE) assert type(onloaded_weight) is type(weight) assert type(onloaded_bias) is type(bias) assert torch.equal(onloaded_weight.to(weight.device), weight) assert torch.equal(onloaded_bias.to(bias.device), bias) @pytest.mark.unit @requires_gpu def test_garbage_collect(offloaded_linear: torch.nn.Linear): weight_ref = ref(offloaded_linear.weight) bias_ref = ref(offloaded_linear.bias) del offloaded_linear gc.collect() assert weight_ref() is None assert bias_ref() is None @pytest.mark.unit @requires_gpu def test_disable_offloading(offloaded_linear: torch.nn.Linear): outside_onloaded = offloaded_linear.weight outside_onloaded_ref = ref(outside_onloaded) assert_device_equal(outside_onloaded.device, ONLOAD_DEVICE) with disable_offloading(): inside_onloaded = offloaded_linear.weight inside_onloaded_ref = ref(inside_onloaded) assert_device_equal(inside_onloaded.device, ONLOAD_DEVICE) del outside_onloaded del inside_onloaded gc.collect() assert outside_onloaded_ref() is None assert inside_onloaded_ref() is not None assert outside_onloaded_ref() is None assert inside_onloaded_ref() is None @pytest.mark.unit @requires_gpu def test_disable_onloading(linear: torch.nn.Linear, cache): offloaded_weight = linear.weight offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) with disable_onloading(): weight = linear.weight assert weight is offloaded_weight # new parameter assignments are direct new_param = torch.nn.Parameter(torch.ones(5, device=ONLOAD_DEVICE)) linear.new_param = new_param assert linear.new_param is new_param assert weight is offloaded_weight @pytest.mark.unit @requires_gpu def test_delete(offloaded_linear: torch.nn.Linear): weight_ref = ref(offloaded_linear.weight) bias_ref = ref(offloaded_linear.bias) del offloaded_linear.weight del offloaded_linear.bias gc.collect() assert weight_ref() is None assert bias_ref() is None @pytest.mark.unit @requires_gpu def test_forward_call(linear: torch.nn.Linear, cache): def forward(self, input: torch.Tensor) -> torch.Tensor: assert_device_equal(input.device, ONLOAD_DEVICE) return torch.nn.functional.linear(input, linear.weight, linear.bias) linear.forward = forward.__get__(linear) offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) with torch.no_grad(): input = torch.zeros(5, device=OFFLOAD_DEVICE) output = linear.forward(input) assert_device_equal(output.device, ONLOAD_DEVICE) @pytest.mark.unit @pytest.mark.parametrize("param_device", (ONLOAD_DEVICE, OFFLOAD_DEVICE)) @pytest.mark.parametrize("use_register_parameter", (True, False)) @pytest.mark.parametrize("requires_grad", (True, False)) def test_register_parameter( offloaded_linear: torch.nn.Linear, param_device, use_register_parameter, requires_grad, ): # register param data = torch.ones(5, device=param_device) param = torch.nn.Parameter(data, requires_grad=requires_grad) if use_register_parameter: offloaded_linear.register_parameter("param_name", param) else: offloaded_linear.param_name = param # new param is correctly onloaded assert_device_equal(offloaded_linear.param_name.device, ONLOAD_DEVICE) assert torch.equal(offloaded_linear.param_name.to(param_device), param) @pytest.mark.unit @pytest.mark.parametrize("param_device", (ONLOAD_DEVICE, OFFLOAD_DEVICE)) @pytest.mark.parametrize("use_register_parameter", (True, False)) @pytest.mark.parametrize("requires_grad", (True, False)) def test_register_parameter_invalidates( offloaded_linear: torch.nn.Linear, param_device, use_register_parameter, requires_grad, ): with disable_offloading(): # original weight is kept onloaded onloaded_weight = offloaded_linear.weight assert onloaded_weight in set(CPUCache.keep_onloaded_values.values()) # add new param data = torch.ones((5, 5), device=param_device) param = torch.nn.Parameter(data, requires_grad=requires_grad) if use_register_parameter: offloaded_linear.register_parameter("weight", param) else: offloaded_linear.weight = param # new param is correct assert_device_equal(offloaded_linear.weight.device, ONLOAD_DEVICE) assert torch.equal(offloaded_linear.weight, param.to(ONLOAD_DEVICE)) @pytest.mark.unit def test_forward_signature(linear: torch.nn.Linear, cache): original_signature = inspect.signature(linear.forward) offload_module(linear, ONLOAD_DEVICE, OFFLOAD_DEVICE) assert inspect.signature(linear.forward) == original_signature @pytest.mark.unit def test_set_item(offloaded_linear: torch.nn.Linear): # update update = torch.nn.Parameter( torch.rand(5, 5, device=OFFLOAD_DEVICE), requires_grad=False ) offloaded_linear.weight = update with disable_onloading(): assert offloaded_linear.weight is not update # overwrite with different size overwrite = torch.nn.Parameter( torch.rand(6, 6, device=OFFLOAD_DEVICE), requires_grad=False ) offloaded_linear.weight = overwrite with disable_onloading(): assert offloaded_linear.weight is overwrite @pytest.mark.unit def test_set_item_buffers(offloaded_linear: torch.nn.Linear): # common case: registering buffers of difference sizes twice new = torch.rand(5) offloaded_linear.register_buffer("buffer", new, persistent=False) with disable_onloading(): assert offloaded_linear.buffer is new overwrite = torch.rand(6) offloaded_linear.register_buffer("buffer", overwrite, persistent=False) with disable_onloading(): assert offloaded_linear.buffer is overwrite vllm-project-compressed-tensors-c18a0fa/tests/test_offload/test_xpu_routing.py000066400000000000000000000100671521257237700303000ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ XPU emulation test (part 1): Mock routing tests for XPU device type. Verifies that routing functions handle ``"xpu"`` correctly by mocking ``torch.accelerator`` — no real tensor operations, no GPU required. These tests are skipped when ``--emulate-xpu`` is active because the global ``is_accelerator_type`` patch conflicts with the per-test monkeypatch. """ from types import SimpleNamespace import pytest import torch from compressed_tensors.offload.cache.base import OffloadCache from compressed_tensors.offload.cache.device import DeviceCache from compressed_tensors.offload.convert.helpers import norm_device from compressed_tensors.utils import is_accelerator_type @pytest.fixture def mock_xpu_accelerator(monkeypatch): """Mock torch.accelerator to report XPU as the current device.""" fake = SimpleNamespace(type="xpu") monkeypatch.setattr(torch.accelerator, "current_accelerator", lambda: fake) monkeypatch.setattr(torch.accelerator, "is_available", lambda: True) monkeypatch.setattr(torch.accelerator, "device_count", lambda: 1) skipif_emulate_xpu = pytest.mark.skipif( "config.getoption('--emulate-xpu', default=False)", reason="Option 1 routing tests conflict with --emulate-xpu global patches", ) @pytest.mark.unit @skipif_emulate_xpu class TestXpuRouting: """Verify that routing functions correctly handle 'xpu' as the accelerator type. These tests mock torch.accelerator without real tensor operations — they validate the decision logic that changed in the torch.accelerator migration. """ def test_is_accelerator_type_xpu(self, mock_xpu_accelerator): assert is_accelerator_type("xpu") is True assert is_accelerator_type("cuda") is False assert is_accelerator_type("cpu") is False def test_is_accelerator_type_unavailable(self, monkeypatch): monkeypatch.setattr(torch.accelerator, "is_available", lambda: False) assert is_accelerator_type("xpu") is False def test_cache_routes_device_cache_for_xpu(self, mock_xpu_accelerator): cache_cls = OffloadCache.cls_from_device(torch.device("xpu", 0)) assert cache_cls is DeviceCache def test_norm_device_resolves_xpu_to_index_0(self, mock_xpu_accelerator): result = norm_device("xpu") assert result == torch.device("xpu", 0) def test_norm_device_preserves_xpu_with_index(self, mock_xpu_accelerator): result = norm_device(torch.device("xpu", 0)) assert result == torch.device("xpu", 0) def test_norm_device_cpu_unaffected(self, mock_xpu_accelerator): result = norm_device("cpu") assert result == torch.device("cpu") def test_get_safe_open_device_xpu(self, mock_xpu_accelerator): from compressed_tensors.offload.cache.disk import _get_safe_open_device # bare "xpu" → current device index (0) result = _get_safe_open_device(torch.device("xpu")) assert result == "xpu:0" def test_get_safe_open_device_xpu_with_index(self, mock_xpu_accelerator): from compressed_tensors.offload.cache.disk import _get_safe_open_device result = _get_safe_open_device(torch.device("xpu", 3)) assert result == "xpu:3" def test_get_safe_open_device_cuda_returns_string(self, monkeypatch): from compressed_tensors.offload.cache.disk import _get_safe_open_device fake = SimpleNamespace(type="cuda") monkeypatch.setattr(torch.accelerator, "current_accelerator", lambda: fake) monkeypatch.setattr(torch.accelerator, "is_available", lambda: True) monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 2) assert _get_safe_open_device(torch.device("cuda")) == "cuda:2" assert _get_safe_open_device(torch.device("cuda", 5)) == "cuda:5" def test_get_safe_open_device_cpu(self, mock_xpu_accelerator): from compressed_tensors.offload.cache.disk import _get_safe_open_device result = _get_safe_open_device(torch.device("cpu")) assert result == "cpu" vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/000077500000000000000000000000001521257237700254145ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/__init__.py000066400000000000000000000001531521257237700275240ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/lifecycle/000077500000000000000000000000001521257237700273535ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/lifecycle/__init__.py000066400000000000000000000001531521257237700314630ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/lifecycle/conftest.py000066400000000000000000000021061521257237700315510ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import List, Optional import pytest import torch from compressed_tensors.quantization.quant_args import QuantizationArgs from compressed_tensors.quantization.quant_config import QuantizationStatus from compressed_tensors.quantization.quant_scheme import QuantizationScheme @pytest.fixture def create_quantization_scheme(): def quantization_scheme( targets: List[str], weights: Optional[QuantizationArgs] = None, input_activations: Optional[QuantizationArgs] = None, output_activations: Optional[QuantizationArgs] = None, ): return QuantizationScheme( targets=targets, weights=weights, input_activations=input_activations, output_activations=output_activations, ) return quantization_scheme @pytest.fixture def mock_frozen(): def update_status(model: torch.nn.Module): model.quantization_status = QuantizationStatus.FROZEN return update_status vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/lifecycle/test_apply.py000066400000000000000000000410101521257237700321050ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import re import shutil from typing import Optional from unittest.mock import MagicMock import pytest import torch from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import ( DEFAULT_QUANTIZATION_METHOD, FP8_E4M3_DATA, QuantizationArgs, QuantizationConfig, QuantizationScheme, QuantizationStatus, QuantizationStrategy, QuantizationType, ) from compressed_tensors.quantization.lifecycle import apply_quantization_config from compressed_tensors.utils import is_match, match_named_modules from transformers import AutoModelForCausalLM @pytest.fixture(scope="module", autouse=True) def cleanup_model_cache(): """Clean up the test model cache directory after all tests complete.""" yield try: shutil.rmtree("test-apply-model-cache", ignore_errors=True) except Exception: pass @pytest.fixture def mock_model(): model = MagicMock() model.named_modules.return_value = [ ("layer1", MagicMock()), ("layer2", MagicMock()), ("layer3", MagicMock()), ] return model @pytest.fixture def mock_module(): return MagicMock() @pytest.fixture def llama_stories_model(): return AutoModelForCausalLM.from_pretrained( "Xenova/llama2.c-stories15M", torch_dtype="auto", cache_dir="test-apply-model-cache", ) def test_target_prioritization(mock_frozen): # tests that the config_groups are applied in the correct order # of priority, where exact layer name > regex > module name config = { "quant_method": "compressed-tensors", "format": "fakequant", "config_groups": { "group_1": { "weights": { "num_bits": 8, }, "targets": ["Linear"], }, "group_2": { "weights": { "num_bits": 4, }, "targets": ["re:.*down_proj"], }, "group_3": { "weights": { "num_bits": 2, }, "targets": ["model.layers.0.mlp.down_proj"], }, }, } model = AutoModelForCausalLM.from_pretrained( "HuggingFaceM4/tiny-random-LlamaForCausalLM", torch_dtype="auto", cache_dir="test-apply-model-cache", ) model.eval() config = QuantizationConfig(**config) config.quantization_status = QuantizationStatus.CALIBRATION apply_quantization_config(model, config) mock_frozen(model) for name, module in model.named_modules(): if name == "model.layers.0.mlp.down_proj": assert module.quantization_scheme.weights.num_bits == 2 elif re.match(".*down_proj", name): assert module.quantization_scheme.weights.num_bits == 4 elif isinstance(module, torch.nn.Linear): assert module.quantization_scheme.weights.num_bits == 8 def test_apply_quantization_config_tinyllama(): quant_config = get_sample_tinyllama_quant_config( status=QuantizationStatus.INITIALIZED ) model = get_tinyllama_model() # check that model is not already quantized for module in model.modules(): _test_layer_quantization_status(module, inputs=False, weights=False) # apply quant config to model apply_quantization_config(model, quant_config) # check for correct application of quant config for quant_scheme in quant_config.config_groups.values(): for name, module in match_named_modules( model, quant_scheme.targets, quant_config.ignore ): _test_layer_quantization_status( module, inputs=quant_scheme.input_activations is not None, weights=quant_scheme.weights is not None, expected_status=QuantizationStatus.INITIALIZED, ) @pytest.mark.parametrize( "config", [ QuantizationConfig( config_groups={ "linear": QuantizationScheme( targets=["Linear"], input_activations=QuantizationArgs( num_bits=8, type="float", strategy="tensor", scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=torch.float, ), ) } ), QuantizationConfig( config_groups={ "linear": QuantizationScheme( targets=["Linear"], input_activations=QuantizationArgs( num_bits=8, type="float", strategy="tensor", scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=torch.float, ), ) }, ignore=[ "model.layers.0.self_attn.q_proj", "model.layers.1.self_attn.k_proj", "model.layers.2.self_attn.v_proj", ], ), QuantizationConfig( config_groups={}, kv_cache_scheme=QuantizationArgs( num_bits=8, type="float", strategy="tensor", scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=torch.float, ), ), QuantizationConfig( config_groups={ "attention": QuantizationScheme( targets=["LlamaAttention"], input_activations=QuantizationArgs( num_bits=8, type="float", strategy="tensor", scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=torch.float, ), ) }, kv_cache_scheme=QuantizationArgs( num_bits=8, type="float", strategy="tensor", scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=torch.float, ), ), ], ) def test_from_pretrained(config: QuantizationConfig): model = AutoModelForCausalLM.from_pretrained("nm-testing/llama2.c-stories15M") apply_quantization_config(model, config) _config = QuantizationConfig.from_pretrained(model) assert list(_config.config_groups.values()) == list(config.config_groups.values()) assert _config.kv_cache_scheme == config.kv_cache_scheme assert _config.ignore == config.ignore def test_serialize_config_tinyllama(): quant_config = get_sample_tinyllama_quant_config() model = get_tinyllama_model() # check that model is not already quantized for module in model.modules(): _test_layer_quantization_status(module, inputs=False, weights=False) # apply quant config to model apply_quantization_config(model, quant_config) serialized_config = QuantizationConfig.from_pretrained(model) assert len(serialized_config.config_groups) == 2 assert serialized_config.config_groups["group_0"].targets == ["Embedding"] assert serialized_config.config_groups["group_0"].input_activations is None assert serialized_config.config_groups["group_1"].targets == ["Linear"] assert serialized_config.config_groups["group_1"].input_activations is not None assert serialized_config.format == CompressionFormat.dense.value assert serialized_config.quant_method == DEFAULT_QUANTIZATION_METHOD assert serialized_config.ignore == ["model.layers.1.mlp.down_proj"] if serialized_config.global_compression_ratio is not None: assert serialized_config.global_compression_ratio > 1.0 assert serialized_config.global_compression_ratio < 8.0 def _test_layer_quantization_status( module, inputs: bool, weights: bool, expected_status: Optional[QuantizationStatus] = None, expected_dtype: Optional[torch.dtype] = None, ): # check if quantization is applied at all (true if inputs or weights targeted) quantized = inputs or weights assert hasattr(module, "quantization_scheme") == quantized assert hasattr(module, "quantization_status") == quantized if expected_status is not None: assert module.quantization_status is expected_status # check inputs matches expected assert hasattr(module, "input_scale") == inputs assert hasattr(module, "input_zero_point") == inputs # check weights matches expected assert hasattr(module, "weight_scale") == weights assert hasattr(module, "weight_zero_point") == weights if weights and expected_dtype is not None: assert module.weight.dtype is expected_dtype def get_tinyllama_model(): return AutoModelForCausalLM.from_pretrained( "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T", torch_dtype="auto", cache_dir="test-apply-model-cache", ) def get_sample_tinyllama_quant_config( status: QuantizationStatus = QuantizationStatus.FROZEN, ): config_dict = { "quant_method": "compressed-tensors", "format": "fakequant", "quantization_status": status, "global_compression_ratio": None, "config_groups": { "group_1": { "weights": { "num_bits": 8, "type": "int", "symmetric": True, "strategy": "tensor", }, "input_activations": { "num_bits": 8, "type": "int", "symmetric": True, "strategy": "tensor", }, "targets": ["Linear"], }, "group_2": { "weights": { "num_bits": 8, "type": "int", "symmetric": False, "strategy": "tensor", }, "input_activations": None, "targets": ["Embedding"], }, }, "ignore": ["LlamaRotaryEmbedding", "model.layers.1.mlp.down_proj"], } return QuantizationConfig.model_validate(config_dict) @pytest.mark.parametrize( "target,should_raise_warning", [ [("Linear",), False], [("Linear", "re:.*foobarbaz"), True], ], ) def test_apply_quantization_config(caplog, target, should_raise_warning): import logging # load a dense, unquantized tiny llama model model = get_tinyllama_model() quantization_config_dict = { "quant_method": "compressed-tensors", "format": "pack-quantized", "global_compression_ratio": None, "config_groups": { "group_1": { "weights": { "num_bits": 4, "type": "int", "symmetric": False, "strategy": "tensor", }, "targets": target, } }, "ignore": ["lm_head", "re:.*gate"], } config = QuantizationConfig(**quantization_config_dict) config.quantization_status = QuantizationStatus.CALIBRATION # mismatch in the ignore key of quantization_config_dict with caplog.at_level(logging.WARNING): apply_quantization_config(model, config) if should_raise_warning: assert len(caplog.text) > 0 else: assert len(caplog.text) == 0 def test_multi_apply_quantization_config(): """ Ensure that multiple quantization configs are applied correctly If quantization config was previously applied to a module, those changes should be reset for newly applied quantization config """ model = get_tinyllama_model() # FP8 applied to self_attn qconfig1 = QuantizationConfig( config_groups={ "group_0": QuantizationScheme( targets=[ r"re:.*self_attn\.(k|q|o|v)_proj$", ], weights=QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.TENSOR, symmetric=True, dynamic=False, ), input_activations=QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.TENSOR, symmetric=True, dynamic=False, ), ) }, ignore=["lm_head"], ) # W4A16_ASYM applied to mlp and self_attn.o_proj to validate overwriting qconfig2 = QuantizationConfig( config_groups={ "group_0": QuantizationScheme( targets=[ r"re:.*mlp\.(down|gate|up)_proj$", r"re:.*self_attn\.o_proj$", ], weights=QuantizationArgs( num_bits=4, type=QuantizationType.INT, strategy=QuantizationStrategy.GROUP, group_size=128, symmetric=False, dynamic=False, ), ) }, ignore=["lm_head"], ) apply_quantization_config(model, qconfig1) apply_quantization_config(model, qconfig2) for name, module in model.named_modules(): if is_match( name, module, qconfig2.config_groups["group_0"].targets, qconfig2.ignore ): # assert W4A16_ASYM parameters are present with correct shape # and FP8 parameters have been removed assert not hasattr(module, "input_scale") assert not hasattr(module, "input_zero_point") weight_scale = getattr(module, "weight_scale", None) assert ( weight_scale is not None and weight_scale.shape[:-1] == module.weight.shape[:-1] and weight_scale.shape[-1] == module.weight.shape[-1] / 128 ) weight_zero_point = getattr(module, "weight_zero_point", None) assert ( weight_zero_point is not None and weight_zero_point.shape[:-1] == module.weight.shape[:-1] and weight_zero_point.shape[-1] == module.weight.shape[-1] / 128 ) elif is_match( name, module, qconfig1.config_groups["group_0"].targets, qconfig1.ignore ): # assert FP8 scheme parameters are present with correct shape input_scale = getattr(module, "input_scale", None) assert input_scale is not None and input_scale.shape == torch.Size([1]) input_zero_point = getattr(module, "input_zero_point", None) assert ( input_zero_point is not None and input_zero_point.shape == torch.Size([1]) ) weight_scale = getattr(module, "weight_scale", None) assert weight_scale is not None and weight_scale.shape == torch.Size([1]) weight_zero_point = getattr(module, "weight_zero_point", None) assert ( weight_zero_point is not None and weight_zero_point.shape == torch.Size([1]) ) def test_apply_kv_cache(): model = AutoModelForCausalLM.from_pretrained("nm-testing/llama2.c-stories15M") args = QuantizationArgs( num_bits=8, type="float", strategy="tensor", scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=torch.float, ) config = QuantizationConfig(config_groups={}, kv_cache_scheme=args) apply_quantization_config(model, config) for layer in model.model.layers: assert getattr(layer.self_attn, "quantization_scheme").input_activations == args assert hasattr(layer.self_attn, "k_scale") assert hasattr(layer.self_attn, "v_scale") def test_apply_attention(): model = AutoModelForCausalLM.from_pretrained("nm-testing/llama2.c-stories15M") scheme = QuantizationScheme( targets=["LlamaAttention"], input_activations=QuantizationArgs( num_bits=8, type="float", strategy="tensor", scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=torch.float, ), ) config = QuantizationConfig(config_groups={"attention": scheme}) apply_quantization_config(model, config) for layer in model.model.layers: assert getattr(layer.self_attn, "quantization_scheme") == scheme assert hasattr(layer.self_attn, "q_scale") assert hasattr(layer.self_attn, "k_scale") assert hasattr(layer.self_attn, "v_scale") vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/lifecycle/test_dynamic_lifecycle.py000066400000000000000000000073431521257237700344360ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch from compressed_tensors.quantization.lifecycle import apply_quantization_config from compressed_tensors.quantization.quant_config import QuantizationConfig from transformers import AutoModelForCausalLM def test_apply_tinyllama_dynamic_activations(): # NOTE: should not calibrate dynamic quant quant_config = get_sample_dynamic_tinyllama_quant_config() model = get_tinyllama_model() # check that model is not already quantized for module in model.modules(): _test_layer_dynamic_quantization_status(module, inputs=False, weights=False) # apply quant config to model apply_quantization_config(model, quant_config) # test linears are dynamically quantized for calibration _test_linears_dynamic_quantization_status(model, quant_config, frozen=False) # verify forward works w/ dynamic during calibration model(torch.zeros((1, 1), dtype=int), torch.zeros((1, 1), dtype=int)) _test_linears_dynamic_quantization_status(model, quant_config, frozen=True) # verify forward works w/ dynamic after freeze model(torch.zeros((1, 1), dtype=int), torch.zeros((1, 1), dtype=int)) def _test_linears_dynamic_quantization_status(model, quant_config, frozen: bool): # check for correct application of quant config num_linears = 0 for name, module in model.named_modules(): if name in quant_config.ignore: continue module_type = module.__class__.__name__ if module_type == "Linear": num_linears += 1 _test_layer_dynamic_quantization_status( module, inputs=True, weights=True, frozen=frozen ) # sanity check correct number of layers targeted assert num_linears == 154 # 155 Linear layers - 1 that gets ignored def _test_layer_dynamic_quantization_status( module, inputs: bool, weights: bool, frozen: bool = False ): # check if quantization is applied at all (true if inputs or weights targeted) quantized = inputs or weights assert hasattr(module, "quantization_scheme") == quantized assert hasattr(module, "quantization_status") == quantized # check inputs always have an observer if quantized but never scale/zp assert not hasattr(module, "input_scale") assert not hasattr(module, "input_zero_point") assert not hasattr(module, "input_observer") # check weights always have scale/zp and observer only if not frozen assert hasattr(module, "weight_scale") == weights assert hasattr(module, "weight_zero_point") == weights def get_tinyllama_model(): return AutoModelForCausalLM.from_pretrained( "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T", torch_dtype=torch.bfloat16, ) def get_sample_dynamic_tinyllama_quant_config(): config_dict = { "quant_method": "compressed-tensors", "format": "fakequant", "quantization_status": "calibration", "global_compression_ratio": None, "config_groups": { "group_1": { "weights": { "num_bits": 8, "type": "int", "symmetric": True, "strategy": "tensor", "dynamic": False, }, "input_activations": { "num_bits": 8, "type": "int", "symmetric": True, "strategy": "tensor", "dynamic": True, }, "targets": ["Linear"], }, }, "ignore": ["LlamaRotaryEmbedding", "model.layers.1.mlp.down_proj"], } return QuantizationConfig.model_validate(config_dict) vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/lifecycle/test_enabled.py000066400000000000000000000025031521257237700323560ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from copy import deepcopy import torch from compressed_tensors.quantization import ( QuantizationConfig, apply_quantization_config, disable_quantization, enable_quantization, ) from torch.nn import Linear def test_quantization_enabled_disabled(): inp = torch.randn(16) model = Linear(16, 16) quantized_model = deepcopy(model) apply_quantization_config( model=quantized_model, config=QuantizationConfig( config_groups=dict(W8A8=["Linear"]), quantization_status="calibration", ), ) # run one calibration pass quantized_model(inp) model_output = model(inp) quantized_model_output = quantized_model(inp) # quantized and non quantized outputs should be different assert not torch.all(model_output == quantized_model_output) # disable quantization quantized_model.apply(disable_quantization) # check that quantized model now matches model output assert torch.all(model_output == quantized_model(inp)) # re-enable quantization quantized_model.apply(enable_quantization) # check that quantized model matches original quantized output assert torch.all(quantized_model_output == quantized_model(inp)) vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/lifecycle/test_forward.py000066400000000000000000000517061521257237700324410ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math import pytest import torch from compressed_tensors.quantization.lifecycle.forward import ( _process_quantization, fake_quantize, forward_quantize, set_forward_quantized, ) from compressed_tensors.quantization.lifecycle.forward_helpers import ( _dequantize, _quantize, _quantize_dequantize, ) from compressed_tensors.quantization.lifecycle.initialize import ( initialize_module_for_quantization, ) from compressed_tensors.quantization.quant_args import ( QuantizationArgs, QuantizationStrategy, ) from compressed_tensors.quantization.quant_config import QuantizationStatus from compressed_tensors.quantization.utils.helpers import calculate_range from torch.nn import Embedding, Linear def make_dummy_g_idx(columns: int, group_size: int) -> torch.Tensor: perm = torch.randperm(columns) return torch.tensor([index // group_size for index in range(columns)])[perm] def test_set_forward_quantized(): layer = Linear(4, 4) func_forward = layer.forward.__func__ # check that the forward call is overwritten set_forward_quantized(layer) assert not func_forward == layer.forward.__func__ def test_set_forward_quantized_embedding(): """Test that set_forward_quantized works with Embedding modules""" embedding = Embedding(num_embeddings=10, embedding_dim=4) func_forward = embedding.forward.__func__ # check that the forward call is overwritten set_forward_quantized(embedding) assert not func_forward == embedding.forward.__func__ def test_set_forward_quantized_embedding_no_quantization(): """ Test forward pass of Embedding when quantization is disabled or scheme is not set """ embedding = Embedding(num_embeddings=10, embedding_dim=4) set_forward_quantized(embedding) input_indices = torch.tensor([0, 1, 2, 3]) expected_output = torch.nn.functional.embedding(input_indices, embedding.weight) # Without quantization scheme, should behave like normal embedding output = embedding(input_indices) assert torch.allclose(output, expected_output) def test_set_forward_quantized_embedding_with_weight_quantization( mock_per_tensor_calibration, create_quantization_scheme ): """Test forward pass with weight quantization on Embedding module""" num_bits = 8 embedding = Embedding(num_embeddings=10, embedding_dim=4) embedding.weight.data *= 10 quantization_scheme = create_quantization_scheme( targets=["*"], weights=QuantizationArgs(num_bits=num_bits, symmetric=True), ) # initialize_module_for_quantization calls set_forward_quantized initialize_module_for_quantization(embedding, quantization_scheme) embedding.quantization_status = QuantizationStatus.CALIBRATION # Calibrate weights mock_per_tensor_calibration(embedding, "weight", value=embedding.weight.data) # Forward pass should quantize weights input_indices = torch.tensor([0, 1, 2, 3]) output = embedding(input_indices) assert output.shape == (4, 4) # Output should be different from unquantized forward unquantized_output = torch.nn.functional.embedding(input_indices, embedding.weight) assert not torch.allclose(output, unquantized_output, atol=1e-3) def test_set_forward_quantized_no_quantization(): """Test forward pass when quantization is disabled or scheme is not set""" layer = Linear(4, 4) set_forward_quantized(layer) input_tensor = torch.randn(2, 4) expected_output = torch.nn.functional.linear(input_tensor, layer.weight, layer.bias) # Without quantization scheme, should behave like normal linear output = layer(input_tensor) assert torch.allclose(output, expected_output) def test_set_forward_quantized_disabled(): """Test forward pass when quantization_enabled is False""" layer = Linear(4, 4) set_forward_quantized(layer) # Set up quantization but disable it layer.quantization_enabled = False layer.quantization_scheme = torch.nn.Module() # dummy scheme layer.quantization_status = QuantizationStatus.INITIALIZED input_tensor = torch.randn(2, 4) expected_output = torch.nn.functional.linear(input_tensor, layer.weight, layer.bias) # With quantization disabled, should behave like normal linear output = layer(input_tensor) assert torch.allclose(output, expected_output) @pytest.mark.parametrize( "quantization_status", [ QuantizationStatus.INITIALIZED, QuantizationStatus.CALIBRATION, QuantizationStatus.FROZEN, ], ) def test_set_forward_quantized_with_input_activations( mock_per_tensor_calibration, create_quantization_scheme, quantization_status ): """Test forward pass with input activation quantization""" num_bits = 8 layer = Linear(4, 4) layer.weight.data *= 10 quantization_scheme = create_quantization_scheme( targets=["*"], input_activations=QuantizationArgs(num_bits=num_bits, symmetric=True), ) # initialize_module_for_quantization calls set_forward_quantized initialize_module_for_quantization(layer, quantization_scheme) layer.quantization_status = quantization_status # Calibrate input activations input_tensor = torch.randn(2, 4) mock_per_tensor_calibration(layer, "input", value=input_tensor) # Forward pass should quantize inputs output = layer(input_tensor) assert output.shape == (2, 4) # Output should be different from unquantized forward unquantized_output = torch.nn.functional.linear( input_tensor, layer.weight, layer.bias ) assert not torch.allclose(output, unquantized_output, atol=1e-3) @pytest.mark.parametrize( "quantization_status", [ QuantizationStatus.INITIALIZED, QuantizationStatus.CALIBRATION, ], ) def test_set_forward_quantized_with_weight_quantization( mock_per_tensor_calibration, create_quantization_scheme, quantization_status ): """Test forward pass with weight quantization (non-FROZEN status)""" num_bits = 8 layer = Linear(4, 4) layer.weight.data *= 10 quantization_scheme = create_quantization_scheme( targets=["*"], weights=QuantizationArgs(num_bits=num_bits, symmetric=True), ) # initialize_module_for_quantization calls set_forward_quantized initialize_module_for_quantization(layer, quantization_scheme) layer.quantization_status = quantization_status # Calibrate weights mock_per_tensor_calibration(layer, "weight", value=layer.weight.data) # Forward pass should quantize weights input_tensor = torch.randn(2, 4) output = layer(input_tensor) assert output.shape == (2, 4) def test_set_forward_quantized_compressed_status( mock_per_tensor_calibration, create_quantization_scheme ): """Test that weight quantization is skipped when status is FROZEN""" num_bits = 8 layer = Linear(4, 4) layer.weight.data *= 10 quantization_scheme = create_quantization_scheme( targets=["*"], weights=QuantizationArgs(num_bits=num_bits, symmetric=True), ) # initialize_module_for_quantization calls set_forward_quantized initialize_module_for_quantization(layer, quantization_scheme) layer.quantization_status = QuantizationStatus.COMPRESSED # Calibrate weights mock_per_tensor_calibration(layer, "weight", value=layer.weight.data) # Forward pass should NOT quantize weights due to FROZEN status input_tensor = torch.randn(2, 4) output = layer(input_tensor) expected_output = torch.nn.functional.linear(input_tensor, layer.weight, layer.bias) assert torch.allclose(output, expected_output) def test_set_forward_quantized_with_output_activations( mock_per_tensor_calibration, create_quantization_scheme ): """Test forward pass with output activation quantization""" num_bits = 8 layer = Linear(4, 4) layer.weight.data *= 10 quantization_scheme = create_quantization_scheme( targets=["*"], output_activations=QuantizationArgs(num_bits=num_bits, symmetric=True), ) # initialize_module_for_quantization calls set_forward_quantized initialize_module_for_quantization(layer, quantization_scheme) layer.quantization_status = QuantizationStatus.CALIBRATION # Need to calibrate output activations input_tensor = torch.randn(2, 4) output_sample = torch.nn.functional.linear(input_tensor, layer.weight, layer.bias) mock_per_tensor_calibration(layer, "output", value=output_sample) # Forward pass should quantize outputs output = layer(input_tensor) assert output.shape == (2, 4) def test_set_forward_quantized_full_quantization( mock_per_tensor_calibration, create_quantization_scheme ): """Test forward pass with input, weight, and output quantization enabled""" num_bits = 8 layer = Linear(4, 4) layer.weight.data *= 10 quantization_scheme = create_quantization_scheme( targets=["*"], input_activations=QuantizationArgs(num_bits=num_bits, symmetric=True), weights=QuantizationArgs(num_bits=num_bits, symmetric=True), output_activations=QuantizationArgs(num_bits=num_bits, symmetric=True), ) # initialize_module_for_quantization calls set_forward_quantized initialize_module_for_quantization(layer, quantization_scheme) layer.quantization_status = QuantizationStatus.CALIBRATION # Calibrate all components input_tensor = torch.randn(2, 4) mock_per_tensor_calibration(layer, "weight", value=layer.weight.data) mock_per_tensor_calibration(layer, "input", value=input_tensor) output_sample = torch.nn.functional.linear(input_tensor, layer.weight, layer.bias) mock_per_tensor_calibration(layer, "output", value=output_sample) # Forward pass should quantize all components output = layer(input_tensor) assert output.shape == (2, 4) # Should be significantly different from unquantized unquantized_output = torch.nn.functional.linear( input_tensor, layer.weight, layer.bias ) assert not torch.allclose(output, unquantized_output, atol=1e-2) @pytest.mark.parametrize("quantization_status", ["initialized", "calibration"]) def test_forward_quantize( mock_per_tensor_calibration, create_quantization_scheme, quantization_status ): num_bits = 8 quantization_scheme = create_quantization_scheme( targets=["*"], weights=QuantizationArgs(num_bits=num_bits, symmetric=True), input_activations=QuantizationArgs(num_bits=num_bits, symmetric=True), ) quantization_args = QuantizationArgs(num_bits=num_bits, symmetric=True) layer = Linear(4, 4) layer.weight.data *= 100 dummy_tensor = torch.randn(8, 4) # (num_tokens, num_features) layer.quantization_status = QuantizationStatus(quantization_status) # only calibration updates the scale and zero-point if layer.quantization_status == QuantizationStatus.INITIALIZED: # Init zp and scales initialize_module_for_quantization(layer, quantization_scheme) # mock weight calibration mock_per_tensor_calibration(layer, "weight", value=layer.weight.data) # call quant/dequant on weights out = forward_quantize(layer, layer.weight, "weight", quantization_args) assert torch.allclose(out, layer.weight.data, atol=0.2) elif layer.quantization_status == QuantizationStatus.CALIBRATION: # init zp/scales initialize_module_for_quantization(layer, quantization_scheme) # run weight and input calibration mock_per_tensor_calibration(layer, "weight", value=layer.weight.data) mock_per_tensor_calibration(layer, "input", value=dummy_tensor) # call quant/dequant on inputs out = forward_quantize(layer, dummy_tensor, "input", quantization_args) assert torch.allclose(out, dummy_tensor, atol=0.2) @pytest.mark.parametrize( "num_bits,type,strategy,group_size,scale,zero_point,g_idx,global_scale", [ ( 4, "int", QuantizationStrategy.TENSOR, None, torch.rand((1,)) * 0.01, torch.zeros((1,)), None, None, ), ( 4, "int", QuantizationStrategy.GROUP, 128, torch.rand((512, 8)) * 0.01, torch.zeros((512, 8)), None, None, ), ( 4, "int", QuantizationStrategy.GROUP, 128, torch.rand((512, 8)) * 0.01, torch.zeros((512, 8)), make_dummy_g_idx(1024, 128), None, ), ( 8, "float", QuantizationStrategy.TENSOR, None, torch.rand((1,)) * 0.01, torch.zeros((1,)), None, None, ), ( 8, "float", QuantizationStrategy.GROUP, 128, torch.rand((512, 8)) * 0.01, torch.zeros((512, 8)), None, None, ), ( 8, "float", QuantizationStrategy.GROUP, 128, torch.rand((512, 8)) * 0.01, torch.zeros((512, 8)), make_dummy_g_idx(1024, 128), None, ), ( 8, "int", QuantizationStrategy.GROUP, 128, torch.rand((512, 8)) * 0.01, torch.zeros((512, 8)), None, None, ), ( 8, "int", QuantizationStrategy.GROUP, 128, torch.rand((512, 8)) * 0.01, torch.zeros((512, 8)), make_dummy_g_idx(1024, 128), None, ), ], ) def test_fake_quantize_2d( num_bits, type, strategy, group_size, scale, zero_point, g_idx, global_scale ): args = QuantizationArgs( num_bits=num_bits, type=type, strategy=strategy, group_size=group_size ) x = torch.rand((512, 1024)) fake_quantize( x=x, scale=scale, zero_point=zero_point, args=args, g_idx=g_idx, global_scale=global_scale, ) # note that reconstruction loss is bad for uncalibrated scales def test_process_quantization_block_static(): """ Static block quantization (QuantizationStrategy.BLOCK) should split a 2D tensor into blocks, quantize each block, and reassemble without changing shape. """ rows, cols = 8, 8 bh, bw = 2, 4 x = torch.randn(rows, cols) args = QuantizationArgs( num_bits=8, type="float", strategy=QuantizationStrategy.BLOCK, symmetric=True, dynamic=False, block_structure=[bh, bw], ) num_rb = math.ceil(rows / bh) num_cb = math.ceil(cols / bw) scale = torch.rand(num_rb, num_cb) + 0.1 zp = torch.zeros_like(scale) q_min, q_max = calculate_range(args, x.device) out = _process_quantization( x=x, scale=scale, zero_point=zp, args=args, do_quantize=True, do_dequantize=False, dtype=None, global_scale=None, ) assert out.shape == x.shape # full fake-quantize roundtrip out2 = _process_quantization( x=x, scale=scale, zero_point=zp, args=args, do_quantize=True, do_dequantize=True, dtype=None, global_scale=None, ) assert out2.shape == x.shape @pytest.mark.parametrize( "rows,cols,block_height,block_width", [ (4544, 768, 128, 128), # Falcon-7B dimensions: 4544 = 64*71 (100, 200, 128, 128), # Both dimensions not divisible (256, 300, 128, 128), # Only cols not divisible (300, 256, 128, 128), # Only rows not divisible (127, 127, 128, 128), # Both dimensions smaller than block size (1, 1, 128, 128), # Minimal tensor ], ) def test_process_quantization_block_non_divisible( rows, cols, block_height, block_width ): """ Block quantization should handle tensor dimensions that are not divisible by the block size by padding internally. """ x = torch.randn(rows, cols) args = QuantizationArgs( num_bits=8, type="float", strategy=QuantizationStrategy.BLOCK, symmetric=True, dynamic=False, block_structure=[block_height, block_width], ) # Calculate number of blocks (with ceiling division for padding) num_rb = math.ceil(rows / block_height) num_cb = math.ceil(cols / block_width) scale = torch.rand(num_rb, num_cb) + 0.1 zp = torch.zeros_like(scale) # Should NOT raise ValueError anymore out = _process_quantization( x=x, scale=scale, zero_point=zp, args=args, do_quantize=True, do_dequantize=False, dtype=None, global_scale=None, ) # Output shape should match original input shape assert out.shape == x.shape, f"Expected {x.shape}, got {out.shape}" # Full fake-quantize roundtrip out2 = _process_quantization( x=x, scale=scale, zero_point=zp, args=args, do_quantize=True, do_dequantize=True, dtype=None, global_scale=None, ) assert out2.shape == x.shape, f"Expected {x.shape}, got {out2.shape}" @pytest.mark.parametrize( "rows,cols,block_height,block_width", [ (100, 200, 128, 128), # Both dimensions not divisible (256, 300, 128, 128), # Only cols not divisible (300, 256, 128, 128), # Only rows not divisible (127, 127, 128, 128), # Both dimensions smaller than block size ], ) def test_process_quantization_block_non_divisible_values( rows, cols, block_height, block_width ): """ Verify that block quantization with non-divisible dimensions produces correct values. Using uniform input (ones) with scale=1.0 should result in zero quantization loss. """ # Use uniform values - quantization with scale=1.0 should be lossless x = torch.ones(rows, cols) args = QuantizationArgs( num_bits=8, type="float", strategy=QuantizationStrategy.BLOCK, symmetric=True, dynamic=False, block_structure=[block_height, block_width], ) num_rb = math.ceil(rows / block_height) num_cb = math.ceil(cols / block_width) # Use scale=1.0 for lossless quantization of values within FP8 range scale = torch.ones(num_rb, num_cb) zp = torch.zeros_like(scale) # Full fake-quantize roundtrip should preserve values exactly out = _process_quantization( x=x, scale=scale, zero_point=zp, args=args, do_quantize=True, do_dequantize=True, dtype=None, global_scale=None, ) # Values should match input (no quantization loss for uniform values) assert out.shape == x.shape, f"Expected shape {x.shape}, got {out.shape}" assert torch.allclose( out, x, atol=1e-6 ), f"Values mismatch: expected all ones, got min={out.min()}, max={out.max()}" # Test with a different uniform value x_val = torch.full((rows, cols), 0.5) out_val = _process_quantization( x=x_val, scale=scale, zero_point=zp, args=args, do_quantize=True, do_dequantize=True, dtype=None, global_scale=None, ) assert torch.allclose( out_val, x_val, atol=1e-6 ), f"Values mismatch for 0.5: got min={out_val.min()}, max={out_val.max()}" @pytest.mark.parametrize( "num_bits,type,symmetric,global_scale", [ (8, "int", True, None), (8, "int", False, None), (4, "int", True, None), (8, "float", True, None), (8, "float", True, torch.tensor([2.0])), (8, "int", False, torch.tensor([2.0])), ], ) def test_quantize_dequantize_matches_sequential( num_bits, type, symmetric, global_scale ): """Verify that the fused _quantize_dequantize produces identical output to calling _quantize then _dequantize sequentially.""" args = QuantizationArgs( num_bits=num_bits, type=type, symmetric=symmetric, strategy=QuantizationStrategy.TENSOR, ) q_min, q_max = calculate_range(args, torch.device("cpu")) x = torch.randn(512, 1024) scale = torch.rand(1) * 0.01 + 0.001 zero_point = None if symmetric else torch.tensor([3.0]) # sequential: quantize then dequantize q = _quantize( x=x, scale=scale, zero_point=zero_point, q_min=q_min, q_max=q_max, args=args, global_scale=global_scale, ) sequential_out = _dequantize( x_q=q, scale=scale, zero_point=zero_point, global_scale=global_scale, ) # fused fused_out = _quantize_dequantize( x=x, scale=scale, zero_point=zero_point, q_min=q_min, q_max=q_max, args=args, global_scale=global_scale, ) assert torch.equal( sequential_out, fused_out ), f"Mismatch: max diff = {(sequential_out - fused_out).abs().max().item()}" vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/lifecycle/test_initialize.py000066400000000000000000000167251521257237700331400ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math import pytest import torch from compressed_tensors.offload import set_onload_device from compressed_tensors.quantization import ( FP8_E4M3_DATA, ActivationOrdering, QuantizationArgs, QuantizationScheme, QuantizationStatus, QuantizationStrategy, ) from compressed_tensors.quantization.lifecycle.initialize import ( initialize_module_for_quantization, ) from tests.testing_utils import requires_gpu from torch.nn import Linear NUM_BITS = 8 Q_PARAM_NAMES = { "input_activations": "input", "weights": "weight", "output_activations": "output", } @pytest.fixture def layer(): return Linear(4, 4) @pytest.mark.parametrize( "weights,input_activations", [ ( QuantizationArgs(num_bits=NUM_BITS, symmetric=True), None, ), ( None, QuantizationArgs(num_bits=NUM_BITS, symmetric=True), ), ( QuantizationArgs(num_bits=NUM_BITS, symmetric=True), QuantizationArgs(num_bits=NUM_BITS, symmetric=True), ), ], ) def test_initialize_module_for_quantization( create_quantization_scheme, weights, input_activations, layer ): quantization_scheme = create_quantization_scheme( targets=["*"], weights=weights, input_activations=input_activations, ) assert not hasattr(layer, "quantization_scheme") assert not hasattr(layer, "quantization_status") # add attributes, zero_points and scale initialize_module_for_quantization(layer, quantization_scheme) registered_params = {"weight", "bias"} if weights is not None: registered_params.add("weight_scale") registered_params.add("weight_zero_point") if input_activations is not None: registered_params.add("input_scale") registered_params.add("input_zero_point") for key in layer.state_dict().keys(): assert key in registered_params registered_params.remove(key) assert len(registered_params) == 0 assert hasattr(layer, "quantization_scheme") assert hasattr(layer, "quantization_status") assert layer.quantization_status == QuantizationStatus.INITIALIZED @requires_gpu @pytest.mark.parametrize( "weights,input_activations", [ ( QuantizationArgs(num_bits=NUM_BITS, symmetric=True), None, ), ( None, QuantizationArgs(num_bits=NUM_BITS, symmetric=True), ), ( QuantizationArgs(num_bits=NUM_BITS, symmetric=True), QuantizationArgs(num_bits=NUM_BITS, symmetric=True), ), ], ) def test_initialize_module_for_quantization_offloaded( create_quantization_scheme, weights, input_activations, layer ): set_onload_device(layer, "cuda:0") test_initialize_module_for_quantization( create_quantization_scheme, weights, input_activations, layer, ) @pytest.mark.parametrize( "weights,input_activations", [ ( QuantizationArgs(strategy="tensor"), QuantizationArgs(strategy="tensor"), ), ( QuantizationArgs(strategy="channel"), None, ), ( QuantizationArgs(strategy="group", group_size=2), None, ), ( QuantizationArgs(strategy="group", group_size=2, actorder="group"), None, ), ( QuantizationArgs(strategy="group", group_size=2, actorder="weight"), None, ), ( QuantizationArgs( strategy="tensor_group", group_size=16, type="float", num_bits=4, scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=FP8_E4M3_DATA.dtype, ), None, ), ( QuantizationArgs( strategy="tensor_group", group_size=16, type="float", num_bits=4, scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=FP8_E4M3_DATA.dtype, ), QuantizationArgs( strategy="tensor_group", group_size=16, type="float", num_bits=4, dynamic="local", scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=FP8_E4M3_DATA.dtype, ), ), ( QuantizationArgs(strategy="block", block_structure=[2, 4]), None, ), ], ) def test_initialize_quantization_parameters(weights, input_activations): quantization_scheme = QuantizationScheme( targets=["*"], weights=weights, input_activations=input_activations, ) layer = Linear(7, 8) initialize_module_for_quantization(layer, quantization_scheme) for q_type in ("input_activations", "weights"): args = getattr(quantization_scheme, q_type) if args is None: continue q_param_name = Q_PARAM_NAMES[q_type] if args.strategy == QuantizationStrategy.TENSOR_GROUP: if q_type == "weights": assert hasattr(layer, "weight_global_scale") assert layer.weight_global_scale.dtype == torch.float32 assert layer.weight_global_scale.numel() == 1 assert layer.weight_scale.dtype == layer.weight.dtype elif q_type == "input_activations": assert hasattr(layer, "input_global_scale") assert layer.input_global_scale.dtype == torch.float32 assert layer.input_global_scale.numel() == 1 else: assert not hasattr(layer, "weight_global_scale") assert not hasattr(layer, "input_global_scale") # scale and zero point if args.strategy == QuantizationStrategy.TENSOR: expected_shape = (1,) elif args.strategy == QuantizationStrategy.CHANNEL: # only weight expected_shape = (layer.weight.shape[0], 1) elif args.strategy in ( QuantizationStrategy.TENSOR_GROUP, QuantizationStrategy.GROUP, ): num_groups = math.ceil(layer.weight.shape[1] / args.group_size) expected_shape = (layer.weight.shape[0], max(num_groups, 1)) elif args.strategy == QuantizationStrategy.BLOCK: # For block quantization, only weights get block-level scales # Activations fall back to tensor-level since shape is unknown at init if q_type == "weights" and args.block_structure is not None: block_height, block_width = args.block_structure rows, cols = layer.weight.shape[-2], layer.weight.shape[-1] num_rows_blocks = math.ceil(rows / block_height) num_cols_blocks = math.ceil(cols / block_width) expected_shape = (num_rows_blocks, num_cols_blocks) else: # For activations or when block_structure is None expected_shape = (1,) if not args.dynamic: assert getattr(layer, f"{q_param_name}_scale").shape == expected_shape assert getattr(layer, f"{q_param_name}_zero_point").shape == expected_shape # g_idx if args.actorder == ActivationOrdering.GROUP: assert getattr(layer, f"{q_param_name}_g_idx").shape == ( layer.weight.shape[1], ) vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/lifecycle/test_lifecycle.py000066400000000000000000000100531521257237700327220ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from copy import deepcopy import torch from compressed_tensors.quantization.lifecycle.initialize import ( initialize_module_for_quantization, ) from compressed_tensors.quantization.quant_args import QuantizationArgs from compressed_tensors.quantization.quant_config import QuantizationStatus from torch.nn import Linear def test_lifecyle(mock_per_tensor_calibration, create_quantization_scheme): torch.manual_seed(42) num_bits = 8 quantization_scheme = create_quantization_scheme( input_activations=QuantizationArgs(num_bits=num_bits, symmetric=False), weights=QuantizationArgs(num_bits=num_bits, symmetric=True), targets=["*"], ) layer = Linear(4, 4, dtype=torch.bfloat16) layer.weight.data *= 100 # updated layer keys check expected_layer_keys = {"weight", "bias"} for key in layer.state_dict().keys(): expected_layer_keys.remove(key) assert len(expected_layer_keys) == 0 # over write forward pass and register zero_point and scale initialize_module_for_quantization(layer, quantization_scheme) expected_layer_keys = { "input_scale", "input_zero_point", "weight_scale", "weight_zero_point", "weight", "bias", } for key in layer.state_dict().keys(): expected_layer_keys.remove(key) assert len(expected_layer_keys) == 0 assert hasattr(layer, "quantization_scheme") assert hasattr(layer, "quantization_status") assert layer.quantization_status == QuantizationStatus.INITIALIZED assert torch.numel(layer.input_zero_point.data) == 1 assert torch.numel(layer.input_scale) == 1 assert torch.numel(layer.weight_scale) == 1 assert torch.numel(layer.weight_zero_point) == 1 random_input = torch.randn(4, 4) random_input[0][0] = 42 # skew distribution to force non-zero zp # do a calibration step mock_per_tensor_calibration(layer, "weight", value=layer.weight) mock_per_tensor_calibration(layer, "input", value=random_input) # zero-points and scale should be updated after forward pass assert torch.numel(layer.input_zero_point.data) > 0 assert torch.numel(layer.input_scale) > 0 assert torch.numel(layer.weight_scale) > 0 assert torch.numel(layer.weight_zero_point) > 0 # symmetric zero points should center at 0 assert layer.weight_zero_point.data == 0 # check high and low bound of the weights assert torch.all(layer.weight.data >= -128) and torch.all(layer.weight.data <= 127) initialized_layer_input_zero_point = deepcopy(layer.input_zero_point) initialized_layer_input_scale = deepcopy(layer.input_scale) initialized_layer_weight_scale = deepcopy(layer.weight_scale) # calibrate the layers with each iteration for _ in range(10): random_input = torch.randn(4, 4) random_input[0][0] = 42 # skew distribution to force non-zero zp mock_per_tensor_calibration(layer, "weight", value=layer.weight) mock_per_tensor_calibration(layer, "input", value=random_input) assert initialized_layer_input_zero_point != 0 assert initialized_layer_input_scale != layer.input_scale assert initialized_layer_weight_scale == layer.weight_scale # check quantization f_q(x) is applied after frozen without update input_check_for_quant = torch.randn(4, 4) out_calibration = layer(input_check_for_quant) layer_before_freeze_input_zero_point = deepcopy(layer.input_zero_point) layer_before_freeze_input_scale = deepcopy(layer.input_scale) layer_before_freeze_weight_scale = deepcopy(layer.weight_scale) for _ in range(10): layer(torch.randn(4, 4)) assert layer_before_freeze_input_zero_point == layer.input_zero_point assert layer_before_freeze_input_scale == layer.input_scale assert layer_before_freeze_weight_scale == layer.weight_scale # check that the same quantization is applied as calibration to frozen assert torch.all(out_calibration == layer(input_check_for_quant)) vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/lifecycle/test_static_lifecycle.py000066400000000000000000000321141521257237700342730ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.quantization import ( FP8_E4M3_DATA, QuantizationScheme, forward_quantize, initialize_module_for_quantization, initialize_qparams, ) from compressed_tensors.quantization.quant_args import QuantizationArgs from compressed_tensors.quantization.quant_config import QuantizationStatus from tests.mock_observer import MockMinMaxObserver @pytest.mark.parametrize( "args,exp_min_val,exp_max_val,exp_quant,exp_loss", [ ( QuantizationArgs( num_bits=4, type="int", symmetric=True, strategy="tensor", # equivalent to token ), torch.tensor([0.0]), torch.tensor([23.0]), torch.tensor( [ [0.0000, 0.0000, 3.0625, 3.0625, 3.0625, 6.1250], [6.1250, 6.1250, 9.1875, 9.1875, 9.1875, 12.2500], [12.2500, 12.2500, 15.3125, 15.3125, 15.3125, 18.3750], [18.3750, 18.3750, 21.5000, 21.5000, 21.5000, 21.5000], ], dtype=torch.bfloat16, ), 0.85, ), # token is not supported ( QuantizationArgs( num_bits=4, type="int", symmetric=True, strategy="channel", ), torch.tensor([[0], [6], [12], [18]]), torch.tensor([[5], [11], [17], [23]]), torch.tensor( [ [0.0000, 1.3359, 2.0000, 2.6719, 4.0000, 4.6875], [5.8750, 7.3438, 7.3438, 8.8125, 10.2500, 10.2500], [11.3125, 13.6250, 13.6250, 15.8750, 15.8750, 15.8750], [18.3750, 18.3750, 21.5000, 21.5000, 21.5000, 21.5000], ], dtype=torch.bfloat16, ), 0.45, ), ( QuantizationArgs( num_bits=4, type="int", symmetric=True, strategy="group", group_size=3, ), torch.tensor([[0, 3], [6, 9], [12, 15], [18, 21]]), torch.tensor([[2, 5], [8, 11], [14, 17], [20, 23]]), torch.tensor( [ [0.0000, 1.0703, 1.8750, 2.6719, 4.0000, 4.6875], [6.4375, 7.5000, 7.5000, 8.8125, 10.2500, 10.2500], [11.1875, 13.0625, 13.0625, 15.8750, 15.8750, 15.8750], [18.7500, 18.7500, 18.7500, 21.5000, 21.5000, 21.5000], ], ), 0.45, ), ( QuantizationArgs( num_bits=4, type="float", # tensor group requires FP4 symmetric=True, strategy="tensor_group", # requires float4 group_size=3, scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=FP8_E4M3_DATA.dtype, ), torch.tensor([[0, 3], [6, 9], [12, 15], [18, 21]]), torch.tensor([[2, 5], [8, 11], [14, 17], [20, 23]]), torch.tensor( [ [0.0000, 1.0234, 2.0469, 3.2812, 3.2812, 4.9375], [5.4688, 8.1875, 8.1875, 10.6875, 10.6875, 10.6875], [9.8750, 14.7500, 14.7500, 16.3750, 16.3750, 16.3750], [19.7500, 19.7500, 19.7500, 23.0000, 23.0000, 23.0000], ], ), 1.1, ), ( QuantizationArgs( num_bits=4, type="int", symmetric=True, strategy="block", block_structure=[2, 3], ), torch.tensor([[0, 3], [12, 15]]), torch.tensor([[8, 11], [20, 23]]), torch.tensor( [ [0.0000, 1.0703, 2.1406, 2.9375, 4.4062, 4.4062], [6.4375, 7.5000, 7.5000, 8.8125, 10.2500, 10.2500], [10.6875, 13.3750, 13.3750, 15.3125, 15.3125, 18.3750], [18.7500, 18.7500, 18.7500, 21.5000, 21.5000, 21.5000], ], ), 0.5, ), ], ) def test_static_weight_quantization( args, exp_min_val, exp_max_val, exp_quant, exp_loss ): """ weight = tensor([[ 0, 1, 2, 3, 4, 5], [ 6, 7, 8, 9, 10, 11], [12, 13, 14, 15, 16, 17], [18, 19, 20, 21, 22, 23]]) """ # set up weight input_size, output_size = 6, 4 linear = torch.nn.Linear(input_size, output_size, bias=False) linear.weight.data = torch.arange( input_size * output_size, dtype=torch.bfloat16 ).reshape(output_size, input_size) # initialize quantization parameters scheme = QuantizationScheme(targets=[], weights=args) initialize_module_for_quantization(linear, scheme) assert getattr(linear, "quantization_scheme") is scheme linear.weight_observer = MockMinMaxObserver("weight", args, linear) # calibrate_global_scale if hasattr(linear, "weight_global_scale"): global_scale = linear.weight_observer.get_global_scale(linear.weight) linear.weight_global_scale.data = global_scale # calibrate quantization parameters scale, zero_point = linear.weight_observer(linear.weight) linear.weight_scale.data = scale linear.weight_zero_point.data = zero_point assert torch.equal(linear.weight_observer.min_vals, exp_min_val) assert torch.equal(linear.weight_observer.max_vals, exp_max_val) # forward pass input = torch.eye(input_size, dtype=torch.bfloat16) output = linear(input) assert torch.allclose(output.T, exp_quant.to(output.dtype)) assert torch.nn.functional.mse_loss(output.T, linear.weight) <= exp_loss @pytest.mark.parametrize( "args,exp_min_val,exp_max_val,exp_quant,exp_loss", [ ( QuantizationArgs( num_bits=4, type="int", symmetric=True, strategy="tensor", ), torch.tensor([0.0]), torch.tensor([11.0]), torch.tensor( [ [ [0.0000, 1.4688, 1.4688, 2.9375, 4.4062, 4.4062], [5.8750, 7.3438, 7.3438, 8.8125, 10.2500, 10.2500], ] ] ), 0.2, ), # static token is not supported # channel is not supported # group is not supported ( QuantizationArgs( num_bits=4, type="float", # must be fp4 symmetric=True, strategy="tensor_group", dynamic="local", group_size=3, scale_dtype=FP8_E4M3_DATA.dtype, zp_dtype=FP8_E4M3_DATA.dtype, ), None, None, torch.tensor( [ [ [0.0000, 0.9844, 1.9688, 3.4062, 3.4062, 5.1250], [5.2500, 7.8750, 7.8750, 7.3438, 11.0000, 11.0000], ] ] ), 0.5, ), # block is not supported # head is not supported ], ) def test_static_activation_quantization( args, exp_min_val, exp_max_val, exp_quant, exp_loss ): """ input = tensor([[ 0, 1, 2, 3, 4, 5] [ 6, 7, 8, 9, 10, 11]]) """ # set up activation (and identity weight) batch_size, seq_len, input_size = 1, 2, 6 input = torch.arange( (batch_size * seq_len * input_size), dtype=torch.bfloat16 ).reshape((batch_size, seq_len, input_size)) linear = torch.nn.Linear(input_size, input_size, bias=False) linear.weight.data = torch.eye(input_size, dtype=torch.bfloat16) # initialize quantization parameters scheme = QuantizationScheme(targets=[], input_activations=args) initialize_module_for_quantization(linear, scheme) assert getattr(linear, "quantization_scheme") is scheme linear.input_observer = MockMinMaxObserver("input", args, linear) # calibrate quantization parameters def calibrate_input_hook(_, args): if hasattr(linear, "input_global_scale"): global_scale = linear.input_observer.get_global_scale(args[0]) linear.input_global_scale.data = global_scale if linear.quantization_scheme.input_activations.dynamic is False: scale, zero_point = linear.input_observer(args[0]) linear.input_scale.data = scale linear.input_zero_point.data = zero_point linear.register_forward_pre_hook(calibrate_input_hook) # calibration forward pass output = linear(input) # check calibration if exp_min_val is not None: assert torch.equal(linear.input_observer.min_vals, exp_min_val) if exp_max_val is not None: assert torch.equal(linear.input_observer.max_vals, exp_max_val) # check forward pass assert torch.allclose(output, exp_quant.to(output.dtype)) assert torch.nn.functional.mse_loss(output, input) <= exp_loss class MockAttention(torch.nn.Module): pass @pytest.mark.filterwarnings("ignore::UserWarning") @pytest.mark.parametrize( "args,exp_min_val,exp_max_val,exp_quant,exp_loss", [ ( QuantizationArgs( num_bits=4, type="int", symmetric=True, strategy="tensor", ), torch.tensor([0.0]), torch.tensor([23.0]), torch.tensor( [ [ [ [0.0000, 0.0000, 3.0625, 3.0625], [3.0625, 6.1250, 6.1250, 6.1250], [9.1875, 9.1875, 9.1875, 12.2500], ], [ [12.2500, 12.2500, 15.3125, 15.3125], [15.3125, 18.3750, 18.3750, 18.3750], [21.5000, 21.5000, 21.5000, 21.5000], ], ] ] ), 0.81, ), # static token is not supported # channel is not supported # group is not supported # tensor group is not supported # block is not supported ( QuantizationArgs( num_bits=4, type="int", symmetric=True, strategy="attn_head", ), torch.tensor([[[0.0]], [[12.0]]]), torch.tensor([[[11.0]], [[23.0]]]), torch.tensor( [ [ [ [0.0000, 1.4688, 1.4688, 2.9375], [4.4062, 4.4062, 5.8750, 7.3438], [7.3438, 8.8125, 10.2500, 10.2500], ], [ [12.2500, 12.2500, 15.3125, 15.3125], [15.3125, 18.3750, 18.3750, 18.3750], [21.5000, 21.5000, 21.5000, 21.5000], ], ] ] ), 0.55, ), ], ) def test_static_attention_quantization( args, exp_min_val, exp_max_val, exp_quant, exp_loss ): """ input = tensor([[[[ 0., 1., 2., 3.], [ 4., 5., 6., 7.], [ 8., 9., 10., 11.]], [[12., 13., 14., 15.], [16., 17., 18., 19.], [20., 21., 22., 23.]]]]) """ # set up attention batch_size, num_heads, seq_len, head_dim = 1, 2, 3, 4 input = torch.arange( (batch_size * num_heads * seq_len * head_dim), dtype=torch.bfloat16 ).reshape((batch_size, num_heads, seq_len, head_dim)) attention = MockAttention() # initialize quantization parameters scheme = QuantizationScheme(targets=[], input_activations=args) initialize_qparams( attention, "k", args, (num_heads, None, head_dim), observed_dtype=torch.bfloat16 ) attention.quantization_scheme = scheme attention.quantization_status = QuantizationStatus.INITIALIZED attention.k_observer = MockMinMaxObserver("k", args, attention) # calibrate quantization parameters if scheme.input_activations.dynamic is False: scale, zero_point = attention.k_observer(input) attention.k_scale.data = scale attention.k_zero_point.data = zero_point # calibration forward pass output = forward_quantize(attention, input, "k", scheme.input_activations) # check calibration if exp_min_val is not None: assert torch.equal(attention.k_observer.min_vals, exp_min_val) if exp_max_val is not None: assert torch.equal(attention.k_observer.max_vals, exp_max_val) # check forward pass assert torch.allclose(output, exp_quant.to(output.dtype)) assert torch.nn.functional.mse_loss(output, input) <= exp_loss vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/test_configs/000077500000000000000000000000001521257237700301035ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/test_configs/__init__.py000066400000000000000000000001531521257237700322130ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/test_configs/test_bit_depths.py000066400000000000000000000130041521257237700336370ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.quantization import ( QuantizationArgs, QuantizationConfig, QuantizationScheme, QuantizationStatus, apply_quantization_config, ) from compressed_tensors.quantization.lifecycle.forward import fake_quantize, quantize from torch.nn import Linear def create_config(bit_depth, quant_type, input_symmetry, weight_symmetry): weights = QuantizationArgs( num_bits=bit_depth, type=quant_type, symmetric=weight_symmetry ) if input_symmetry is not None: inputs = QuantizationArgs( num_bits=bit_depth, type=quant_type, symmetric=input_symmetry ) else: inputs = None config_groups = { "group_1": QuantizationScheme( targets=["Linear"], weights=weights, input_activations=inputs ) } config = QuantizationConfig( config_groups=config_groups, quantization_status=QuantizationStatus.CALIBRATION ) return config @torch.no_grad @pytest.mark.parametrize("bit_depth", [4, 8]) @pytest.mark.parametrize("quant_type", ["int"]) @pytest.mark.parametrize("input_symmetry", [True, False, None]) @pytest.mark.parametrize("weight_symmetry", [True, False]) def test_bit_depths( mock_per_tensor_calibration, bit_depth, quant_type, input_symmetry, weight_symmetry ): model = Linear(64, 64) quant_config = create_config(bit_depth, quant_type, input_symmetry, weight_symmetry) apply_quantization_config(model, quant_config) min = -1 * int(2**bit_depth / 2) max = int(2**bit_depth / 2) - 1 inputs = torch.randn(32, 64) model.apply( lambda module: mock_per_tensor_calibration( module, base_name="weight", value=model.weight ) ) if input_symmetry is not None: model.apply( lambda module: mock_per_tensor_calibration( module, base_name="input", value=inputs ) ) assert model.input_zero_point >= min assert model.input_zero_point <= max input_max = torch.max(inputs) input_min = torch.min(inputs) diff_from_max = abs( abs(model.input_scale * (max - model.input_zero_point)) - abs(input_max) ) diff_from_min = abs( abs(model.input_scale * abs(min - model.input_zero_point)) - abs(input_min) ) assert diff_from_max < model.input_scale or diff_from_min < model.input_scale assert model.weight_zero_point >= min assert model.weight_zero_point <= max weight_max = torch.max(model.weight) weight_min = torch.min(model.weight) diff_from_max = abs( abs(model.weight_scale * (max - model.weight_zero_point)) - abs(weight_max) ) diff_from_min = abs( abs(model.weight_scale * abs(min - model.weight_zero_point)) - abs(weight_min) ) assert diff_from_max < model.weight_scale or diff_from_min < model.weight_scale quantized_weight = fake_quantize( model.weight, model.weight_scale, model.weight_zero_point, model.quantization_scheme.weights, ) assert not torch.any(quantized_weight < min).item() assert not torch.any(quantized_weight > max).item() @torch.no_grad @pytest.mark.parametrize("bit_depth", [8]) @pytest.mark.parametrize("quant_type", ["float"]) @pytest.mark.parametrize("input_symmetry", [True, False, None]) @pytest.mark.parametrize("weight_symmetry", [True, False]) def test_fp8( mock_per_tensor_calibration, bit_depth, quant_type, input_symmetry, weight_symmetry ): model = Linear(64, 64) quant_config = create_config(bit_depth, quant_type, input_symmetry, weight_symmetry) apply_quantization_config(model, quant_config) dtype_info = torch.finfo(torch.float8_e4m3fn) min = dtype_info.min max = dtype_info.max inputs = torch.randn(32, 64) model.apply( lambda module: mock_per_tensor_calibration( module, base_name="weight", value=model.weight ) ) assert model.weight_zero_point.dtype == torch.float8_e4m3fn model.weight_zero_point.data = model.weight_zero_point.to(model.weight.dtype) if input_symmetry is not None: model.apply( lambda module: mock_per_tensor_calibration( module, base_name="input", value=inputs ) ) assert model.input_zero_point.dtype == torch.float8_e4m3fn model.input_zero_point.data = model.input_zero_point.to(model.weight.dtype) assert model.input_zero_point >= min assert model.input_zero_point <= max inputs_fake_quant = quantize( inputs, model.input_scale, model.input_zero_point, model.quantization_scheme.input_activations, ) input_max = torch.max(inputs_fake_quant) input_min = torch.min(inputs_fake_quant) diff_from_max = abs(input_max - max) diff_from_min = abs(input_min - min) assert diff_from_max.item() == 0.0 or diff_from_min.item() == 0.0 assert model.weight_zero_point >= min assert model.weight_zero_point <= max weight_fake_quant = quantize( model.weight, model.weight_scale, model.weight_zero_point, model.quantization_scheme.weights, ) weight_max = torch.max(weight_fake_quant) weight_min = torch.min(weight_fake_quant) diff_from_max = abs(weight_max - max) diff_from_min = abs(weight_min - min) assert diff_from_max.item() == 0.0 or diff_from_min.item() == 0.0 test_compression_format.py000066400000000000000000000075511521257237700353560ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/test_configs# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import pytest from compressed_tensors.config import CompressionFormat from compressed_tensors.quantization import QuantizationConfig, QuantizationScheme def test_compression_format_serializable(): """Test that CompressionFormat can be serialized to JSON""" format = CompressionFormat.int_quantized # Test direct JSON serialization json_str = json.dumps(format) assert json_str == '"int-quantized"' # Test deserialization deserialized = CompressionFormat(json.loads(json_str)) assert deserialized == format def test_compression_format_all_values(): """Test that all CompressionFormat values are serializable""" for format in CompressionFormat: # Serialize to JSON json_str = json.dumps(format) assert isinstance(json_str, str) # Deserialize from JSON deserialized = CompressionFormat(json.loads(json_str)) assert deserialized == format def test_compression_format_in_dict(): """Test that CompressionFormat can be serialized in a dict""" test_dict = { "format": CompressionFormat.pack_quantized, "other_field": "value", } # Serialize to JSON json_str = json.dumps(test_dict, default=str) parsed = json.loads(json_str) assert parsed["format"] == "pack-quantized" assert parsed["other_field"] == "value" def test_compression_format_in_scheme(): """Test that CompressionFormat serializes properly in QuantizationScheme""" scheme = QuantizationScheme( targets=["Linear"], format=CompressionFormat.int_quantized ) # Serialize to dict scheme_dict = scheme.model_dump() assert scheme_dict["format"] == "int-quantized" assert isinstance(scheme_dict["format"], str) # Serialize to JSON json_str = json.dumps(scheme_dict) parsed = json.loads(json_str) assert parsed["format"] == "int-quantized" # Deserialize from dict reloaded = QuantizationScheme.model_validate(parsed) assert reloaded.format == CompressionFormat.int_quantized def test_compression_format_in_config(): """Test that CompressionFormat serializes properly in QuantizationConfig""" config = QuantizationConfig( config_groups={"group_1": QuantizationScheme(targets=[])}, format=CompressionFormat.float_quantized.value, ) # Serialize to dict config_dict = config.to_dict() assert config_dict["format"] == "float-quantized" assert isinstance(config_dict["format"], str) # Serialize to JSON json_str = json.dumps(config_dict) parsed = json.loads(json_str) assert parsed["format"] == "float-quantized" # Deserialize from dict reloaded = QuantizationConfig.model_validate(parsed) assert reloaded.format == "float-quantized" @pytest.mark.parametrize( "format_value", [ CompressionFormat.dense, CompressionFormat.int_quantized, CompressionFormat.float_quantized, CompressionFormat.pack_quantized, CompressionFormat.naive_quantized, CompressionFormat.mixed_precision, CompressionFormat.nvfp4_pack_quantized, CompressionFormat.mxfp4_pack_quantized, ], ) def test_compression_format_round_trip(format_value): """Test round-trip serialization for each CompressionFormat value""" # Create a config with the format config = QuantizationConfig( config_groups={"group_1": QuantizationScheme(targets=["Linear"])}, format=format_value.value, ) # Serialize to dict then JSON config_dict = config.to_dict() json_str = json.dumps(config_dict) # Deserialize from JSON parsed = json.loads(json_str) reloaded = QuantizationConfig.model_validate(parsed) # Verify format is preserved assert reloaded.format == format_value.value assert config == reloaded vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/test_configs/test_strategies.py000066400000000000000000000061151521257237700336710ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.quantization import ( QuantizationArgs, QuantizationConfig, QuantizationScheme, QuantizationStatus, QuantizationStrategy, apply_quantization_config, ) from torch.nn import Linear def create_config( input_symmetry, weight_symmetry, w_strategy, i_strategy=None, group_size=None ): weights = QuantizationArgs( symmetric=weight_symmetry, strategy=w_strategy, group_size=group_size ) if input_symmetry is not None: inputs = QuantizationArgs( symmetric=input_symmetry, strategy=i_strategy, group_size=group_size ) else: inputs = None config_groups = { "group_1": QuantizationScheme( targets=["Linear"], weights=weights, input_activations=inputs ) } config = QuantizationConfig( config_groups=config_groups, quantization_status=QuantizationStatus.CALIBRATION ) return config @torch.no_grad @pytest.mark.parametrize("input_symmetry", [None]) @pytest.mark.parametrize("weight_symmetry", [True, False]) @pytest.mark.parametrize("model_shape", [(64, 128), (300, 200), (400, 400)]) def test_channelwise( mock_per_channel_calibration, input_symmetry, weight_symmetry, model_shape ): model = Linear(model_shape[0], model_shape[1]) quant_config = create_config( input_symmetry, weight_symmetry, w_strategy=QuantizationStrategy.CHANNEL ) apply_quantization_config(model, quant_config) inputs = torch.randn(32, model_shape[0]) mock_per_channel_calibration(model, base_name="weight", value=model.weight) if input_symmetry is not None: mock_per_channel_calibration(model, base_name="input", value=inputs) assert model.weight_scale.shape == (model_shape[1], 1) assert model.weight_zero_point.shape == (model_shape[1], 1) @torch.no_grad @pytest.mark.parametrize("input_symmetry", [None]) @pytest.mark.parametrize("weight_symmetry", [True, False]) @pytest.mark.parametrize("model_shape", [(128, 256), (256, 512), (512, 1024)]) @pytest.mark.parametrize("group_size", [32, 128]) def test_group( mock_per_group_calibration, input_symmetry, weight_symmetry, model_shape, group_size ): model = Linear(model_shape[0], model_shape[1]) quant_config = create_config( input_symmetry, weight_symmetry, w_strategy=QuantizationStrategy.GROUP, group_size=group_size, ) apply_quantization_config(model, quant_config) inputs = torch.randn(128, model_shape[0]) mock_per_group_calibration( model, base_name="weight", value=model.weight, group_size=group_size ) if input_symmetry is not None: mock_per_group_calibration( model, base_name="input", value=inputs, group_size=group_size ) assert model.weight_scale.shape == ( model_shape[1], int(model_shape[0] / group_size), ) assert model.weight_zero_point.shape == ( model_shape[1], int(model_shape[0] / group_size), ) vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/test_quant_args.py000066400000000000000000000126621521257237700312000ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest from compressed_tensors.quantization import ( ActivationOrdering, QuantizationArgs, QuantizationStrategy, QuantizationType, ) from pydantic import ValidationError def test_defaults(): default = QuantizationArgs() assert default.num_bits == 8 assert default.type == QuantizationType.INT assert default.symmetric assert default.strategy == QuantizationStrategy.TENSOR assert default.group_size is None assert default.block_structure is None def test_group(): kwargs = {"strategy": "group", "group_size": 128} group = QuantizationArgs(**kwargs) assert group.strategy == QuantizationStrategy.GROUP assert group.group_size == kwargs["group_size"] with pytest.raises(ValueError): QuantizationArgs(strategy=QuantizationStrategy.GROUP, group_size=-1) args = QuantizationArgs(group_size=128, strategy="group") assert args.group_size == 128 assert args.strategy == "group" with pytest.raises(ValueError): QuantizationArgs(strategy=QuantizationStrategy.GROUP) with pytest.raises(ValueError): QuantizationArgs(strategy="tensor", group_size=128) def test_block(): kwargs = {"strategy": "block", "block_structure": "2x4"} block = QuantizationArgs(**kwargs) assert block.strategy == QuantizationStrategy.BLOCK assert block.block_structure == [2, 4] assert block.block_structure != kwargs["block_structure"] # "2x4" != [2, 4] def test_infer_strategy(): args = QuantizationArgs(group_size=128) assert args.strategy == QuantizationStrategy.GROUP args = QuantizationArgs(group_size=-1) assert args.strategy == QuantizationStrategy.CHANNEL def test_enums(): assert QuantizationArgs( type=QuantizationType.INT, strategy=QuantizationStrategy.GROUP, actorder=ActivationOrdering.WEIGHT, group_size=1, ) == QuantizationArgs(type="InT", strategy="GROUP", actorder="weight", group_size=1) def test_actorder(): # test group inference with actorder args = QuantizationArgs(group_size=128, actorder=ActivationOrdering.GROUP) assert args.strategy == QuantizationStrategy.GROUP args = QuantizationArgs(group_size=128, actorder=ActivationOrdering.DYNAMIC) assert args.strategy == QuantizationStrategy.GROUP # test invalid pairings with pytest.raises(ValueError): QuantizationArgs(group_size=None, actorder="group") with pytest.raises(ValueError): QuantizationArgs(group_size=-1, actorder="group") with pytest.raises(ValueError): QuantizationArgs(strategy="tensor", actorder="group") # test boolean and none defaulting assert ( QuantizationArgs(group_size=1, actorder=True).actorder == ActivationOrdering.GROUP ) assert QuantizationArgs(group_size=1, actorder=False).actorder is None assert QuantizationArgs(group_size=1, actorder=None).actorder is None def test_actorder_aliases(): assert ( ActivationOrdering.GROUP == ActivationOrdering.DYNAMIC == ActivationOrdering.GROUP ) assert ( ActivationOrdering.WEIGHT == ActivationOrdering.STATIC == ActivationOrdering.WEIGHT ) assert ActivationOrdering.GROUP == "dynamic" == ActivationOrdering.GROUP assert ActivationOrdering.DYNAMIC == "dynamic" == ActivationOrdering.DYNAMIC assert ActivationOrdering.GROUP == "group" == ActivationOrdering.GROUP assert ActivationOrdering.DYNAMIC == "group" == ActivationOrdering.DYNAMIC assert ActivationOrdering.WEIGHT == "static" == ActivationOrdering.WEIGHT assert ActivationOrdering.STATIC == "static" == ActivationOrdering.STATIC assert ActivationOrdering.WEIGHT == "weight" == ActivationOrdering.WEIGHT assert ActivationOrdering.STATIC == "weight" == ActivationOrdering.STATIC assert ActivationOrdering.WEIGHT != "dynamic" != ActivationOrdering.WEIGHT assert ActivationOrdering.STATIC != "dynamic" != ActivationOrdering.STATIC assert ActivationOrdering.WEIGHT != "group" != ActivationOrdering.WEIGHT assert ActivationOrdering.STATIC != "group" != ActivationOrdering.STATIC assert ActivationOrdering.GROUP != "static" != ActivationOrdering.GROUP assert ActivationOrdering.DYNAMIC != "static" != ActivationOrdering.DYNAMIC assert ActivationOrdering.GROUP != "weight" != ActivationOrdering.GROUP assert ActivationOrdering.DYNAMIC != "weight" != ActivationOrdering.DYNAMIC def test_invalid(): with pytest.raises(ValidationError): QuantizationArgs(type="invalid") with pytest.raises(ValidationError): QuantizationArgs(strategy="invalid") with pytest.raises(ValidationError): QuantizationArgs(strategy=QuantizationStrategy.GROUP) def test_serialize_args(): """Test serialization of QuantizationArgs""" args = QuantizationArgs( num_bits=4, type=QuantizationType.INT, symmetric=True, group_size=128, actorder=ActivationOrdering.GROUP, ) # Serialize to dict args_dict = args.model_dump() assert args_dict["num_bits"] == 4 assert args_dict["type"] == "int" assert args_dict["symmetric"] is True assert args_dict["group_size"] == 128 assert args_dict["strategy"] == "group" assert args_dict["actorder"] == "group" # Deserialize from dict reloaded = QuantizationArgs.model_validate(args_dict) assert reloaded == args vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/test_quant_config.py000066400000000000000000000060661521257237700315120ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest from compressed_tensors.quantization import ( DEFAULT_QUANTIZATION_FORMAT, DEFAULT_QUANTIZATION_METHOD, QuantizationConfig, QuantizationScheme, QuantizationStatus, ) from pydantic import ValidationError def test_basic_config(): config_groups = {"group_1": QuantizationScheme(targets=[])} config = QuantizationConfig(config_groups=config_groups) assert config.config_groups == config_groups assert config.quant_method == DEFAULT_QUANTIZATION_METHOD assert config.format == DEFAULT_QUANTIZATION_FORMAT assert config.quantization_status == QuantizationStatus.INITIALIZED assert config.global_compression_ratio is None assert isinstance(config.ignore, list) and len(config.ignore) == 0 def test_full_config(): config_groups = { "group_1": QuantizationScheme(targets=[]), "group_2": QuantizationScheme(targets=[]), } global_compression_ratio = 3.5 ignore = ["model.layers.0"] quantization_status = "compressed" config = QuantizationConfig( config_groups=config_groups, global_compression_ratio=global_compression_ratio, ignore=ignore, quantization_status=quantization_status, ) assert config.config_groups == config_groups assert config.global_compression_ratio == global_compression_ratio assert config.ignore == ignore assert config.quantization_status == QuantizationStatus.COMPRESSED def test_need_config_groups(): with pytest.raises(ValidationError): _ = QuantizationScheme() @pytest.mark.parametrize( "scheme_name", ["W8A8", "W8A16", "W4A16", "FP8"], ) def test_load_scheme_from_preset(scheme_name: str): targets = ["Linear"] config = QuantizationConfig(config_groups={scheme_name: targets}) assert scheme_name in config.config_groups assert isinstance(config.config_groups[scheme_name], QuantizationScheme) assert config.config_groups[scheme_name].targets == targets def test_to_dict(): """Test serialization of QuantizationConfig including format""" from compressed_tensors.quantization import QuantizationArgs config_groups = { "group_1": QuantizationScheme( targets=["Linear"], weights=QuantizationArgs(num_bits=4, symmetric=True, group_size=128), ), "group_2": QuantizationScheme( targets=["Conv2d"], weights=QuantizationArgs(num_bits=8), ), } config = QuantizationConfig( config_groups=config_groups, global_compression_ratio=3.5, ignore=["model.layers.0"], quantization_status="compressed", format="int-quantized", ) # Serialize to dict config_dict = config.to_dict() assert "config_groups" in config_dict assert config_dict["format"] == "int-quantized" assert config_dict["quantization_status"] == "compressed" # Deserialize from dict reloaded = QuantizationConfig.model_validate(config_dict) assert config == reloaded vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/test_quant_metadata.py000066400000000000000000000030551521257237700320200ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.offload import offload_module from compressed_tensors.quantization import ( QuantizationMetadata, initialize_module_for_quantization, preset_name_to_scheme, ) @pytest.mark.parametrize("offloaded", (True, False)) def test_clear(offloaded): module = torch.nn.Linear(16, 16) scheme = preset_name_to_scheme("NVFP4", ["Linear"]) base_forward = module.forward # offload module if offloaded: offload_module(module, "cpu", "cpu") offloaded_forward = module.forward assert module._original_forward_func is base_forward.__func__ # add quantized forward (inside offloaded forward) initialize_module_for_quantization(module, scheme) qparams = ["weight_scale", "weight_global_scale", "input_global_scale"] for name in qparams: assert hasattr(module, name) if offloaded: assert module.forward is offloaded_forward assert module._original_forward_func.__wrapped__ is base_forward.__func__ else: assert module.forward.__wrapped__ is base_forward.__func__ # remove quantized forward QuantizationMetadata.clear_quantization(module) for name in qparams: assert not hasattr(module, name) if offloaded: assert module.forward is offloaded_forward assert module._original_forward_func is base_forward.__func__ else: assert module.forward.__func__ is base_forward.__func__ vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/test_quant_scheme.py000066400000000000000000000052721521257237700315070ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme from pydantic import ValidationError def test_basic_scheme(): targets = ["model.layer.0", "model.layer.3"] weights = QuantizationArgs() scheme = QuantizationScheme(targets=targets, weights=weights) assert scheme.targets == targets assert scheme.weights == weights assert scheme.input_activations is None assert scheme.output_activations is None assert scheme.format is None def test_full_scheme(): targets = ["Linear"] weights = QuantizationArgs() input_activations = QuantizationArgs(num_bits=8) output_activations = QuantizationArgs(num_bits=8, type="float", symmetric=False) scheme = QuantizationScheme( targets=targets, weights=weights, input_activations=input_activations, output_activations=output_activations, format="float-quantized", ) assert scheme.targets == targets assert scheme.weights == weights assert scheme.input_activations == input_activations assert scheme.output_activations == output_activations assert scheme.format == "float-quantized" def test_needs_targets(): with pytest.raises(ValidationError): _ = QuantizationScheme() def test_defaults(): targets = ["Linear"] output = QuantizationScheme(targets=targets) assert output.weights is None assert output.input_activations is None assert output.output_activations is None assert output.format is None def test_serialize_scheme(): """Test serialization of QuantizationScheme including format""" from compressed_tensors.config import CompressionFormat targets = ["Linear"] weights = QuantizationArgs(num_bits=4, symmetric=True, group_size=128) input_activations = QuantizationArgs(num_bits=8, dynamic=True) output_activations = QuantizationArgs(num_bits=8, type="float", symmetric=False) scheme = QuantizationScheme( targets=targets, weights=weights, input_activations=input_activations, output_activations=output_activations, format=CompressionFormat.pack_quantized, ) # Serialize to dict scheme_dict = scheme.model_dump() assert scheme_dict["targets"] == targets assert scheme_dict["format"] == "pack-quantized" assert "weights" in scheme_dict assert scheme_dict["weights"]["num_bits"] == 4 assert "input_activations" in scheme_dict assert "output_activations" in scheme_dict # Deserialize from dict reloaded = QuantizationScheme.model_validate(scheme_dict) assert reloaded == scheme vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/test_utils/000077500000000000000000000000001521257237700276135ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/test_utils/test_helpers.py000066400000000000000000000124451521257237700326740ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.quantization import ( FP4_E2M1_DATA, FP8_E4M3_DATA, QuantizationArgs, QuantizationStrategy, ) from compressed_tensors.quantization.utils import ( calculate_block_padding, calculate_qparams, compute_dynamic_scales_and_zp, generate_gparam, maybe_pad_tensor_for_block_quant, ) @pytest.mark.parametrize( "keepdims,strategy,exp_shape", [ ( False, "tensor", torch.Size( [ 1, ] ), ), (True, "channel", torch.Size([1, 1])), (True, "group", torch.Size([1, 1])), ( False, "block", torch.Size( [ 1, ] ), ), ], ) def test_calculate_qparams(keepdims, strategy, exp_shape): value = torch.empty(5, 6) min_val = torch.amin(value, dim=tuple(), keepdims=keepdims) max_val = torch.amax(value, dim=tuple(), keepdims=keepdims) if strategy == QuantizationStrategy.GROUP: args = QuantizationArgs(strategy=strategy, group_size=2) elif strategy == QuantizationStrategy.BLOCK: args = QuantizationArgs(strategy=strategy, block_structure=[1, 3]) else: args = QuantizationArgs( strategy=strategy, group_size=(2 if strategy == "group" else None), block_structure=([1, 3] if strategy == "block" else None), ) scale, zp = calculate_qparams(min_val, max_val, args) assert scale.shape == exp_shape assert zp.shape == exp_shape def test_fused_global_scales(): layer = torch.nn.Linear(7, 8) max_tensor_value = torch.abs(layer.weight.data).max() # use defaults min_val, max_val = torch.aminmax(layer.weight) global_scale = generate_gparam(min_val.data, max_val.data) # max value should be = (448 * 6) / global_scale assert max_tensor_value.item() == pytest.approx( FP4_E2M1_DATA.max * FP8_E4M3_DATA.max / global_scale, abs=0.001 ) @pytest.mark.parametrize( "shape,group_size,exp_shape", [ # Only batch size =1 is supported for dynamic GROUP quantization ((1, 4, 8), 4, torch.Size([1, 4, 2])), ], ) def test_compute_dynamic_scales_and_zp_group(shape, group_size, exp_shape): """ Dynamic group quantization should reduce activations in groups, producing scales and zero points of shape [batch, num_groups]. """ value = torch.randn(*shape) args = QuantizationArgs( strategy=QuantizationStrategy.GROUP, group_size=group_size, dynamic=True, ) scale, zp = compute_dynamic_scales_and_zp(value, args, module=torch.nn.Module()) assert scale.shape == exp_shape assert zp.shape == exp_shape # Tests for block quantization padding utilities @pytest.mark.parametrize( "shape,block_structure,expected_padding", [ # DeepSeek-V2-Lite intermediate_size (10944 % 128 = 64) ((10944, 2048), (128, 128), (64, 0)), # Both dimensions non-divisible ((100, 200), (128, 128), (28, 56)), # Only rows non-divisible ((300, 256), (128, 128), (84, 0)), # Only cols non-divisible ((256, 300), (128, 128), (0, 84)), # Divisible dimensions (no padding needed) ((256, 384), (128, 128), (0, 0)), # Smaller than block size ((100, 100), (128, 128), (28, 28)), ], ) def test_calculate_block_padding(shape, block_structure, expected_padding): """Test that calculate_block_padding computes correct padding amounts.""" pad_rows, pad_cols = calculate_block_padding(shape, block_structure) assert (pad_rows, pad_cols) == expected_padding, ( f"For shape {shape} with block {block_structure}, " f"expected padding {expected_padding}, got ({pad_rows}, {pad_cols})" ) @pytest.mark.parametrize( "rows,cols,block_height,block_width", [ (10944, 2048, 128, 128), # DeepSeek-V2-Lite (100, 200, 128, 128), # Both non-divisible (256, 256, 128, 128), # Divisible (no padding) (50, 50, 128, 128), # Smaller than block ], ) def test_maybe_pad_tensor_for_block_quant(rows, cols, block_height, block_width): """Test that maybe_pad_tensor_for_block_quant correctly pads tensors.""" tensor = torch.randn(rows, cols) block_structure = (block_height, block_width) padded = maybe_pad_tensor_for_block_quant(tensor, block_structure) # Check padded dimensions are divisible by block size assert ( padded.shape[-2] % block_height == 0 ), f"Padded rows {padded.shape[-2]} should be divisible by {block_height}" assert ( padded.shape[-1] % block_width == 0 ), f"Padded cols {padded.shape[-1]} should be divisible by {block_width}" # Check that original values are preserved assert torch.equal( padded[:rows, :cols], tensor ), "Original values should be preserved in padded tensor" # Check that padding is zeros if padded.shape[-2] > rows: assert torch.all(padded[rows:, :] == 0), "Row padding should be zeros" if padded.shape[-1] > cols: assert torch.all(padded[:, cols:] == 0), "Column padding should be zeros" vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/test_utils/test_mxfp4_utils.py000066400000000000000000000060021521257237700335000ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.quantization import round_to_quantized_type_dtype from compressed_tensors.quantization.utils import ( generate_mx_scales, maybe_convert_from_mx_exp, round_to_power_2, ) def test_round_power_2_noise(): powers = torch.Tensor( [ [2**-10, 2**-9, 2**-8, 2**-7, 2**-6], [2**-5, 2**-4, 2**-3, 2**-2, 2**-1], [2**0, 2**1, 2**-10, 2**-9, 2**-8], [2**-7, 2**-6, 2**-5, 2**-4, 2**-3], [2**-2, 2**-1, 2**0, 2**1, 2**-10], ] ).to(torch.bfloat16) noise = torch.rand_like(powers) * 0.2 powers_noisy = powers * (1 + noise) rounded = round_to_power_2(powers_noisy) assert torch.equal(rounded, powers) def test_round_power_2(): x = torch.Tensor( ( [5.687891, -8.291567, -1.540329, -0.315635, 0.965272], [-6.944130, 0.073246, -0.451778, 8.571118, -9.856593], [-0.040571, -0.708509, 2.485657, -4.003352, -0.995600], [0.224199, 5.032586, -1.309816, -0.621958, 7.290238], [-9.848001, -0.290731, 1.501562, 0.379829, -5.312081], ) ).to(torch.bfloat16) x_rounded = torch.Tensor( ( [4.000000, -8.000000, -1.000000, -0.250000, 1.000000], [-4.000000, 0.062500, -0.500000, 8.000000, -8.000000], [-0.0312, -0.500000, 2.000000, -4.000000, -1.000000], [0.250000, 4.000000, -1.000000, -0.500000, 8.000000], [-8.000000, -0.250000, 1.000000, 0.250000, -4.000000], ) ).to(torch.bfloat16) rounded = round_to_power_2(x) assert torch.equal(rounded, x_rounded) @pytest.mark.parametrize( "dtype", [torch.bfloat16, torch.float16, torch.float32, torch.float64] ) def test_mxfp4_scales_e2e(dtype): from compressed_tensors.quantization.quant_args import ( QuantizationArgs, QuantizationStrategy, QuantizationType, ) mock_weight = torch.normal(mean=0.0002, std=0.0576, size=(2880, 2880)) x = mock_weight.reshape(*mock_weight.shape[:-1], -1, 32).to(dtype) min_vals = torch.amin(x, dim=-1) max_vals = torch.amax(x, dim=-1) min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) block_max = torch.max(torch.abs(min_vals), torch.abs(max_vals)) args = QuantizationArgs( num_bits=4, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.GROUP, group_size=32, scale_dtype=torch.uint8, zp_dtype=torch.uint8, ) scales = generate_mx_scales(block_max, num_bits=4) scales = round_to_quantized_type_dtype(scales, dtype=args.scale_dtype) converted_ct = maybe_convert_from_mx_exp(args=args, scale=scales) scales_exp = torch.log2(converted_ct) block_max_exp = torch.floor(torch.log2(round_to_power_2(block_max))) - 2 assert torch.equal(scales_exp, block_max_exp) vllm-project-compressed-tensors-c18a0fa/tests/test_quantization/test_utils/test_mxfp8_utils.py000066400000000000000000000061561521257237700335160ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.quantization import round_to_quantized_type_dtype from compressed_tensors.quantization.quant_args import ( QuantizationArgs, QuantizationStrategy, QuantizationType, ) from compressed_tensors.quantization.utils import ( generate_mx_scales, maybe_convert_from_mx_exp, round_to_power_2, should_generate_mx_scales, ) def test_should_generate_mx_scales_mxfp8(): """Test that should_generate_mx_scales returns True for MXFP8 args.""" args = QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.GROUP, group_size=32, scale_dtype=torch.uint8, zp_dtype=torch.uint8, ) assert should_generate_mx_scales(args) is True def test_should_generate_mx_scales_mxfp4(): """Test that should_generate_mx_scales returns True for MXFP4 args.""" args = QuantizationArgs( num_bits=4, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.GROUP, group_size=32, scale_dtype=torch.uint8, zp_dtype=torch.uint8, ) assert should_generate_mx_scales(args) is True def test_should_generate_mx_scales_regular_fp8(): """Test that should_generate_mx_scales returns False for regular FP8.""" args = QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.TENSOR, ) assert should_generate_mx_scales(args) is False def test_should_generate_mx_scales_wrong_group_size(): """Test that should_generate_mx_scales returns False for non-32 group size.""" args = QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.GROUP, group_size=128, ) assert should_generate_mx_scales(args) is False @pytest.mark.parametrize( "dtype", [torch.bfloat16, torch.float16, torch.float32, torch.float64] ) def test_mxfp8_scales_e2e(dtype): """End-to-end test for MXFP8 scale generation and conversion.""" mock_weight = torch.normal(mean=0.0002, std=0.0576, size=(2880, 2880)) x = mock_weight.reshape(*mock_weight.shape[:-1], -1, 32).to(dtype) min_vals = torch.amin(x, dim=-1) max_vals = torch.amax(x, dim=-1) min_vals = torch.min(min_vals, torch.zeros_like(min_vals)) max_vals = torch.max(max_vals, torch.zeros_like(max_vals)) block_max = torch.max(torch.abs(min_vals), torch.abs(max_vals)) args = QuantizationArgs( num_bits=8, type=QuantizationType.FLOAT, strategy=QuantizationStrategy.GROUP, group_size=32, scale_dtype=torch.uint8, zp_dtype=torch.uint8, ) scales = generate_mx_scales(block_max, num_bits=8) scales = round_to_quantized_type_dtype(scales, dtype=args.scale_dtype) converted_ct = maybe_convert_from_mx_exp(args=args, scale=scales) scales_exp = torch.log2(converted_ct) block_max_exp = torch.floor(torch.log2(round_to_power_2(block_max))) - 8 assert torch.equal(scales_exp, block_max_exp) vllm-project-compressed-tensors-c18a0fa/tests/test_transform/000077500000000000000000000000001521257237700247015ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_transform/conftest.py000066400000000000000000000111771521257237700271070ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.transform import TransformArgs from transformers import PretrainedConfig, PreTrainedModel class TransformableModel(PreTrainedModel): config_class = PretrainedConfig def __init__(self, *sizes): super().__init__(config=PretrainedConfig()) self.fcs = torch.nn.ModuleList( [ torch.nn.Linear(sizes[index], sizes[index + 1], bias=False) for index in range(0, len(sizes) - 1) ] ) def forward(self, x): for layer in self.fcs: x = layer(x) return x class MockAttention(torch.nn.Module): def __init__( self, hidden_size: int, num_attention_heads: int, num_key_value_heads: int ): super().__init__() self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.num_key_value_groups = num_attention_heads // num_key_value_heads self.head_dim = hidden_size // num_attention_heads self.scaling = self.head_dim**-0.5 assert hidden_size >= num_attention_heads * self.head_dim self.q_proj = torch.nn.Linear( hidden_size, num_attention_heads * self.head_dim, bias=False ) self.k_proj = torch.nn.Linear( hidden_size, num_key_value_heads * self.head_dim, bias=False ) self.v_proj = torch.nn.Linear( hidden_size, num_key_value_heads * self.head_dim, bias=False ) self.o_proj = torch.nn.Linear( num_attention_heads * self.head_dim, hidden_size, bias=False ) def forward( self, hidden_states: torch.Tensor, past_key_values=None ) -> torch.Tensor: batch_size, seq_len, hidden_size = hidden_states.shape hidden_shape = (batch_size, seq_len, -1, self.head_dim) query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) if past_key_values is not None: past_key_values.update(key_states, value_states, 0, {}) key_states = self.repeat_kv(key_states, self.num_key_value_groups) value_states = self.repeat_kv(value_states, self.num_key_value_groups) attn_weights = ( torch.matmul(query_states, key_states.transpose(2, 3)) * self.scaling ) attn_weights = torch.nn.functional.softmax( attn_weights, dim=-1, dtype=torch.float32 ).to(query_states.dtype) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape((batch_size, seq_len, -1)).contiguous() return self.o_proj(attn_output) def repeat_kv(self, hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand( batch, num_key_value_heads, n_rep, slen, head_dim ) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) class MockAttentionModel(PreTrainedModel): config_class = PretrainedConfig def __init__(self, hidden_size, num_attention_heads, num_key_value_heads): super().__init__(PretrainedConfig()) self.self_attn = MockAttention( hidden_size=hidden_size, num_attention_heads=num_attention_heads, num_key_value_heads=num_key_value_heads, ) def forward(self, x): return self.self_attn(x) @pytest.fixture(scope="function") def model_apply(): model = TransformableModel(2, 4, 8, 16, 32, 64) apply = [ # weight output -> input TransformArgs(targets="fcs.0", location="weight_output"), TransformArgs(targets="fcs.1", location="input", inverse=True), # output -> weight input TransformArgs(targets="fcs.1", location="output"), TransformArgs(targets="fcs.2", location="weight_input", inverse=True), # output -> input TransformArgs(targets="fcs.2", location="output"), TransformArgs(targets="fcs.3", location="input", inverse=True), # weight output -> weight input TransformArgs(targets="fcs.3", location="weight_output"), TransformArgs(targets="fcs.4", location="weight_input", inverse=True), ] return model, apply vllm-project-compressed-tensors-c18a0fa/tests/test_transform/factory/000077500000000000000000000000001521257237700263505ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_transform/factory/test_correctness.py000066400000000000000000000242541521257237700323220ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.offload import set_onload_device from compressed_tensors.transform import ( TransformArgs, TransformConfig, TransformFactory, TransformScheme, apply_transform_config, ) from tests.test_transform.conftest import MockAttention, MockAttentionModel from tests.testing_utils import requires_gpu @pytest.mark.parametrize("type", ("hadamard", "random-hadamard", "random-matrix")) @pytest.mark.parametrize("randomize", (True, False)) @pytest.mark.parametrize("head_dim", (None, 2, 4)) @pytest.mark.parametrize("input_batch_size", (1, 5, 17)) def test_correctness_linear(type, randomize, head_dim, input_batch_size): size = (4, 8) module = torch.nn.Linear(*size, bias=False) scheme = TransformScheme(type=type, randomize=randomize, head_dim=head_dim) factory = TransformFactory.from_scheme(scheme, name="") input_tfm = factory.create_transform( module, TransformArgs(targets="Linear", location="input", inverse=True) ) w_in_tfm = factory.create_transform( module, TransformArgs(targets="Linear", location="weight_input") ) w_out_tfm = factory.create_transform( module, TransformArgs(targets="Linear", location="weight_output") ) output_tfm = factory.create_transform( module, TransformArgs(targets="Linear", location="output", inverse=True) ) input = torch.rand((input_batch_size, 5, size[0])) true_output = input @ module.weight.T input_transformed = input_tfm(input) weight_transformed = w_out_tfm(w_in_tfm(module.weight)) output = output_tfm(input_transformed @ weight_transformed.T) assert torch.allclose(true_output, output, atol=1e-5, rtol=0.0) @pytest.mark.parametrize("type", ("hadamard", "random-hadamard", "random-matrix")) @pytest.mark.parametrize("randomize", (True, False)) @pytest.mark.parametrize("embed_loc", ("weight_output", "output")) @pytest.mark.parametrize("linear_loc", ("input", "weight_input")) def test_correctness_embedding(type, randomize, embed_loc, linear_loc): model = torch.nn.Sequential( torch.nn.Embedding(2, 4), torch.nn.Linear(4, 8, bias=False), ) input = torch.randint(high=1, low=0, size=(17, 5, 2)) true_output = model(input) config = TransformConfig( config_groups={ "": TransformScheme( type=type, randomize=randomize, apply=[ TransformArgs(targets="Embedding", location=embed_loc), TransformArgs(targets="Linear", location=linear_loc, inverse=True), ], ) } ) apply_transform_config(model, config) # compare outputs output = model(input) assert torch.allclose(true_output, output, atol=1e-5, rtol=0.0) @requires_gpu @pytest.mark.parametrize("type", ("hadamard", "random-hadamard", "random-matrix")) @pytest.mark.parametrize("randomize", (True, False)) @pytest.mark.parametrize("input_batch_size", (1, 5, 17)) @pytest.mark.parametrize("offload", (True, False)) def test_correctness_model(type, randomize, input_batch_size, model_apply, offload): # load model model = model_apply[0] if offload: set_onload_device(model, torch.device("cuda")) # get output input = torch.rand((input_batch_size, 5, model.fcs[0].in_features)) if offload: input = input.to(torch.device("cuda")) true_output = model(input) # apply transforms config = TransformConfig( config_groups={ "": TransformScheme(type=type, randomize=randomize, apply=model_apply[1]) } ) apply_transform_config(model, config) # compare outputs output = model(input) assert torch.allclose(true_output, output, atol=1e-5, rtol=0.0) @pytest.mark.parametrize("type", ("hadamard", "random-hadamard", "random-matrix")) @pytest.mark.parametrize("randomize", (True, False)) @pytest.mark.parametrize("head_dim", (4, 8)) @pytest.mark.parametrize("input_batch_size", (1, 5, 17)) def test_correctness_attention_heads(type, randomize, head_dim, input_batch_size): hidden_size = 64 num_attention_heads = 8 attention = MockAttention( hidden_size=hidden_size, num_attention_heads=num_attention_heads, num_key_value_heads=head_dim, ) input = torch.rand(input_batch_size, 5, hidden_size) true_output = attention(input) config = TransformConfig( config_groups={ "R2": TransformScheme( type=type, randomize=randomize, head_dim=head_dim, apply=[ TransformArgs(targets="v_proj", location="weight_output"), TransformArgs( targets="o_proj", location="weight_input", inverse=True ), ], ) } ) apply_transform_config(attention, config) output = attention(input) assert torch.allclose(true_output, output, atol=1e-5, rtol=0.0) @pytest.mark.parametrize("type", ("hadamard", "random-hadamard", "random-matrix")) @pytest.mark.parametrize("randomize", (True, False)) @pytest.mark.parametrize("head_dim", (None, 2, 4)) @pytest.mark.parametrize("input_batch_size", (1, 5)) def test_correctness_linear_with_bias(type, randomize, head_dim, input_batch_size): """Test that WEIGHT_OUTPUT transforms correctly handle bias. For WEIGHT_OUTPUT on a Linear with bias: y' = R @ (W @ x + b) = (R @ W) @ x + (R @ b) Both weight and bias must be transformed. """ size = (4, 8) module = torch.nn.Linear(*size, bias=True) input = torch.rand((input_batch_size, 5, size[0])) # Apply WEIGHT_OUTPUT (transforms weight AND bias) then WEIGHT_INPUT inverse config = TransformConfig( config_groups={ "": TransformScheme( type=type, randomize=randomize, head_dim=head_dim, apply=[ TransformArgs(targets="0", location="weight_output"), TransformArgs(targets="1", location="weight_input", inverse=True), ], ) } ) model = torch.nn.Sequential(module, torch.nn.Linear(size[1], 16, bias=False)) true_output_full = model(input) apply_transform_config(model, config) output = model(input) assert torch.allclose(true_output_full, output, atol=1e-5, rtol=0.0) @pytest.mark.parametrize("type", ("hadamard", "random-hadamard", "random-matrix")) @pytest.mark.parametrize("randomize", (True, False)) @pytest.mark.parametrize("head_dim", (4, 8)) @pytest.mark.parametrize("input_batch_size", (1, 5)) def test_correctness_attention_heads_with_bias( type, randomize, head_dim, input_batch_size ): """Test R2 head-wise rotation with attention bias (e.g. Qwen2 v_proj). When v_proj has bias and R2 WEIGHT_OUTPUT is applied, the bias must also be rotated for the o_proj WEIGHT_INPUT inverse to correctly cancel out. """ hidden_size = 64 num_attention_heads = 8 attention = MockAttention( hidden_size=hidden_size, num_attention_heads=num_attention_heads, num_key_value_heads=head_dim, ) # Add bias to v_proj to simulate Qwen2-like architecture attention.v_proj.bias = torch.nn.Parameter( torch.randn(attention.v_proj.out_features) ) input = torch.rand(input_batch_size, 5, hidden_size) true_output = attention(input) config = TransformConfig( config_groups={ "R2": TransformScheme( type=type, randomize=randomize, head_dim=head_dim, apply=[ TransformArgs(targets="v_proj", location="weight_output"), TransformArgs( targets="o_proj", location="weight_input", inverse=True ), ], ) } ) apply_transform_config(attention, config) output = attention(input) assert torch.allclose(true_output, output, atol=1e-5, rtol=0.0) @pytest.mark.parametrize("type", ("hadamard", "random-hadamard")) @pytest.mark.parametrize("randomize", (True, False)) @pytest.mark.parametrize("head_dim", (4, 8)) @pytest.mark.parametrize("input_batch_size", (1, 5, 17)) def test_correctness_query_key_locations(type, randomize, head_dim, input_batch_size): hidden_size = 64 num_attention_heads = 8 model = MockAttentionModel( hidden_size=hidden_size, num_attention_heads=num_attention_heads, num_key_value_heads=head_dim, ) input = torch.rand(input_batch_size, 5, hidden_size) true_output = model(input) config = TransformConfig( config_groups={ "R3": TransformScheme( type=type, randomize=randomize, head_dim=head_dim, apply=[ TransformArgs(targets="self_attn", location="q_attn"), TransformArgs(targets="self_attn", location="k_cache"), ], ) } ) apply_transform_config(model, config) output = model(input) assert torch.allclose(true_output, output, atol=1e-5, rtol=0.0) @requires_gpu @pytest.mark.parametrize("cuda_default", (True, False)) def test_random_matrix_device_handling(cuda_default): """ Test that random-matrix transforms can be created on CUDA. """ seed = 0 size = (4, 8) cur_default = torch.get_default_device() if cuda_default: torch.set_default_device("cuda") module = torch.nn.Linear(*size, bias=False).cuda() scheme = TransformScheme(type="random-matrix", randomize=True) factory = TransformFactory.from_scheme(scheme, name="", seed=seed) # Create transforms - this should work despite CPU generator and CUDA module input_tfm = factory.create_transform( module, TransformArgs(targets="Linear", location="input", inverse=True) ) # Verify transforms work correctly on CUDA input = torch.rand((5, 3, size[0])).cuda() input_tfm(input) # Verify that transforms were created on CUDA assert input_tfm.weight.device.type == "cuda" torch.set_default_device(cur_default) vllm-project-compressed-tensors-c18a0fa/tests/test_transform/factory/test_memory.py000066400000000000000000000041121521257237700312670ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import Counter import pytest import torch from compressed_tensors.offload import ( disable_offloading, disable_onloading, set_onload_device, ) from compressed_tensors.transform import ( TransformArgs, TransformBase, TransformConfig, TransformScheme, apply_transform_config, ) from tests.test_transform.conftest import TransformableModel from tests.testing_utils import requires_gpu @requires_gpu @pytest.mark.parametrize("type", ("hadamard", "random-hadamard")) @pytest.mark.parametrize("randomize", (True, False)) @pytest.mark.parametrize("requires_grad", (True, False)) # @pytest.mark.parametrize("offload", (True, False)) @pytest.mark.parametrize("offload", (True,)) def test_memory_sharing(type, randomize, requires_grad, offload): # load model (maybe with offloading) model = TransformableModel(2, 2, 4, 4, 8, 8) if offload: set_onload_device(model, torch.device("cuda")) # add transforms to model config = TransformConfig( config_groups={ "": TransformScheme( type=type, randomize=randomize, requires_grad=requires_grad, apply=[ TransformArgs(targets="Linear", location="input"), TransformArgs(targets="Linear", location="output"), ], ) } ) apply_transform_config(model, config) for context in disable_onloading, disable_offloading: with context(): weights = [ m.weight for m in model.modules() if isinstance(m, TransformBase) ] weight_to_count = Counter(weights) size_to_weight = {weight.size(0): weight for weight in weight_to_count} assert len(weight_to_count) == len(size_to_weight) == 3 assert weight_to_count[size_to_weight[2]] == 3 assert weight_to_count[size_to_weight[4]] == 4 assert weight_to_count[size_to_weight[8]] == 3 vllm-project-compressed-tensors-c18a0fa/tests/test_transform/factory/test_serialization.py000066400000000000000000000052241521257237700326410ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os import pytest import torch from compressed_tensors.offload import set_onload_device from compressed_tensors.transform import ( TransformConfig, TransformScheme, apply_transform_config, ) from safetensors import safe_open from tests.testing_utils import requires_gpu from transformers import AutoModelForCausalLM, AutoTokenizer @pytest.mark.parametrize("type", ("hadamard", "random-hadamard")) @pytest.mark.parametrize("randomize", (True, False)) @pytest.mark.parametrize("offload", (True, False)) def test_serialization(type, randomize, model_apply, tmp_path, offload): # get model, maybe offload model, apply = model_apply if offload: set_onload_device(model, torch.device("cuda")) # apply transforms to model config = TransformConfig( config_groups={"": TransformScheme(type=type, randomize=randomize, apply=apply)} ) apply_transform_config(model, config) # save model model_path = os.path.join(tmp_path, "test_model_path") model.save_pretrained(model_path) # check that saved values match model values # note that shared weights are only serialized once safetensors_path = os.path.join(model_path, "model.safetensors") device = "cuda:0" if offload else "cpu" with safe_open(safetensors_path, framework="pt", device=device) as file: saved_keys = set(file.keys()) assert { "fcs.0.weight", "fcs.1.weight", "fcs.2.weight", "fcs.3.weight", "fcs.4.weight", } <= saved_keys for key in saved_keys: param = model.get_parameter(key) saved_param = file.get_tensor(key) assert torch.equal(param, saved_param) @pytest.mark.skip("Requires transformers#40673") @requires_gpu @pytest.mark.parametrize( "model_stub,exp_perplexity", [ ("nm-testing/Llama-3.2-1B-Instruct-spinquantR1R2R4-w4a16", 10.0), ("nm-testing/Llama-3.2-1B-Instruct-quip-w4a16", 10.0), ], ) def test_load_perplexity(model_stub, exp_perplexity): model = AutoModelForCausalLM.from_pretrained(model_stub, device_map="cuda") tokenizer = AutoTokenizer.from_pretrained(model_stub) prompt = "The capital of France is Paris, the capital of Germany is Berlin" inputs = tokenizer(prompt, return_tensors="pt") inputs = {key: value.to(model.device) for key, value in inputs.items()} labels = inputs["input_ids"] with torch.no_grad(): outputs = model(**inputs, labels=labels) perplexity = torch.exp(outputs.loss) assert perplexity <= exp_perplexity vllm-project-compressed-tensors-c18a0fa/tests/test_transform/test_transform_args.py000066400000000000000000000020431521257237700313400ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from compressed_tensors.transform import TransformArgs def test_basic(): targets = ["Embedding"] location = "input" args = TransformArgs(targets=targets, location=location) assert args.targets == targets assert args.location == location assert len(args.ignore) == 0 def test_args_full(): targets = ["Linear"] location = "weight_input" inverse = True ignore = ["model.layers.2"] args = TransformArgs( targets=targets, location=location, inverse=inverse, ignore=ignore, ) args.targets = targets args.location == location args.inverse == inverse args.ignore == ignore def test_singleton_targets(): target = "target" location = "input" ignore = "ignore" args = TransformArgs(targets=target, location=location, ignore=ignore) assert args.targets == [target] assert args.location == location assert args.ignore == [ignore] vllm-project-compressed-tensors-c18a0fa/tests/test_transform/test_transform_config.py000066400000000000000000000030011521257237700316440ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest from compressed_tensors.transform import TransformArgs, TransformConfig, TransformScheme @pytest.fixture(scope="module") def scheme(): targets = ["Embedding"] location = "input" basic_args = TransformArgs(targets=targets, location=location) return TransformScheme( type="hadamard", apply=[basic_args], ) @pytest.fixture(scope="module") def config(scheme): return TransformConfig( config_groups={ "transform_0": scheme, } ) def test_basic(config): assert isinstance(config.config_groups.get("transform_0"), TransformScheme) def test_to_dict(config): config_dict = config.model_dump() assert "config_groups" in config_dict.keys() def test_multiple_groups(): location = "weight_input" targets_1 = ["model.layers.0.attn.v_proj"] linear_args_1 = TransformArgs(targets=targets_1, location=location) targets_2 = ["model.layers.0.attn.q_proj"] linear_args_2 = TransformArgs(targets=targets_2, location=location) scheme_1 = TransformScheme( type="hadamard", apply=[linear_args_1], ) scheme_2 = TransformScheme( type="hadamard", apply=[linear_args_2], ) _ = TransformConfig( config_groups={"transform_0": scheme_1, "transform_1": scheme_2} ) def test_reload(config): assert config == TransformConfig.model_validate(config.model_dump()) vllm-project-compressed-tensors-c18a0fa/tests/test_transform/test_transform_scheme.py000066400000000000000000000034231521257237700316530ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from compressed_tensors.transform import TransformArgs, TransformScheme def test_basic_scheme(): targets = ["Linear"] location = "input" basic_args = TransformArgs(targets=targets, location=location) scheme = TransformScheme( type="hadamard", apply=[basic_args], ) assert not scheme.randomize assert scheme.type == "hadamard" assert len(scheme.apply) == 1 assert isinstance(scheme.apply[0], TransformArgs) def test_multiple_groups_global(): targets = ["Embedding"] location = "input" embedding_args = TransformArgs(targets=targets, location=location) targets = ["Linear"] location = "weight_input" linear_args = TransformArgs(targets=targets, location=location) # same transform applied to multiple groups scheme = TransformScheme( type="hadamard", apply=[embedding_args, linear_args], randomize=True, ) assert scheme.randomize assert scheme.type == "hadamard" assert len(scheme.apply) == 2 assert isinstance(scheme.apply[0], TransformArgs) assert isinstance(scheme.apply[1], TransformArgs) def test_multiple_groups(): apply = [] location = "weight_output" for i in range(20): targets = [f"model.layers.{i}.attn.v_proj", f"model.layers.{i}.attn.o_proj"] args = TransformArgs(targets=targets, location=location) apply.append(args) # global is False, different hadamard transform applied to each group # same dimension/hidden dim scheme = TransformScheme( type="hadamard", apply=apply, ) assert not scheme.randomize assert scheme.type == "hadamard" assert len(scheme.apply) == 20 vllm-project-compressed-tensors-c18a0fa/tests/test_transform/utils/000077500000000000000000000000001521257237700260415ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_transform/utils/test_hadamard.py000066400000000000000000000043731521257237700312220ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.transform.utils.hadamard import ( deterministic_hadamard_matrix, is_pow2, random_hadamard_matrix, ) from tests.testing_utils import requires_gpu _sizes_to_test = [ 768, # gpt2 small 1024, # gpt2 medium 1280, # qwen_2_5_vl vision 1600, # gpt2 xl 2048, # gpt3 small 3584, # qwen_2_5_vl 3840, # qwen_2_5_vl vision qkv 4096, # llama3 7168, # deepseek_v3 14336, # llama3 intermediate 18432, # deepseek_v3 intermediate 18944, # qwen_2_5_vl intermediate ] _atol = 1e-1 # bfloat16 is low precision for large matrices @requires_gpu @pytest.mark.parametrize("size", _sizes_to_test) def test_random_hadamard_matrix_compliant(size): # (H / sqrt(n))(H.T / sqrt(n)) == I matrix = random_hadamard_matrix(size, device="cuda") product = (matrix @ matrix.T) / matrix.size(0) eye = torch.eye(size, dtype=product.dtype, device="cuda") assert torch.allclose(product, eye, atol=_atol) def test_random_hadamard_generator(): # check that generation is deterministic with a seed generator = torch.Generator().manual_seed(42) one = random_hadamard_matrix(2048, gen=generator) two = random_hadamard_matrix(2048, gen=generator) one_true = torch.tensor( [ [-1, -1, -1], [+1, -1, +1], [-1, -1, +1], ] ) two_true = torch.tensor( [ [-1, -1, -1], [-1, +1, -1], [+1, +1, -1], ] ) assert torch.all(one[:3, :3].sign() == one_true.sign()) assert torch.all(two[:3, :3].sign() == two_true.sign()) @requires_gpu @pytest.mark.parametrize("size", _sizes_to_test) def test_deterministic_hadamard_compliant(size): if not is_pow2(size): with pytest.raises(ValueError): matrix = deterministic_hadamard_matrix(size, device="cuda") return # (H / sqrt(n))(H.T / sqrt(n)) == I matrix = deterministic_hadamard_matrix(size, device="cuda") product = (matrix @ matrix.T) / matrix.size(0) eye = torch.eye(size, dtype=product.dtype, device="cuda") assert torch.allclose(product, eye, atol=_atol) vllm-project-compressed-tensors-c18a0fa/tests/test_utils/000077500000000000000000000000001521257237700240265ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/tests/test_utils/__init__.py000066400000000000000000000001531521257237700261360ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project vllm-project-compressed-tensors-c18a0fa/tests/test_utils/test_helpers.py000066400000000000000000000032411521257237700271010ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace from compressed_tensors import ParameterizedDefaultDict, patch_attr, patch_attrs def test_patch_attr(): # patch, original value obj = SimpleNamespace() obj.attribute = "original" with patch_attr(obj, "attribute", "patched"): assert obj.attribute == "patched" obj.attribute = "modified" assert obj.attribute == "original" # patch, no original attribute obj = SimpleNamespace() with patch_attr(obj, "attribute", "patched"): assert obj.attribute == "patched" obj.attribute = "modified" assert not hasattr(obj, "attribute") def test_patch_attrs(): num_objs = 4 objs = [SimpleNamespace() for _ in range(num_objs)] for idx, obj in enumerate(objs): if idx % 2 == 0: obj.attribute = f"original_{idx}" with patch_attrs(objs, "attribute", [f"patched_{idx}" for idx in range(num_objs)]): for idx, obj in enumerate(objs): assert obj.attribute == f"patched_{idx}" obj.attribute = "modified" for idx, obj in enumerate(objs): if idx % 2 == 0: assert obj.attribute == f"original_{idx}" else: assert not hasattr(obj, "attribute") def test_parameterized_default_dict(): def add_one(value): return value + 1 add_dict = ParameterizedDefaultDict(add_one) assert add_dict[0] == 1 assert add_dict[1] == 2 def sum_vals(a, b): return a + b sum_dict = ParameterizedDefaultDict(sum_vals) assert sum_dict[0, 1] == 1 assert sum_dict[5, 7] == 12 vllm-project-compressed-tensors-c18a0fa/tests/test_utils/test_match.py000066400000000000000000001007531521257237700265410ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from unittest.mock import patch import pytest import torch import torch.nn as nn # Assuming the module is named "module_matching" - adjust import as needed from compressed_tensors.utils import ( InternalModule, get_lowest_common_ancestor_name, is_match, is_narrow_match, match_modules_set, match_name, match_named_modules, match_named_parameters, match_quantizable_tensors, ) from compressed_tensors.utils.match import _match_class from transformers import AutoModelForCausalLM @pytest.fixture def llama_stories_model(): return AutoModelForCausalLM.from_pretrained( "Xenova/llama2.c-stories15M", torch_dtype="auto", ) class DummyModel(nn.Module): """Test model for unit tests. Weights are initialized on meta device""" def __init__(self): super().__init__() self.layer1 = nn.Linear(10, 20) self.layer2 = nn.Linear(20, 30) self.norm = nn.LayerNorm(30) self.attention = nn.MultiheadAttention(30, 2) # Create nested structure self.transformer = nn.ModuleDict( { "layers": nn.ModuleList( [ nn.ModuleDict( { "self_attn": nn.ModuleDict( { "q_proj": nn.Linear(30, 30), "k_proj": nn.Linear(30, 30), "v_proj": nn.Linear(30, 30), } ), "norm": nn.LayerNorm(30), "mlp": nn.Linear(30, 30), } ) for _ in range(3) ] ) } ) class DummyMoEModel(nn.Module): """Test MoE model for unit tests. Weights are initialized on meta device""" def __init__(self, num_layers=2, num_experts=4): super().__init__() self.layers = nn.ModuleList( [ nn.ModuleDict( { "post_attention_layernorm": nn.LayerNorm(3), "mlp": nn.ModuleDict( { "experts": nn.ModuleList( [ nn.ModuleDict( { "gate_proj": nn.Linear(3, 6), "up_proj": nn.Linear(3, 6), "down_proj": nn.Linear(6, 3), } ) for _ in range(num_experts) ] ), } ), } ) for _ in range(num_layers) ] ) class TestMatchName: """Test cases for match_name function""" def test_exact_match(self): """Test exact string matching""" assert match_name("layer1", "layer1") assert not match_name("layer1", "layer2") assert match_name( "transformer.layers.0.self_attn.q_proj", "transformer.layers.0.self_attn.q_proj", ) def test_regex_match(self): """Test regex matching with "re:" prefix""" assert match_name("layer1", "re:layer.*") assert match_name("layer1", "re:^layer1$") assert not match_name("layer1", "re:layer2") assert match_name("transformer.layers.0.self_attn.q_proj", "re:.*q_proj") assert match_name( "transformer.layers.0.self_attn.q_proj", "re:transformer\\.layers\\.\\d+\\.self_attn\\..*_proj$", ) def test_empty_strings(self): """Test edge cases with empty strings""" assert match_name("", "") assert not match_name("layer1", "") assert not match_name("", "layer1") def test_regex_special_characters(self): """Test regex with special characters""" assert match_name("layer.1", "re:layer\\.1") assert match_name("layer.1", "re:layer.1") # . matches any char assert match_name("layer_1", "re:layer_1") class TestMatchClass: """Test cases for _match_class function""" def test_direct_class_match(self): """Test matching direct class names""" linear = nn.Linear(10, 20) assert _match_class(linear, "Linear") assert not _match_class(linear, "Conv2d") norm = nn.LayerNorm(10) assert _match_class(norm, "LayerNorm") assert not _match_class(norm, "BatchNorm1d") def test_parent_class_match(self): """Test matching parent class names""" linear = nn.Linear(10, 20) assert _match_class(linear, "Module") conv = nn.Conv2d(3, 16, 3) assert _match_class(conv, "Module") assert _match_class(conv, "_ConvNd") def test_non_torch_module(self): """Test with non-torch modules""" regular_object = object() assert not _match_class(regular_object, "object") # not a torch.nn.Module def test_custom_module(self): """Test with custom module classes""" model = DummyModel() assert _match_class(model, "DummyModel") assert _match_class(model, "Module") def test_linear_base(self): """Test matching against vllm's LinearBase class""" class LinearBase(nn.Module): pass linear = LinearBase() assert _match_class(linear, "Linear") class TestIsMatch: """Test cases for is_match function""" def test_name_match(self): """Test matching by name""" linear = nn.Linear(10, 20) assert is_match("layer1", linear, "layer1") assert not is_match("layer1", linear, "layer2") def test_class_match(self): """Test matching by class""" linear = nn.Linear(10, 20) assert is_match("layer1", linear, "Linear") assert not is_match("layer1", linear, "Conv2d") def test_combined_match(self): """Test that either name or class match works""" linear = nn.Linear(10, 20) assert is_match("layer1", linear, "layer1") # name match assert is_match("layer1", linear, "Linear") # class match assert not is_match("layer1", linear, "layer2") # no match def test_regex_in_name_match(self): """Test regex matching in name""" linear = nn.Linear(10, 20) assert is_match("layer1", linear, "re:layer.*") assert not is_match("layer1", linear, "re:conv.*") def test_internal_module_match(self): """Test not matching internal modules""" class InternalLinear(InternalModule, nn.Linear): pass linear = InternalLinear(10, 20) assert not is_match("layer1", linear, "re:layer.*") def test_fused_mapping(self): """""" linear = nn.Linear(10, 20) mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], } assert is_match("dummy.qkv_proj", linear, "re:.*q_proj", fused=mapping) assert is_match("dummy.qkv_proj", linear, "re:.*k_proj", fused=mapping) assert is_match("dummy.qkv_proj", linear, "re:.*v_proj", fused=mapping) assert is_match("dummy.qkv_proj", linear, "Linear", fused=mapping) assert is_match("dummy.gate_up_proj", linear, "re:.*gate_proj", fused=mapping) assert is_match("dummy.gate_up_proj", linear, "re:.*up_proj", fused=mapping) assert is_match("dummy.gate_up_proj", linear, "Linear", fused=mapping) class TestMatchNamedModules: """Test cases for match_named_modules function""" def test_exact_module_match(self): """Test matching modules by exact name""" model = DummyModel() matches = list(match_named_modules(model, ["layer1", "layer2"])) assert len(matches) == 2 names = [name for name, _ in matches] assert "layer1" in names assert "layer2" in names def test_class_module_match(self): """Test matching modules by class name""" model = DummyModel() matches = list(match_named_modules(model, ["Linear"])) # Should find all Linear layers linear_modules = [ module for _, module in matches if isinstance(module, nn.Linear) ] assert len(linear_modules) > 0 def test_regex_module_match(self): """Test matching modules with regex patterns""" model = DummyModel() matches = list(match_named_modules(model, ["re:.*linear.*"])) # Should find layers with "linear" in name (case insensitive depends on model) assert len(matches) >= 0 # May be 0 if no "linear" in names def test_ignore_parameter(self): """Test ignoring specific modules""" model = DummyModel() matches_without_ignore = list(match_named_modules(model, ["Linear"])) matches_with_ignore = list( match_named_modules(model, ["Linear"], ignore=["layer1"]) ) # Should have fewer matches when ignoring layer1 assert len(matches_with_ignore) <= len(matches_without_ignore) # layer1 should not be in ignored results ignored_names = [name for name, _ in matches_with_ignore] assert "layer1" not in ignored_names def test_empty_targets(self): """Test with empty targets list""" model = DummyModel() matches = list(match_named_modules(model, [])) assert len(matches) == 0 @patch("compressed_tensors.utils.match._LOGGER") def test_warn_on_fail(self, mock_logger): """Test warning when targets don"t match""" model = DummyModel() list(match_named_modules(model, ["nonexistent_module"], warn_on_fail=True)) mock_logger.warning.assert_called_once() warning_msg = mock_logger.warning.call_args[0][0] assert "Could not match" in warning_msg assert "nonexistent_module" in warning_msg def test_internal_match(self): """Test not matching internal modules""" class InternalLinear(InternalModule, nn.Linear): pass linear = InternalLinear(10, 20) matches = list(match_named_modules(linear, ["re:.*"])) assert len(matches) == 0 @pytest.mark.parametrize( "targets, ignore, expected_targets", [ ( ["re:model.layers.[01].self_attn.q_proj"], ["re:model.layers.1.self_attn.q_proj"], set(["model.layers.0.self_attn.q_proj"]), ), ( ["re:model.layers.[01].self_attn.q_proj"], [], set( [ "model.layers.0.self_attn.q_proj", "model.layers.1.self_attn.q_proj", ] ), ), ( ["re:model.layers.[0-2].self_attn.q_proj"], ["re:model.layers.1.self_attn.q_proj"], set( [ "model.layers.0.self_attn.q_proj", "model.layers.2.self_attn.q_proj", ] ), ), ( ["model.layers.0.self_attn.q_proj"], ["model.layers.0.self_attn.q_proj"], set(), ), ( ["re:model.layers.*.self_attn.q_proj"], ["re:model.layers.[01].self_attn.q_proj"], set( f"model.layers.{layer_idx}.self_attn.q_proj" for layer_idx in range(2, 6) ), ), ], ) def test_expand_targets_with_llama_stories( self, llama_stories_model, targets, ignore, expected_targets ): expanded_targets = { name for name, _ in match_named_modules(llama_stories_model, targets, ignore) } assert expanded_targets == expected_targets class TestMatchNamedParameters: """Test cases for match_named_parameters function""" def test_parameter_match(self): """Test matching parameters by name""" model = DummyModel() matches = list(match_named_parameters(model, ["layer1.weight", "layer1.bias"])) assert len(matches) == 2 param_names = [name for name, _, _ in matches] assert "layer1.weight" in param_names assert "layer1.bias" in param_names def test_regex_parameter_match(self): """Test matching parameters with regex""" model = DummyModel() matches = list(match_named_parameters(model, ["re:.*weight$"])) # Should find all weight parameters weight_params = [name for name, _, _ in matches if name.endswith(".weight")] assert len(weight_params) > 0 def test_ignore_parameters(self): """Test ignoring specific parameters""" model = DummyModel() matches_without_ignore = list(match_named_parameters(model, ["re:.*weight$"])) matches_with_ignore = list( match_named_parameters(model, ["re:.*weight$"], ignore=["layer1.weight"]) ) # Should have fewer matches when ignoring assert len(matches_with_ignore) < len(matches_without_ignore) # layer1.weight should not be in ignored results ignored_names = [name for name, _, _ in matches_with_ignore] assert "layer1.weight" not in ignored_names def test_parameter_return_values(self): """Test that function returns correct tuple values""" model = DummyModel() matches = list(match_named_parameters(model, ["layer1.weight"])) assert len(matches) == 1 param_name, parent_module, param = matches[0] assert param_name == "layer1.weight" assert parent_module is model.layer1 assert isinstance(param, nn.Parameter) assert param is model.layer1.weight @patch("compressed_tensors.utils.match._LOGGER") def test_warn_on_fail_parameters(self, mock_logger): """Test warning when parameter targets don"t match""" model = DummyModel() list(match_named_parameters(model, ["nonexistent.param"], warn_on_fail=True)) mock_logger.warning.assert_called_once() warning_msg = mock_logger.warning.call_args[0][0] assert "Could not match" in warning_msg assert "nonexistent.param" in warning_msg def test_internal_match(self): """Test not matching internal modules""" class InternalLinear(InternalModule, nn.Linear): pass linear = InternalLinear(10, 20) matches = list(match_named_parameters(linear, ["re:.*"])) assert len(matches) == 0 class TestGetLowestCommonModuleName: """Test cases for get_lowest_common_ancestor_name function""" def test_multiple_modules(self): assert "abc" == get_lowest_common_ancestor_name( [ "abc.a", "abc.b", "abc.c", ] ) def test_single_module(self): assert "abc.abc" == get_lowest_common_ancestor_name( [ "abc.abc", ] ) def test_substring_modules(self): assert "abc" == get_lowest_common_ancestor_name( [ "abc.abc", "abc.ab", ] ) def test_parent_and_child_modules(self): assert "abc.abc" == get_lowest_common_ancestor_name( [ "abc.abc.ab", "abc.abc", ] ) def test_root(self): assert "" == get_lowest_common_ancestor_name( [ "abc.abc", "b.abc", ] ) def test_ignore_none(self): assert "abc.abc" == get_lowest_common_ancestor_name( [ "abc.abc", None, ] ) class TestMatchModulesSet: """Test cases for match_modules_set function""" def test_simple_module_set(self): """Test matching simple module sets""" model = DummyModel() targets = [ "re:.*self_attn.q_proj$", "re:.*self_attn.k_proj$", "re:.*self_attn.v_proj$", ] matches = list(match_modules_set(model, targets)) # Should have 3 sets (one for each layer) assert len(matches) == 3 # Each set should have 3 modules for module_set in matches: assert len(module_set) == 3 assert all(isinstance(*m, nn.Linear) for m in module_set) def test_moe_module_match(self): """Test matching MoE modules with multiple experts per layer""" model = DummyMoEModel(num_layers=2, num_experts=4) # Test matching expert projections - each expert becomes its own set # because the parent context differs between experts targets = [ "re:.*gate_proj$", "re:.*up_proj$", ] matches = list(match_modules_set(model, targets)) # Should have 8 sets (2 layers * 4 experts) assert len(matches) == 8 # Each set should have 2 target lists (gate_proj, up_proj) for expert_group in matches: assert len(expert_group) == 2 gate_modules, up_modules = expert_group # Each target should have matched 1 module (single expert) assert len(gate_modules) == 1 assert len(up_modules) == 1 # All modules should be Linear layers assert isinstance(gate_modules[0], nn.Linear) assert isinstance(up_modules[0], nn.Linear) def test_moe_with_layernorm_match(self): """ Test matching MoE modules with their corresponding layer norms. Including a layer-level module (layernorm) groups all experts in that layer together. """ model = DummyMoEModel(num_layers=2, num_experts=3) # Match layer norm with expert projections - the layernorm is at layer level, # so it establishes a common parent context for all experts in that layer targets = [ "re:.*post_attention_layernorm$", "re:.*gate_proj$", "re:.*up_proj$", ] matches = list(match_modules_set(model, targets)) # Should have 2 layer groups (one per layer) assert len(matches) == 2 for layer_group in matches: assert len(layer_group) == 3 norm_modules, gate_modules, up_modules = layer_group # LayerNorm should have 1 module (single per layer) assert len(norm_modules) == 1 assert isinstance(norm_modules[0], nn.LayerNorm) # Each projection should have 3 experts (all experts in the layer) assert len(gate_modules) == 3 assert len(up_modules) == 3 assert all(isinstance(m, nn.Linear) for m in gate_modules) assert all(isinstance(m, nn.Linear) for m in up_modules) def test_module_set_ordering(self): """Test that module sets maintain target ordering""" model = DummyModel() targets = [ "re:.*v_proj$", # v first "re:.*self_attn.q_proj$", # q second "re:.*self_attn.k_proj$", ] # k third matches = list(match_modules_set(model, targets)) for module_set in matches: # Check that modules are returned in target order (v, q, k) v_proj, q_proj, k_proj = module_set v_proj, q_proj, k_proj = *v_proj, *q_proj, *k_proj # We can't easily check the exact modules, but can check they're all Linear assert all(isinstance(m, nn.Linear) for m in [v_proj, q_proj, k_proj]) def test_incomplete_set_error(self): """Test error when unable to complete a set""" model = DummyModel() targets = ["layer1", "nonexistent_module"] with pytest.raises( ValueError, match="Found a final incomplete set with matches found for keys" ): list(match_modules_set(model, targets)) def test_empty_targets_set(self): """Test with empty targets""" model = DummyModel() matches = list(match_modules_set(model, [])) # Should yield one empty set for each module traversed? # Actually, with empty targets, we expect no matches assert len(matches) == 0 def test_module_set_with_ignore(self): """Test module set matching with ignore parameter""" model = DummyModel() targets = ["re:.*self_attn.q_proj$", "re:.*self_attn.k_proj$"] ignore = ["re:transformer.layers.0.*"] # Ignore first layer matches = list(match_modules_set(model, targets, ignore=ignore)) # Should have 2 sets (layers 1 and 2, but not 0) assert len(matches) == 2 def test_internal_match(self): """Test not matching internal modules""" class InternalLinear(InternalModule, nn.Linear): pass linear = InternalLinear(10, 20) matches = list(match_modules_set(linear, ["re:.*"])) assert len(matches) == 0 class TestIsNarrowMatch: def test_narrow_match_true_child_only(self): """ Target matches the child module name but NOT its parent name. Should return True. """ model = DummyModel() name = "transformer.layers.0.self_attn.q_proj" # Matches "...q_proj" but not "...self_attn" target = r"re:.*q_proj$" assert is_narrow_match(model, target, name) def test_narrow_match_false_when_parent_also_matches(self): """ Target matches both the child and its parent name. Should return False because it's not a 'narrow' match. """ model = DummyModel() name = "transformer.layers.0.self_attn.q_proj" # Broad target that also matches the parent "transformer.layers.0.self_attn" target = r"re:transformer\.layers\.0\..*" assert not is_narrow_match(model, target, name) def test_narrow_match_false_when_neither_matches(self): """ Target matches neither the child nor the parent. Should return False. """ model = DummyModel() name = "transformer.layers.0.self_attn.q_proj" target = r"re:this_does_not_exist$" assert not is_narrow_match(model, target, name) def test_narrow_match_iterable_targets_any_true(self): """ With multiple targets: if any target narrowly matches the child, the function should return True. """ model = DummyModel() name = "transformer.layers.0.self_attn.q_proj" # First target is broad (matches both child & parent -> narrow False), # second target is narrow (matches child only -> narrow True). targets = [ r"re:transformer\.layers\.0\..*", r"re:.*q_proj$", ] assert is_narrow_match(model, targets, name) def test_narrow_match_with_explicit_module_argument(self): """ Passing the module explicitly should behave the same as when it's retrieved from the model by name. """ model = DummyModel() name = "transformer.layers.0.self_attn.q_proj" module = model.get_submodule(name) target = r"re:.*q_proj$" # Both ways should be True assert is_narrow_match(model, target, name) assert is_narrow_match(model, target, name, module=module) def test_narrow_match_top_level_behavior_documented(self): """ (Behavior check) For a top-level module name without a dot, the current implementation derives parent_name == name, so parent==child. Then 'narrow' cannot be True because parent match mirrors child match. This test documents current behavior to guard against regressions. """ model = DummyModel() name = "layer1" # top-level module in DummyModel target = r"re:^layer1$" assert not is_narrow_match(model, target, name) class TestIntegration: """Integration tests combining multiple functions""" def test_complex_model_matching(self): """Test matching on a more complex model structure""" model = DummyModel() # Test that we can find attention projection layers q_matches = list(match_named_modules(model, ["re:.*q_proj$"])) k_matches = list(match_named_modules(model, ["re:.*k_proj$"])) v_matches = list(match_named_modules(model, ["re:.*v_proj$"])) assert len(q_matches) == 3 # 3 layers assert len(k_matches) == 3 assert len(v_matches) == 3 def test_parameter_and_module_consistency(self): """Test that parameter and module matching are consistent""" model = DummyModel() # Get modules module_matches = list(match_named_modules(model, ["layer1"])) assert len(module_matches) == 1 module_name, module = module_matches[0] # Get parameters from that module param_matches = list(match_named_parameters(model, [f"{module_name}.weight"])) assert len(param_matches) == 1 param_name, parent_module, param = param_matches[0] # Check consistency assert parent_module is module assert param is module.weight def test_all_functions_with_regex(self): """Test all functions work with regex patterns""" model = DummyModel() regex_target = "re:.*Linear.*" # Should not crash and should handle regex consistently modules = list(match_named_modules(model, [regex_target])) params = list(match_named_parameters(model, [regex_target])) # Basic sanity checks assert isinstance(modules, list) assert isinstance(params, list) @pytest.fixture def sample_tensors(): """Create a sample set of tensors mimicking a model's state dict.""" return { "model.layers.0.self_attn.q_proj.weight": torch.randn(128, 128), "model.layers.0.self_attn.k_proj.weight": torch.randn(128, 128), "model.layers.0.self_attn.v_proj.weight": torch.randn(128, 128), "model.layers.0.mlp.gate_proj.weight": torch.randn(256, 128), "model.layers.0.mlp.up_proj.weight": torch.randn(256, 128), "model.layers.0.mlp.down_proj.weight": torch.randn(128, 256), "model.layers.0.input_layernorm.weight": torch.randn(128), "model.layers.0.post_attention_layernorm.weight": torch.randn(128), "model.embed_tokens.weight": torch.randn(32000, 128), "lm_head.weight": torch.randn(32000, 128), "model.layers.0.self_attn.q_proj.bias": torch.randn(128), } @pytest.mark.parametrize( "ignore,targets,param_targets,allow_nonquantizable,expected_names", [ # Test case: basic matching without ignore or targets ( [], [], ("weight",), False, { "model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.k_proj.weight", "model.layers.0.self_attn.v_proj.weight", "model.layers.0.mlp.gate_proj.weight", "model.layers.0.mlp.up_proj.weight", "model.layers.0.mlp.down_proj.weight", "model.embed_tokens.weight", "lm_head.weight", }, ), # Test case: ignore attention layers ( ["re:.*self_attn.*"], [], ("weight",), False, { "model.layers.0.mlp.gate_proj.weight", "model.layers.0.mlp.up_proj.weight", "model.layers.0.mlp.down_proj.weight", "model.embed_tokens.weight", "lm_head.weight", }, ), # Test case: ignore attention and mlp layers ( ["re:.*self_attn.*", "re:.*mlp.*"], [], ("weight",), False, { "model.embed_tokens.weight", "lm_head.weight", }, ), # Test case: target only mlp gate_proj and up_proj ( [], ["re:.*mlp.*gate_proj", "re:.*mlp.*up_proj"], ("weight",), False, { "model.layers.0.mlp.gate_proj.weight", "model.layers.0.mlp.up_proj.weight", }, ), # Test case: empty targets (all-inclusive) ( [], [], ("weight",), False, { "model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.k_proj.weight", "model.layers.0.self_attn.v_proj.weight", "model.layers.0.mlp.gate_proj.weight", "model.layers.0.mlp.up_proj.weight", "model.layers.0.mlp.down_proj.weight", "model.embed_tokens.weight", "lm_head.weight", }, ), # Test case: Linear targets (all-inclusive) ( [], ["Linear"], ("weight",), False, { "model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.k_proj.weight", "model.layers.0.self_attn.v_proj.weight", "model.layers.0.mlp.gate_proj.weight", "model.layers.0.mlp.up_proj.weight", "model.layers.0.mlp.down_proj.weight", "model.embed_tokens.weight", "lm_head.weight", }, ), # Test case: allow_nonquantizable includes bias and layernorm ( [], [], ("weight", "bias"), True, { "model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.k_proj.weight", "model.layers.0.self_attn.v_proj.weight", "model.layers.0.mlp.gate_proj.weight", "model.layers.0.mlp.up_proj.weight", "model.layers.0.mlp.down_proj.weight", "model.layers.0.input_layernorm.weight", "model.layers.0.post_attention_layernorm.weight", "model.embed_tokens.weight", "lm_head.weight", "model.layers.0.self_attn.q_proj.bias", }, ), # Test case: ignore takes precedence over targets ( ["re:.*self_attn.*"], ["re:.*self_attn.*q_proj"], ("weight",), False, set(), ), # Test case: regex pattern matching all proj layers ( [], ["re:.*proj$"], ("weight",), False, { "model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.k_proj.weight", "model.layers.0.self_attn.v_proj.weight", "model.layers.0.mlp.gate_proj.weight", "model.layers.0.mlp.up_proj.weight", "model.layers.0.mlp.down_proj.weight", }, ), ], ids=[ "basic_matching", "ignore_attention", "ignore_attention_and_mlp", "target_mlp_gate_up", "empty_targets", "linear_targets", "allow_nonquantizable", "ignore_precedence", "regex_all_proj", ], ) def test_match_quantizable_tensors( sample_tensors, ignore, targets, param_targets, allow_nonquantizable, expected_names ): """ Parameterized test for match_quantizable_tensors function. Tests various combinations of ignore patterns, target patterns, and flags to verify that the function returns the expected set of tensor names. """ matches = list( match_quantizable_tensors( sample_tensors, ignore=ignore, targets=targets, param_targets=param_targets, allow_nonquantizable=allow_nonquantizable, ) ) # Extract full names from results result_names = {full_name for _, full_name in matches} # Assert the result matches expected assert result_names == expected_names # Additionally verify all results are tuples with correct format for module_name, full_name in matches: assert isinstance(module_name, str) assert isinstance(full_name, str) # module_name should be full_name without the last component assert full_name.startswith(module_name) assert full_name.rsplit(".", 1)[0] == module_name vllm-project-compressed-tensors-c18a0fa/tests/test_utils/test_safetensors_load.py000066400000000000000000000042121521257237700307710ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from unittest.mock import patch import pytest from compressed_tensors.utils.safetensors_load import get_nested_weight_mappings mock_weight_mappings = { "layer1.weight": "file1", "layer1.bias": "file2", "layer2.weight": "file3", "layer2.bias": "file4", "layer3.weight": "file5", } @pytest.fixture def mock_get_weight_mappings(): with patch( "compressed_tensors.utils.safetensors_load.get_weight_mappings", return_value=mock_weight_mappings, ): yield @pytest.mark.usefixtures("mock_get_weight_mappings") class TestGetNestedWeightMappings: """ Tests for the get_nested_weight_mappings function in different scenarios, such as single and multiple parameters to nest, and returning other parameters """ def test_single_param(self): params_to_nest = ["weight"] result = get_nested_weight_mappings("dummy_path", params_to_nest) expected = { "layer1": {"weight": "file1"}, "layer2": {"weight": "file3"}, "layer3": {"weight": "file5"}, } assert result == expected def test_multiple_params(self): params_to_nest = ["weight", "bias"] result = get_nested_weight_mappings("dummy_path", params_to_nest) expected = { "layer1": {"weight": "file1", "bias": "file2"}, "layer2": {"weight": "file3", "bias": "file4"}, "layer3": {"weight": "file5"}, } assert result == expected def test_return_other_params(self): params_to_nest = ["weight"] result, other_params = get_nested_weight_mappings( "dummy_path", params_to_nest, return_unmatched_params=True ) expected_nested = { "layer1": {"weight": "file1"}, "layer2": {"weight": "file3"}, "layer3": {"weight": "file5"}, } expected_other = { "layer1.bias": "file2", "layer2.bias": "file4", } assert result == expected_nested assert other_params == expected_other vllm-project-compressed-tensors-c18a0fa/tests/test_utils/test_save_mtp_tensors.py000066400000000000000000000212371521257237700310370ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import os import pytest import torch from compressed_tensors.utils.mtp import save_mtp_tensors_to_checkpoint from safetensors import safe_open from safetensors.torch import save_file from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_safetensors(path: str, tensors: dict) -> None: """Save a dict of tensors as a safetensors file.""" save_file({k: v.contiguous() for k, v in tensors.items()}, path) def _read_safetensors(path: str) -> dict: """Load all tensors from a safetensors file into a plain dict.""" result = {} with safe_open(path, framework="pt", device="cpu") as f: for key in f.keys(): result[key] = f.get_tensor(key) return result # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @pytest.fixture() def source_dir(tmp_path): """ Multi-shard source checkpoint containing both regular model tensors and MTP tensors (prefixed with "mtp"), along with a model.safetensors.index.json that maps each key to its shard file. """ src = tmp_path / "source" src.mkdir() shard1 = { "model.layer0.weight": torch.randn(4, 4), "mtp.layer0.weight": torch.randn(3, 3), } shard2 = { "model.layer1.weight": torch.randn(4, 4), "mtp.layer1.weight": torch.randn(3, 3), } _make_safetensors(str(src / "model-00001-of-00002.safetensors"), shard1) _make_safetensors(str(src / "model-00002-of-00002.safetensors"), shard2) index = { "metadata": {}, "weight_map": { "model.layer0.weight": "model-00001-of-00002.safetensors", "mtp.layer0.weight": "model-00001-of-00002.safetensors", "model.layer1.weight": "model-00002-of-00002.safetensors", "mtp.layer1.weight": "model-00002-of-00002.safetensors", }, } with open(src / SAFE_WEIGHTS_INDEX_NAME, "w") as f: json.dump(index, f) return src @pytest.fixture() def dest_dir_with_index(tmp_path): """ Destination checkpoint directory pre-populated with a single model shard and a model.safetensors.index.json. Represents the common quantized-model output where an index already exists and should be updated in place. """ dest = tmp_path / "dest_index" dest.mkdir() tensors = {"model.layer0.weight": torch.randn(4, 4)} shard_name = "model-00001-of-00001.safetensors" _make_safetensors(str(dest / shard_name), tensors) index = { "metadata": {}, "weight_map": {"model.layer0.weight": shard_name}, } with open(dest / SAFE_WEIGHTS_INDEX_NAME, "w") as f: json.dump(index, f) return dest # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- class TestSaveMtpTensorsToCheckpoint: """ Tests for save_mtp_tensors_to_checkpoint, which extracts tensors whose keys start with a given prefix (default "mtp") from a source checkpoint and appends them as a new shard to a destination checkpoint directory, updating the destination's model.safetensors.index.json accordingly. """ def test_mtp_tensors_saved_correctly(self, source_dir, dest_dir_with_index): """ Verify that the MTP shard is created in the destination directory and that its tensor values are numerically identical to those in the source. Also checks that only MTP-prefixed keys are included in the shard (no regular model weights leak through). """ # Collect expected MTP tensors directly from source shards expected = {} for shard in ( "model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors", ): with safe_open(str(source_dir / shard), framework="pt", device="cpu") as f: for key in f.keys(): if key.startswith("mtp"): expected[key] = f.get_tensor(key) save_mtp_tensors_to_checkpoint(str(source_dir), str(dest_dir_with_index)) mtp_shard_path = str(dest_dir_with_index / "model_mtp.safetensors") assert os.path.exists(mtp_shard_path) saved = _read_safetensors(mtp_shard_path) assert set(saved.keys()) == set(expected.keys()) for key in expected: assert torch.equal(saved[key], expected[key]) assert all(k.startswith("mtp") for k in saved) def test_index_updated(self, source_dir, dest_dir_with_index): """ Verify that after the call the destination index file contains MTP keys pointing to the new shard, existing keys are preserved, and metadata.total_size reflects the sizes of all shards in the weight_map. """ save_mtp_tensors_to_checkpoint(str(source_dir), str(dest_dir_with_index)) with open(dest_dir_with_index / SAFE_WEIGHTS_INDEX_NAME) as f: index = json.load(f) weight_map = index["weight_map"] # MTP keys added, pointing to the new shard assert weight_map.get("mtp.layer0.weight") == "model_mtp.safetensors" assert weight_map.get("mtp.layer1.weight") == "model_mtp.safetensors" # Pre-existing keys must not be removed assert "model.layer0.weight" in weight_map # total_size must equal the sum of all referenced shard sizes on disk expected_size = sum( os.path.getsize(dest_dir_with_index / s) for s in set(weight_map.values()) ) assert index["metadata"]["total_size"] == expected_size def test_single_shard_dest_creates_index(self, source_dir, tmp_path): """ Verify that when the destination has no index file (only a model.safetensors), the function synthesises one that includes both the original model keys and the newly added MTP keys. """ dest = tmp_path / "dest_single" dest.mkdir() _make_safetensors( str(dest / SAFE_WEIGHTS_NAME), {"model.layer0.weight": torch.randn(4, 4)} ) save_mtp_tensors_to_checkpoint(str(source_dir), str(dest)) index_path = dest / SAFE_WEIGHTS_INDEX_NAME assert index_path.exists() with open(index_path) as f: index = json.load(f) assert index["weight_map"].get("model.layer0.weight") == SAFE_WEIGHTS_NAME assert index["weight_map"].get("mtp.layer0.weight") == "model_mtp.safetensors" def test_no_mtp_tensors_raises(self, dest_dir_with_index, tmp_path): """ Verify that a ValueError is raised when the source checkpoint has no tensors matching the MTP prefix, preventing a silent empty-shard write. """ src = tmp_path / "src_no_mtp" src.mkdir() _make_safetensors( str(src / SAFE_WEIGHTS_NAME), {"model.weight": torch.randn(4, 4)} ) with pytest.raises(ValueError, match="No tensors with prefix"): save_mtp_tensors_to_checkpoint(str(src), str(dest_dir_with_index)) def test_missing_dest_files_raises(self, source_dir, tmp_path): """ Verify that a FileNotFoundError is raised when the destination directory contains neither model.safetensors.index.json nor model.safetensors. """ empty_dest = tmp_path / "dest_empty" empty_dest.mkdir() with pytest.raises(ValueError): save_mtp_tensors_to_checkpoint(str(source_dir), str(empty_dest)) def test_custom_mtp_prefix(self, dest_dir_with_index, tmp_path): """ Verify that when a non-default mtp_prefix is supplied, only tensors whose keys start with that prefix are extracted. Tensors matching the default "mtp" prefix but not the custom one must be excluded. """ src = tmp_path / "src_custom" src.mkdir() _make_safetensors( str(src / SAFE_WEIGHTS_NAME), { "model.weight": torch.randn(4, 4), "speculative.layer0.weight": torch.randn(3, 3), "mtp.layer0.weight": torch.randn(3, 3), # should be ignored }, ) save_mtp_tensors_to_checkpoint( str(src), str(dest_dir_with_index), mtp_prefix="speculative" ) saved = _read_safetensors(str(dest_dir_with_index / "model_mtp.safetensors")) assert set(saved.keys()) == {"speculative.layer0.weight"} vllm-project-compressed-tensors-c18a0fa/tests/test_utils/test_type.py000066400000000000000000000036151521257237700264250ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest import torch from compressed_tensors.utils.type import TorchDtype from pydantic import BaseModel, Field from pydantic_core._pydantic_core import ValidationError class DummyModel(BaseModel): dtype: TorchDtype = Field(default=torch.float32) @pytest.mark.unit def test_default_value(): model = DummyModel() assert model.dtype == torch.float32 @pytest.mark.unit def test_value_override(): model = DummyModel() model.dtype = torch.float16 assert model.dtype == torch.float16 @pytest.mark.unit def test_validation(): DummyModel(dtype=torch.float16) DummyModel(dtype="torch.float16") DummyModel(dtype="float16") with pytest.raises(ValidationError): _ = DummyModel(dtype="notatype") @pytest.mark.unit def test_serialization(): model = DummyModel() assert model.model_dump()["dtype"] == "torch.float32" assert DummyModel.model_validate(model.model_dump()) == model model = DummyModel(dtype=torch.float16) assert model.model_dump()["dtype"] == "torch.float16" assert DummyModel.model_validate(model.model_dump()) == model model = DummyModel() model.dtype = torch.float16 assert model.model_dump()["dtype"] == "torch.float16" assert DummyModel.model_validate(model.model_dump()) == model @pytest.mark.unit def test_deserialization(): dummy_dict = {"dtype": "torch.float16"} assert DummyModel.model_validate(dummy_dict).dtype == torch.float16 dummy_dict = {"dtype": "float16"} assert DummyModel.model_validate(dummy_dict).dtype == torch.float16 with pytest.raises(ValueError): dummy_dict = {"dtype": "notatype"} DummyModel.model_validate(dummy_dict) with pytest.raises(ValueError): dummy_dict = {"dtype": "torch.notatype"} DummyModel.model_validate(dummy_dict) vllm-project-compressed-tensors-c18a0fa/tests/testing_utils.py000066400000000000000000000100651521257237700251000ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # flake8: noqa import pytest import torch def compressed_tensors_config_available(): try: from transformers.utils.quantization_config import ( # noqa: F401 CompressedTensorsConfig, ) return True except ImportError: return False _is_compressed_tensors_config_available = compressed_tensors_config_available() def requires_hf_quantizer(): return pytest.mark.skipif( not _is_compressed_tensors_config_available, reason="requires transformers>=4.45 to support CompressedTensorsHfQuantizer", ) def get_random_mat(M, K, dtype) -> "torch.Tensor": """ :param M: number of rows :param K: number of columns :param dtype: data type of the matrix :return: random matrix of shape (M, K) with non-zero values """ import torch from compressed_tensors.quantization import FP8_DTYPE rand_tensor_dtype = dtype if dtype in [torch.int8, FP8_DTYPE]: rand_tensor_dtype = torch.float16 mat = torch.rand(M, K, dtype=rand_tensor_dtype).cuda() mat = mat.masked_fill_(mat == 0, 1) return mat.to(dtype) def generate_pruned_semi_structured_mat(M, K, dtype) -> "torch.Tensor": """ :param M: number of rows :param K: number of columns :param dtype: data type of the matrix :return: random matrix of shape (M, K) with 2:4 sparsity pattern """ import torch from compressed_tensors.quantization import FP8_E4M3_DATA mask = torch.Tensor([0, 0, 1, 1]).tile((M, K // 4)).bool() rand_tensor_dtype = dtype if dtype in [torch.int8, FP8_E4M3_DATA.dtype]: rand_tensor_dtype = torch.float16 mat = torch.rand(M, K, dtype=rand_tensor_dtype) mat = mat.masked_fill_(mat == 0, 1) if dtype == FP8_E4M3_DATA.dtype: # some float8_e4m3fn operations are not supported on CPU mat = mat.cuda() mask = mask.cuda() mat = mat * mask return mat.to(dtype) def induce_sparsity(tensor, sparsity_ratio) -> "torch.Tensor": """ Makes a tensor sparse by zeroing out a given fraction of its smallest absolute values. :param: weight_tensor (torch.Tensor): The input weight tensor. :param: sparsity_ratio (float): Fraction of weights to be zeroed (0 <= sparsity_ratio <= 1). :returns: torch.Tensor: Sparse version of the input tensor. """ import torch if not (0 <= sparsity_ratio <= 1): raise ValueError("Sparsity ratio must be between 0 and 1.") # Flatten the tensor and compute the threshold for sparsity flattened = tensor.view(-1) k = int(sparsity_ratio * flattened.numel()) if k > 0: threshold = torch.topk(flattened.abs(), k, largest=False).values.max() sparse_tensor = torch.where( tensor.abs() > threshold, tensor, torch.zeros_like(tensor) ) else: sparse_tensor = tensor return sparse_tensor def is_gpu_available(): """ :return: True if an accelerator device is available, False otherwise """ try: import torch # noqa: F401 return torch.accelerator.device_count() > 0 except ImportError: return False def requires_gpu(test_case_or_num): """ Pytest decorator to skip based on number of available GPUs. Designed for backwards compatibility with the old requires_gpu decorator Usage: @requires_gpu def test_something(): # only runs if there is at least 1 GPU available pass @requires_gpu(2) def test_something_else(): # only runs if there are at least 2 GPUs available pass """ if isinstance(test_case_or_num, int): num_required_gpus = test_case_or_num else: num_required_gpus = 1 decorator = pytest.mark.skipif( (torch.accelerator.device_count() < num_required_gpus), reason=f"Not enough GPUs available, {num_required_gpus} GPUs required", ) if isinstance(test_case_or_num, int): return decorator else: return decorator(test_case_or_num) vllm-project-compressed-tensors-c18a0fa/utils/000077500000000000000000000000001521257237700216255ustar00rootroot00000000000000vllm-project-compressed-tensors-c18a0fa/utils/copyright.py000066400000000000000000000225541521257237700242170ustar00rootroot00000000000000# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse import glob import os import sys from typing import List, NamedTuple COPYRIGHT_LINES = [ "SPDX-License-Identifier: Apache-2.0", "SPDX-FileCopyrightText: Copyright contributors to the vLLM project", ] NO_COPYRIGHT_LINE = "vllm: no copyright" QUALITY_COMMAND = "quality" STYLE_COMMAND = "style" def parse_args(): """ Setup and parse command line arguments for using the script """ parser = argparse.ArgumentParser( description=( "Add vLLM copyright to the beginning of all " "files under the given glob patterns. " "Currently assumes Python files using '#' as the commenting prefix." ) ) subparsers = parser.add_subparsers(dest="command") quality_parser = subparsers.add_parser( QUALITY_COMMAND, description=( "Run check across the files in the given patterns and " "fail if any do not have a copyright in them" ), ) style_parser = subparsers.add_parser( STYLE_COMMAND, description=( "Add the copyright to any files in the given patterns if it is not present" ), ) for sub in [quality_parser, style_parser]: sub.add_argument( "patterns", type=str, default=[], nargs="+", help="the patterns to search through", ) return parser.parse_args() def quality(patterns: List[str]): """ Run a quality check across all files in the given glob patterns. This checks to make sure all matching files have the NM copyright present. If any do not, it will list them out and exit with an error. :param patterns: The glob file patterns to run quality check on """ check_files = _get_files(patterns) error_files = [] for file in check_files: if not _dont_copyright(file) and not _contains_copyright(file): print(f"would add copyright to {file}") error_files.append(file) if error_files: sys.exit( f"{len(error_files)} would be copyrighted, " f"{len(check_files) - len(error_files)} would be left unchanged." ) else: print(f"{len(check_files)} files have copyrights") def style(patterns: List[str]): """ Run a style application across all files in the given glob patterns. This checks to make sure all matching files have the NM copyright present. If any do not, it will append the copyright to above the file after any already contained headers such as shebang lines. :param patterns: The glob file patterns to run quality check on """ check_files = _get_files(patterns) copyrighted_files = [] for file in check_files: if not _dont_copyright(file) and not _contains_copyright(file): _add_copyright(file) print(f"copyrighted {file}") copyrighted_files.append(file) if copyrighted_files: print( f"{len(copyrighted_files)} file(s) copyrighted, " f"{len(check_files) - len(copyrighted_files)} files unchanged" ) else: print(f"{len(check_files)} files unchanged") def _get_files(patterns: List[str]) -> List[str]: files = [] for pattern in patterns: for file in glob.glob(pattern, recursive=True): files.append(os.path.abspath(os.path.expanduser(file))) files.sort() return files def _dont_copyright(file_path: str) -> bool: if file_path.endswith("compressed_tensors/version.py"): return True with open(file_path, "r") as file: content = file.read() try: content.index(NO_COPYRIGHT_LINE) return True except ValueError: return False def _contains_copyright(file_path: str) -> bool: with open(file_path, "r") as file: content = file.read() try: for line in COPYRIGHT_LINES: content.index(line) return True except ValueError: return False def _add_copyright(file_path: str): file_type = _file_type(file_path) if file_type == "unknown": raise ValueError( f"unsupported file_type given to be copyrighted at {file_path}" ) with open(file_path, "r+") as file: lines = file.readlines() header_info = _file_header_info(lines, file_type) inject_index = 0 if header_info.end_index > -1: # if there is already a header, we want to inject the copyright after it # additionally we'll need a new line between the prev header and copyright inject_index = header_info.end_index + 1 lines.insert(inject_index, "\n") inject_index += 1 # add the copyright at the inject index file_copyright = _file_copyright(file_type) lines.insert(inject_index, file_copyright) if not header_info.new_line_after: # if there wasn't a new line after the header, # add in a new line after to create space between the code and copyright inject_index += 1 lines.insert(inject_index, "\n") file.seek(0) file.writelines(lines) file.truncate() def _file_copyright(file_type: str) -> str: comment_formatting = _code_comment_formatting(file_type) lines = [] if comment_formatting.block_prefix: lines.append(comment_formatting.block_prefix) for line in COPYRIGHT_LINES: lines.append( f"{comment_formatting.line_prefix} {line}" if comment_formatting.line_prefix else line ) if comment_formatting.block_suffix: lines.append(comment_formatting.block_suffix) # make sure there is a new line after last line of the copyright lines.append("") return "\n".join(lines) _HeaderInfo = NamedTuple( "HeaderInfo", [ ("start_index", int), ("end_index", int), ("new_line_before", bool), ("new_line_after", bool), ], ) def _file_header_info(lines: List[str], file_type: str) -> _HeaderInfo: start_index = -1 end_index = -1 new_line_before = False new_line_after = False comment_formatting = _code_comment_formatting(file_type) prefix_found = False suffix_found = False for index, line in enumerate(lines): line = line.strip() if not line: # empty line, record the state of new lines before and after header if not prefix_found: new_line_before = True elif prefix_found and (suffix_found or not comment_formatting.block_suffix): new_line_after = True elif ( comment_formatting.block_prefix and line.startswith(comment_formatting.block_prefix) ) or ( not comment_formatting.block_prefix and line.startswith(comment_formatting.line_prefix) ): # start of header prefix_found = True start_index = index end_index = index suffix_found = comment_formatting.block_suffix and line.endswith( comment_formatting.block_suffix ) elif comment_formatting.block_suffix and line.endswith( comment_formatting.block_suffix ): # end of header suffix_found = True end_index = index elif prefix_found and comment_formatting.block_suffix and not suffix_found: # in the middle of the header, searching for the end # reset new_line_after in case there was a break in the header new_line_after = True else: # first non header line, break out break return _HeaderInfo(start_index, end_index, new_line_before, new_line_after) _CommentFormatting = NamedTuple( "CommentFormatting", [ ("line_prefix", str), ("block_prefix", str), ("block_suffix", str), ], ) def _code_comment_formatting(file_type: str) -> _CommentFormatting: if file_type == "python": return _CommentFormatting("#", "", "") elif file_type == "html" or file_type == "markdown": return _CommentFormatting("", "") elif file_type == "css" or file_type == "javascript": return _CommentFormatting("", "/*", "*/") elif file_type == "restructuredtext": return _CommentFormatting(" ", "..", "") raise ValueError(f"unsupported file_type given for code prefix suffix: {file_type}") def _file_type(file_path: str) -> str: if file_path.endswith(".py"): return "python" elif ( file_path.endswith(".js") or file_path.endswith(".jsx") or file_path.endswith(".ts") or file_path.endswith(".tsx") or file_path.endswith(".jss") ): return "javascript" elif file_path.endswith(".html"): return "html" elif file_path.endswith(".css"): return "css" elif file_path.endswith(".md"): return "markdown" elif file_path.endswith(".rst"): return "restructuredtext" return "unknown" def main(): args = parse_args() if args.command == QUALITY_COMMAND: quality(args.patterns) elif args.command == STYLE_COMMAND: style(args.patterns) else: raise ValueError(f"unknown command given: {args.command}") if __name__ == "__main__": main()