[
{
"path": ".gitignore",
"content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\ncover/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\n.pybuilder/\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n# For a library or package, you might want to ignore these files since the code is\n# intended to run in multiple environments; otherwise, check them in:\n# .python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# poetry\n# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.\n# This is especially recommended for binary packages to ensure reproducibility, and is more\n# commonly ignored for libraries.\n# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control\npoetry.lock\n\n# pdm\n# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.\n#pdm.lock\n# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it\n# in version control.\n# https://pdm.fming.dev/#use-with-ide\n.pdm.toml\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# pytype static type analyzer\n.pytype/\n\n# Cython debug symbols\ncython_debug/\n\n# PyCharm\n# JetBrains specific template is maintained in a separate JetBrains.gitignore that can\n# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore\n# and can be added to the global gitignore or merged into this file. For a more nuclear\n# option (not recommended) you can uncomment the following to ignore the entire idea folder.\n.idea/\n\n# VS Code\n.vscode/\n!.vscode/settings.json\n\n.DS_Store\n*.log\n*.pt\n.tmp/\nruns\nexps\nruns/\nexps/\nwandb\nwandb/\n"
},
{
"path": "LICENSE",
"content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [2024] Yujun Lin, Muyang Li, Zhekai Zhang, Haotian Tang, Shang Yang, Song Han\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License."
},
{
"path": "README.md",
"content": "\n![\"DeepCompressor](\"assets/deepcompressor.png\")
\n
\n\nModel Compression Toolbox for Large Language Models and Diffusion Models
\n\n\n \n ![\"Apache](\"https://img.shields.io/github/license/mit-han-lab/deepcompressor\")
\n \n \n \n
\n\n## News\n- [2025/02] ๐ [**QServe**](https://arxiv.org/abs/2405.04532) has been accepted to MLSys 2025!\n- [2025/01] ๐ [**SVDQuant**](https://arxiv.org/abs/2411.05007) has been accepted to ICLR 2025 (Spotlight)!\n- [2024/12] ๐ [**QServe**](https://github.com/mit-han-lab/qserve) has been integratedd into NVIDIA [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/llama)!\n- [2024/11] ๐ฅ Our latest **W4A4** diffusion model quantization work [**SVDQuant**](https://arxiv.org/abs/2411.05007) algorithm and [**Nunchaku**](https://github.com/mit-han-lab/nunchaku) system is publicly released! Check our [paper](http://arxiv.org/abs/2411.05007)!\n- [2024/05] ๐ฅ Our latest **W4A8KV4** LLM quantization work **QoQ** algorithm and **QServe** system is publicly released! **QoQ** is short for *quattuor-octล-quattuor* which is 4-8-4 in latin. Check our [paper](https://arxiv.org/abs/2405.04532)!\n\n## Key Features\n\n***DeepCompressor*** is an open source model compression toolbox for large language models and diffusion models based on PyTorch. DeepCompressor currently supports fake quantization with any integer and floating-point data type within 8 bits, e.g., INT8, INT4 and FP4_E2M1. Here are examples that implement the following algorithms.\n\n+ [Post-training quantization for large language models](/examples/llm/):\n + Weight-only Quantization\n + [AWQ (W4A16)](/examples/llm/configs/awq.yaml)\n + [GPTQ (W4A16)](/examples/llm/configs/gptq.yaml)\n + Weight-Activation Quantization\n + [SmoothQuant (W8A8)](/examples/llm/configs/smoothquant-static.yaml)\n + Weight-Activation and KV-Cache Quantization\n + [QoQ (W4A8KV4)](/examples/llm/)\n+ [Post-training quantization for diffusion models](/examples/diffusion/):\n + Weight-Activation Quantization\n + [SVDQuant (W4A4)](/examples/diffusion/)\n\nDeepCompressor also contains examples that integrate with other inference libraries.\n + [Deploy weight-only quantized LLMs with TinyChat](/examples/llm/)\n + [Deploy quantized LLMs with QServe]((/examples/llm/))\n + [Deploy quantized diffusion models with Nunchaku](/examples/diffusion/)\n\n## Installation\n\n### Install from Source\n\n1. Clone this repository and navigate to deepcompressor folder\n```\ngit clone https://github.com/mit-han-lab/deepcompressor\ncd deepcompressor\n```\n\n2. Install Package\n```\nconda env create -f environment.yml\npoetry install\n```\n\n## Highlights\n### SVDQuant: Absorbing Outliers by Low-Rank Components for 4-Bit Diffusion Models\n\n[[Website](https://hanlab.mit.edu/projects/svdquant)][[Paper](http://arxiv.org/abs/2411.05007)][[Nunchaku Inference System](https://github.com/mit-han-lab/nunchaku)]\n\nDiffusion models have been proven highly effective at generating high-quality images. However, as these models grow larger, they require significantly more memory and suffer from higher latency, posing substantial challenges for deployment. In this work, we aim to accelerate diffusion models by quantizing their weights and activations to 4 bits. At such an aggressive level, both weights and activations are highly sensitive, where conventional post-training quantization methods for large language models like smoothing become insufficient. To overcome this limitation, we propose **SVDQuant**, a new 4-bit quantization paradigm. Different from smoothing which redistributes outliers between weights and activations, our approach absorbs these outliers using a low-rank branch. We first consolidate the outliers by shifting them from activations to weights, then employ a high-precision low-rank branch to take in the weight outliers with Singular Value Decomposition (SVD). This process eases the quantization on both sides. However, naรฏvely running the low-rank branch independently incurs significant overhead due to extra data movement of activations, negating the quantization speedup. To address this, we co-design an inference engine **Nunchaku** that fuses the kernels of the low-rank branch into those of the low-bit branch to cut off redundant memory access. It can also seamlessly support off-the-shelf low-rank adapters (LoRAs) without the need for re-quantization. Extensive experiments on SDXL, PixArt-โ, and FLUX.1 validate the effectiveness of SVDQuant in preserving image quality. We reduce the memory usage for the 12B FLUX.1 models by 3.5ร, achieving 3.0ร speedup over the 4-bit weight-only quantized baseline on the 16GB laptop 4090 GPU, paving the way for more interactive applications on PCs.\n\n\n\n\n#### Quality Evaluation\n\nBelow is the quality and similarity evaluated with 5000 samples from MJHQ-30K dataset. IR means ImageReward. Our 4-bit results outperform other 4-bit baselines, effectively preserving the visual quality of 16-bit models.\n\n| Model | Precision | Method | FID ($\\downarrow$) | IR ($\\uparrow$) | LPIPS ($\\downarrow$) | PSNR( $\\uparrow$) |\n|----------------------------|-----------|-----------|--------------------|-----------------|----------------------|-------------------|\n| FLUX.1-dev (50 Steps) | BF16 | -- | 20.3 | 0.953 | -- | -- |\n| | W4A16 | NF4 | 20.6 | 0.910 | 0.272 | 19.5 |\n| | INT W4A4 | | 20.2 | 0.908 | 0.322 | 18.5 |\n| | INT W4A4 | SVDQuant | 19.9 | 0.935 | 0.223 | 21.0 |\n| | NVFP4 | | 20.3 | 0.961 | 0.345 | 16.3 |\n| | NVFP4 | SVDQuant | 20.3 | 0.945 | 0.205 | 21.5 |\n| FLUX.1-schnell (4 Steps) | BF16 | -- | 19.2 | 0.938 | -- | -- |\n| | W4A16 | NF4 | 18.9 | 0.943 | 0.257 | 18.2 |\n| | INT W4A4 | | 18.1 | 0.962 | 0.345 | 16.3 |\n| | INT W4A4 | SVDQuant | 18.3 | 0.951 | 0.257 | 18.3 |\n| | NVFP4 | | 19.0 | 0.952 | 0.276 | 17.6 |\n| | NVFP4 | SVDQuant | 18.9 | 0.966 | 0.228 | 19.0 |\n| SANA-1.6b (20 Steps) | BF16 | -- | 20.6 | 0.952 | -- | -- |\n| | INT W4A4 | | 20.5 | 0.894 | 0.339 | 15.3 |\n| | INT W4A4 | SVDQuant | 19.3 | 0.935 | 0.220 | 17.8 |\n| | NVFP4 | | 19.7 | 0.929 | 0.236 | 17.4 |\n| | NVFP4 | SVDQuant | 20.2 | 0.941 | 0.176 | 19.0 |\n| PixArt-Sigma (20 Steps) | FP16 | -- | 16.6 | 0.944 | -- | -- |\n| | INT W4A8 | ViDiT-Q | 37.3 | 0.573 | 0.611 | 12.0 |\n| | INT W4A4 | SVDQuant | 19.2 | 0.878 | 0.323 | 17.6 |\n| | NVFP4 | | 31.8 | 0.660 | 0.517 | 14.8 |\n| | NVFP4 | SVDQuant | 16.6 | 0.940 | 0.271 | 18.5 |\n\n### QServe: W4A8KV4 Quantization for Efficient LLM Serving\n\n[[Website](https://hanlab.mit.edu/projects/qserve)][[Paper](https://arxiv.org/abs/2405.04532)][[QoQ Algorithm Code](/examples/llm)][[QServe GPU System](https://github.com/mit-han-lab/qserve)]\n\nQuantization can accelerate large language model (LLM) inference. Going beyond INT8 quantization, the research community is actively exploring even lower precision, such as INT4. Nonetheless, state-of-the-art INT4 quantization techniques only accelerate low-batch, edge LLM inference, failing to deliver performance gains in large-batch, cloud-based LLM serving. We uncover a critical issue: existing INT4 quantization methods suffer from significant runtime overhead (20-90%) when **dequantizing either weights or partial sums** on GPUs. To address this challenge, we introduce **QoQ**, a W4A8KV4 quantization algorithm with 4-bit weight, 8-bit activation, and 4-bit KV cache. QoQ stands for **quattuor-octo-quattuor**, which represents 4-8-4 in Latin. QoQ is implemented by the **QServe** inference library that achieves measured speedup. The key insight driving QServe is that the efficiency of LLM serving on GPUs is critically influenced by **operations on low-throughput CUDA cores**. Building upon this insight, in QoQ algorithm, we introduce progressive quantization that can allow low dequantization overhead in W4A8 GEMM. Additionally, we develop SmoothAttention to effectively mitigate the accuracy degradation incurred by 4-bit KV quantization. In the QServe system, we perform compute-aware weight reordering and take advantage of register-level parallelism to reduce dequantization latency. We also make fused attention memory-bound, harnessing the performance gain brought by KV4 quantization. As a result, QServe improves the maximum achievable serving throughput of Llama-3-8B by **1.2ร** on A100, **1.4ร** on L40S; and Qwen1.5-72B by **2.4ร** on A100, **3.5ร** on L40S, compared to TensorRT-LLM.\n\n\n\n\n\n#### Perplexity Evaluation\n\nBelow is the WikiText2 perplexity evaluated with 2048 sequence length. The lower is the better.\n\n| Methods | Precision | Llama-3.1 70B | Llama-3.1 8B | Llama-3 70B | Llama-3 8B | Llama-2 7B | Llama-2 13B | Llama-2 70B | Llama 7B | Llama 13B | Llama 30B | Mistral 7B | Yi 34B |\n|-------------|--------------|---------------|--------------|-------------| ------------|------------|-------------|-------------|----------|-----------|-----------|------------|--------|\n| FP16 | | 2.81 | 6.24 | 2.85 | 6.14 | 5.47 | 4.88 | 3.32 | 5.68 | 5.09 | 4.10 | 5.25 | 4.60 |\n| SmoothQuant | W8A8 | 3.23 | 6.38 | 3.14 | 6.28 | 5.54 | 4.95 | 3.36 | 5.73 | 5.13 | 4.23 | 5.29 | 4.69 |\n| GPTQ-R | W4A16 g128 | 3.46 | 6.64 | 3.42 | 6.56 | 5.63 | 4.99 | 3.43 | 5.83 | 5.20 | 4.22 | 5.39 | 4.68 |\n| AWQ | W4A16 g128 | 3.22 | 6.60 | 3.20 | 6.54 | 5.60 | 4.97 | 3.41 | 5.78 | 5.19 | 4.21 | 5.37 | 4.67 |\n| QuaRot | W4A4 | 5.97 | 8.32 | 6.75 | 8.33 | 6.19 | 5.45 | 3.83 | 6.34 | 5.58 | 4.64 | 5.77 | - |\n| SpinQuant | W4A4 | 4.80 | 7.42 | 6.27 | 7.37 | 5.96 | 5.24 | 3.71 | 6.14 | 5.39 | 4.56 | - | - |\n| Atom | W4A4 g128 | - | - | 4.33 | 7.78 | 6.12 | 5.31 | 3.73 | 6.25 | 5.52 | 4.61 | 5.76 | 4.97 |\n| QoQ | W4A8KV4 | 3.68 | 6.87 | 3.65 | 6.81 | 5.75 | 5.11 | 3.50 | 5.92 | 5.27 | 4.31 | 5.44 | 4.73 |\n| QoQ | W4A8KV4 g128 | 3.51 | 6.77 | 3.50 | 6.70 | 5.67 | 5.06 | 3.46 | 5.88 | 5.23 | 4.27 | 5.41 | 4.73 |\n\n\\* SmoothQuant is evaluated with per-tensor static KV cache quantization.\n\n#### Efficiency Benchmarks\n\nWhen serving the large language models Llama-3-8B and Qwen1.5-72B on L40S and A100 GPUs, QServe demonstrates superior performance, achieving **1.2x-1.4x higher throughput** compared to the leading industry solution, TensorRT-LLM, for Llama-3-8B, and a **2.4x-3.5x higher throughput** for Qwen1.5-72B.\n\nSee more about benchmarking setting in [QServe GPU Inference System](https://github.com/mit-han-lab/qserve).\n| L40S (48G) | Llama-3-8B | Llama-2-7B | Mistral-7B | Llama-2-13B | Llama-30B | Yi-34B | Llama-2-70B | Qwen-1.5-72B |\n|----------------------|------------|------------|------------|-------------|-----------|-----------|-------------|--------------|\n| TRT-LLM-FP16 | 1326 | 444 | 1566 | 92 | OOM | OOM | OOM | OOM |\n| TRT-LLM-W4A16 | 1431 | 681 | 1457 | 368 | 148 | 313 | 119 | 17 |\n| TRT-LLM-W8A8 | 2634 | 1271 | 2569 | 440 | 123 | 364 | OOM | OOM |\n| Atom-W4A4 | -- | 2120 | -- | -- | -- | -- | -- | -- |\n| QuaRot-W4A4 | -- | 805 | -- | 413 | 133 | -- | -- | 15 |\n| QServe-W4A8KV4 | **3656** | **2394** | **3774** | **1327** | **504** | **869** | **286** | **59** |\n| Throughput Increase* | **1.39x** | **1.13x** | **1.47x** | **3.02x** | **3.41x** | **2.39x** | **2.40x** | **3.47x** |\n\n| A100 (80G) | Llama-3-8B | Llama-2-7B | Mistral-7B | Llama-2-13B | Llama-30B | Yi-34B | Llama-2-70B | Qwen-1.5-72B |\n|----------------------|------------| -----------|------------|-------------|-----------|-----------|-------------|--------------|\n| TRT-LLM-FP16 | 2503 | 1549 | 2371 | 488 | 80 | 145 | OOM | OOM |\n| TRT-LLM-W4A16 | 2370 | 1549 | 2403 | 871 | 352 | 569 | 358 | 143 |\n| TRT-LLM-W8A8 | 2396 | 2334 | 2427 | 1277 | 361 | 649 | 235 | 53 |\n| Atom-W4A4 | -- | 1160 | -- | -- | -- | -- | -- | -- |\n| QuaRot-W4A4 | -- | 1370 | -- | 289 | 267 | -- | -- | 68 |\n| QServe-W4A8KV4 | **3005** | **2908** | **2970** | **1741** | **749** | **803** | **419** | **340** |\n| Throughput Increase* | **1.20x** | **1.25x** | **1.22x** | **1.36x** | **2.07x** | **1.23x** | **1.17x** | **2.38x** |\n\nThe absolute token generation throughputs of QServe and baseline systems (Unit: tokens/second. `--` means unsupported). All experiments were conducted under the same device memory budget. Throughput increase of QServe is calculated with regard to the best baseline in each column.\n\n## Reference\n\nIf you find `deepcompressor` useful or relevant to your research, please kindly cite our paper:\n\n```bibtex\n@article{lin2024qserve,\n title={QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving},\n author={Lin*, Yujun and Tang*, Haotian and Yang*, Shang and Zhang, Zhekai and Xiao, Guangxuan and Gan, Chuang and Han, Song},\n journal={arXiv preprint arXiv:2405.04532},\n year={2024}\n}\n\n@article{\n li2024svdquant,\n title={SVDQuant: Absorbing Outliers by Low-Rank Components for 4-Bit Diffusion Models},\n author={Li*, Muyang and Lin*, Yujun and Zhang*, Zhekai and Cai, Tianle and Li, Xiuyu and Guo, Junxian and Xie, Enze and Meng, Chenlin and Zhu, Jun-Yan and Han, Song},\n journal={arXiv preprint arXiv:2411.05007},\n year={2024}\n}\n```\n\n## Related Projects\n\nThe following projects are highly related to QServe. Our group has developed full-stack application-algorithm-system-hardware support for efficient large models, receiving **9k+ GitHub stars** and **over 1M Huggingface community downloads**.\n\nYou are also welcome to check out [MIT HAN Lab](https://hanlab.mit.edu) for other exciting projects on **Efficient Generative AI**!\n\n- [**System**] [QServe: W4A8KV4 Quantization for Efficient LLM Serving](https://github.com/mit-han-lab/qserve)\n\n- [**System**] [TinyChat: Efficient and Lightweight Chatbot with AWQ](https://github.com/mit-han-lab/llm-awq/tree/main/tinychat)\n\n- [**Application**] [VILA: On Pretraining of Visual-Language Models](https://github.com/Efficient-Large-Model/VILA)\n\n- [**Algorithm**] [AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration](https://github.com/mit-han-lab/llm-awq)\n\n- [**Algorithm**] [SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models](https://github.com/mit-han-lab/smoothquant)\n\n- [**Algorithm**] [DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models](https://github.com/mit-han-lab/distrifuser)\n\n- [**Hardware**] [SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning](https://arxiv.org/abs/2012.09852)\n\n\n## Acknowledgments\n\nDeepCompressor is inspired by many open-source libraries, including (but not limited to) [GPTQ](https://arxiv.org/abs/2210.17323), [QuaRot](https://arxiv.org/abs/2404.00456) and [Atom](https://arxiv.org/abs/2310.19102). "
},
{
"path": "assets/diffusion/.gitkeep",
"content": ""
},
{
"path": "assets/llm/.gitkeep",
"content": ""
},
{
"path": "deepcompressor/__init__.py",
"content": "from .version import __version__ # noqa: F401\n"
},
{
"path": "deepcompressor/app/__init__.py",
"content": ""
},
{
"path": "deepcompressor/app/diffusion/__init__.py",
"content": ""
},
{
"path": "deepcompressor/app/diffusion/cache/__init__.py",
"content": "from .config import DiffusionPtqCacheConfig, DiffusionQuantCacheConfig\n"
},
{
"path": "deepcompressor/app/diffusion/cache/config.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"LLM quantization cache configuration.\"\"\"\n\nimport functools\nimport re\nimport typing as tp\nfrom dataclasses import dataclass, field\n\nfrom omniconfig import configclass\n\nfrom deepcompressor.utils.config.path import BasePathConfig\n\nfrom ..nn.struct import DiffusionModelStruct\n\n__all__ = [\"DiffusionQuantCacheConfig\", \"DiffusionPtqCacheConfig\"]\n\n\n@dataclass\nclass DiffusionQuantCacheConfig(BasePathConfig):\n \"\"\"Denoising diffusion model quantization cache path.\n\n Args:\n smooth (`str`, *optional*, default=`\"\"`):\n The smoothing scales cache path.\n branch (`str`, *optional*, default=`\"\"`):\n The low-rank branches cache path.\n wgts (`str`, *optional*, default=`\"\"`):\n The weight quantizers state dict cache path.\n acts (`str`, *optional*, default=`\"\"`):\n The activation quantizers state dict cache path\n \"\"\"\n\n smooth: str = \"\"\n branch: str = \"\"\n wgts: str = \"\"\n acts: str = \"\"\n\n @staticmethod\n def simplify_path(path: str, key_map: dict[str, set[str]]) -> str:\n \"\"\"Simplify the cache path.\"\"\"\n to_replace = {}\n # we first extract all the parts matching the pattern \"(skip|include).\\[[a-zA-Z0-9_\\+]+\\]\"\n for part in re.finditer(r\"(skip|include)\\.\\[[a-zA-Z0-9_\\+]+\\]\", path):\n # remove the \"skip.\" or \"include.\" prefix\n part = part.group(0)\n if part[0] == \"s\":\n prefix, keys = part[:4], part[6:-1]\n else:\n prefix, keys = part[:7], part[9:-1]\n # simplify the keys\n keys = \"+\".join(\n (\n \"\".join((s[0] for s in x.split(\"_\")))\n for x in DiffusionModelStruct._simplify_keys(keys.split(\"+\"), key_map=key_map)\n )\n )\n to_replace[part] = f\"{prefix}.[{keys}]\"\n # we then replace the parts\n for key, value in to_replace.items():\n path = path.replace(key, value)\n return path\n\n def simplify(self, key_map: dict[str, set[str]]) -> tp.Self:\n \"\"\"Simplify the cache paths.\"\"\"\n return self.apply(functools.partial(self.simplify_path, key_map=key_map))\n\n\n@configclass\n@dataclass\nclass DiffusionPtqCacheConfig:\n root: str\n dirpath: DiffusionQuantCacheConfig = field(init=False)\n path: DiffusionQuantCacheConfig = field(init=False)\n"
},
{
"path": "deepcompressor/app/diffusion/config.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Top-level config of post-training quantization for a diffusion model.\"\"\"\n\nimport os\nfrom dataclasses import dataclass, field\n\nimport diffusers.training_utils\nimport omniconfig\nimport torch\nfrom omniconfig import ConfigParser, configclass\n\nfrom deepcompressor.app.llm.config import LlmCacheConfig, LlmQuantConfig\nfrom deepcompressor.data.utils import ScaleUtils\nfrom deepcompressor.utils.config.output import OutputConfig\n\nfrom .cache import DiffusionPtqCacheConfig, DiffusionQuantCacheConfig\nfrom .eval import DiffusionEvalConfig\nfrom .nn.struct import DiffusionModelStruct\nfrom .pipeline import DiffusionPipelineConfig\nfrom .quant import DiffusionQuantConfig\n\n__all__ = [\n \"DiffusionPtqRunConfig\",\n \"DiffusionPtqCacheConfig\",\n \"DiffusionQuantCacheConfig\",\n \"DiffusionEvalConfig\",\n \"DiffusionPipelineConfig\",\n \"DiffusionQuantConfig\",\n]\n\n\n@configclass\n@dataclass\nclass DiffusionPtqRunConfig:\n \"\"\"Top-level config of post-training quantization for a diffusion model.\n\n Args:\n cache (`DiffusionPtqCacheConfig`):\n The cache configuration.\n output (`OutputConfig`):\n The output directory configuration.\n pipeline (`DiffusionPipelineConfig`):\n The diffusion pipeline configuration\n eval (`DiffusionEvalConfig`):\n The evaluation configuration.\n quant (`DiffusionQuantConfig`):\n The post-training quantization configuration.\n seed (`int`, *optional*, defaults to `12345`):\n The seed for reproducibility.\n skip_gen (`bool`, *optional*, defaults to `False`):\n Whether to skip generation.\n skip_eval (`bool`, *optional*, defaults to `False`):\n Whether to skip evaluation.\n load_model (`str`, *optional*, defaults to `\"\"`):\n Directory path to load the model checkpoint.\n save_model (`str`, *optional*, defaults to `\"\"`):\n Directory path to save the model checkpoint.\n copy_on_save (`bool`, *optional*, defaults to `False`):\n Whether to copy the quantization cache on save.\n \"\"\"\n\n cache: DiffusionPtqCacheConfig | None\n output: OutputConfig\n pipeline: DiffusionPipelineConfig\n eval: DiffusionEvalConfig\n quant: DiffusionQuantConfig = field(metadata={omniconfig.ARGPARSE_KWARGS: {\"prefix\": \"\"}})\n text: LlmQuantConfig | None = None\n text_cache: LlmCacheConfig = field(default_factory=LlmCacheConfig)\n seed: int = 12345\n skip_gen: bool = False\n skip_eval: bool = False\n load_from: str = \"\"\n save_model: str = \"\"\n copy_on_save: bool = False\n\n def __post_init__(self):\n # region set text encoder quanatization scale default dtype\n if self.text is not None and self.text.enabled_wgts:\n self.text.wgts.scale_dtypes = tuple(\n ScaleUtils.infer_scale_dtypes(self.text.wgts.scale_dtypes, default_dtype=self.pipeline.dtype)\n )\n if self.text is not None and self.text.enabled_ipts:\n self.text.ipts.scale_dtypes = tuple(\n ScaleUtils.infer_scale_dtypes(self.text.ipts.scale_dtypes, default_dtype=self.pipeline.dtype)\n )\n if self.text is not None and self.text.enabled_opts:\n self.text.opts.scale_dtypes = tuple(\n ScaleUtils.infer_scale_dtypes(self.text.opts.scale_dtypes, default_dtype=self.pipeline.dtype)\n )\n # endregion\n self.eval.num_gpus = min(torch.cuda.device_count(), self.eval.num_gpus)\n if self.eval.batch_size_per_gpu is None:\n self.eval.batch_size_per_gpu = max(1, self.eval.batch_size // self.eval.num_gpus)\n self.eval.batch_size = self.eval.batch_size_per_gpu * self.eval.num_gpus\n else:\n self.eval.batch_size = self.eval.batch_size_per_gpu * self.eval.num_gpus\n # region setup calib dataset path\n self.quant.calib.path = self.quant.calib.path.format(\n dtype=self.pipeline.dtype,\n family=self.pipeline.family,\n model=self.pipeline.name,\n protocol=self.eval.protocol,\n data=self.quant.calib.data,\n )\n if self.quant.calib.path:\n self.quant.calib.path = os.path.abspath(os.path.expanduser(self.quant.calib.path))\n # endregion\n # region setup eval reference root\n self.eval.ref_root = self.eval.ref_root.format(\n dtype=self.pipeline.dtype,\n family=self.pipeline.family,\n model=self.pipeline.name,\n protocol=self.eval.protocol,\n )\n if self.eval.ref_root:\n self.eval.ref_root = os.path.abspath(os.path.expanduser(self.eval.ref_root))\n # endregion\n # region setup cache directory\n if self.cache is not None:\n if self.quant.enabled_wgts or self.quant.enabled_ipts or self.quant.enabled_opts:\n self.cache.dirpath = self.quant.generate_cache_dirpath(\n root=self.cache.root, shift=self.pipeline.shift_activations, default_dtype=self.pipeline.dtype\n )\n self.cache.path = self.cache.dirpath.clone().add_children(f\"{self.pipeline.name}.pt\")\n else:\n self.cache.dirpath = self.cache.path = None\n if self.text is not None and self.text.is_enabled():\n if not self.text_cache.root:\n self.text_cache.root = os.path.join(self.cache.root, \"diffusion\")\n self.text_cache.dirpath = self.text.generate_cache_dirpath(root=self.text_cache.root, seed=self.seed)\n self.text_cache.path = self.text_cache.dirpath.clone().add_children(f\"{self.pipeline.name}.pt\")\n # endregion\n # region setup output directory\n if self.output.dirname == \"reference\":\n assert self.eval.ref_root\n self.output.job = f\"run-{self.eval.num_samples}\"\n self.output.dirpath = self.eval.ref_root\n self.eval.ref_root = \"\"\n self.eval.gen_root = \"{output}\"\n else:\n if self.output.dirname == \"default\":\n self.output.dirname = self.generate_default_dirname()\n calib_dirname = self.quant.generate_calib_dirname() or \"-\"\n self.output.dirpath = os.path.join(\n self.output.root,\n \"diffusion\",\n self.pipeline.family,\n self.pipeline.name,\n *self.quant.generate_dirnames(default_dtype=self.pipeline.dtype)[:-1],\n calib_dirname,\n self.output.dirname,\n )\n if (self.eval.chunk_start > 0 or self.eval.chunk_step > 1) and not self.eval.chunk_only:\n self.output.job += f\".c{self.eval.chunk_start}.{self.eval.chunk_step}\"\n # endregion\n diffusers.training_utils.set_seed(self.seed)\n\n def generate_default_dirname(self) -> str:\n name = \"-shift\" if self.pipeline.shift_activations else \"\"\n if self.quant.is_enabled():\n name += f\"-{self.quant.generate_default_dirname()}\"\n if self.text is not None and self.text.is_enabled():\n name += f\"-text-{self.text.generate_default_dirname()}\"\n size_name = \"\"\n if self.eval.height:\n size_name += f\".h{self.eval.height}\"\n if self.eval.width:\n size_name += f\".w{self.eval.width}\"\n if size_name:\n name += f\"-{size_name[1:]}\"\n sampling_name = \"\"\n if self.eval.num_steps is not None:\n sampling_name += f\".t{self.eval.num_steps}\"\n if self.eval.guidance_scale is not None:\n sampling_name += f\".g{self.eval.guidance_scale}\"\n if sampling_name:\n name += f\"-{sampling_name[1:]}\"\n if self.eval.num_samples != -1:\n name += f\"-s{self.eval.num_samples}\"\n if self.eval.chunk_only:\n name += f\".c{self.eval.chunk_start}.{self.eval.chunk_step}\"\n assert name[0] == \"-\"\n return name[1:]\n\n @classmethod\n def get_parser(cls) -> ConfigParser:\n \"\"\"Get a parser for post-training quantization of a diffusion model.\n\n Returns:\n `ConfigParser`:\n A parser for post-training quantization of a diffusion model.\n \"\"\"\n parser = ConfigParser(\"Diffusion Run configuration\")\n DiffusionQuantConfig.set_key_map(DiffusionModelStruct._get_default_key_map())\n parser.add_config(cls)\n return parser\n"
},
{
"path": "deepcompressor/app/diffusion/dataset/__init__.py",
"content": "# -*- coding: utf-8 -*-\n\nfrom .base import DiffusionDataset\nfrom .calib import DiffusionCalibCacheLoader, DiffusionCalibCacheLoaderConfig\n"
},
{
"path": "deepcompressor/app/diffusion/dataset/base.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Dataset for diffusion models.\"\"\"\n\nimport os\nimport random\nimport typing as tp\n\nimport numpy as np\nimport torch\nimport torch.utils.data\nfrom torch.nn import functional as F\n\nfrom deepcompressor.utils.common import tree_collate, tree_map\n\n__all__ = [\"DiffusionDataset\"]\n\n\nclass DiffusionDataset(torch.utils.data.Dataset):\n path: str\n filenames: list[str]\n filepaths: list[str]\n\n def __init__(self, path: str, num_samples: int = -1, seed: int = 0, ext: str = \".npy\") -> None:\n if os.path.exists(path):\n self.path = path\n if \"caches\" in os.listdir(path):\n path = os.path.join(path, \"caches\")\n filenames = [f for f in sorted(os.listdir(path)) if f.endswith(ext)]\n if num_samples > 0 and num_samples < len(filenames):\n random.Random(seed).shuffle(filenames)\n filenames = filenames[:num_samples]\n filenames = sorted(filenames)\n self.filenames = filenames\n self.filepaths = [os.path.join(path, f) for f in filenames]\n else:\n raise ValueError(f\"Invalid data path: {path}\")\n\n def __len__(self) -> int:\n return len(self.filepaths)\n\n def __getitem__(self, idx) -> dict[str, tp.Any]:\n data = np.load(self.filepaths[idx], allow_pickle=True).item()\n if isinstance(data[\"input_args\"][0], str):\n name = data[\"input_args\"][0]\n latent = np.load(os.path.join(self.path, \"latents\", name))\n data[\"input_args\"][0] = latent\n if isinstance(data[\"input_kwargs\"][\"encoder_hidden_states\"], str):\n name = data[\"input_kwargs\"][\"encoder_hidden_states\"]\n text_emb = np.load(os.path.join(self.path, \"text_embs\", name))\n data[\"input_kwargs\"][\"encoder_hidden_states\"] = text_emb\n data = tree_map(lambda x: torch.from_numpy(x), data)\n\n # Pad encoder_hidden_states to 300 for pixart\n if \"encoder_attention_mask\" in data[\"input_kwargs\"]:\n encoder_attention_mask = data[\"input_kwargs\"][\"encoder_attention_mask\"]\n encoder_hidden_states = data[\"input_kwargs\"][\"encoder_hidden_states\"]\n encoder_hidden_states = F.pad(\n encoder_hidden_states,\n (0, 0, 0, encoder_attention_mask.shape[1] - encoder_hidden_states.shape[1]),\n )\n data[\"input_kwargs\"][\"encoder_hidden_states\"] = encoder_hidden_states\n\n return data\n\n def build_loader(self, **kwargs) -> torch.utils.data.DataLoader:\n return torch.utils.data.DataLoader(self, collate_fn=tree_collate, **kwargs)\n"
},
{
"path": "deepcompressor/app/diffusion/dataset/calib.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Calibration dataset for diffusion models.\"\"\"\n\nimport random\nimport typing as tp\nfrom collections import OrderedDict\nfrom dataclasses import MISSING, dataclass\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.data\nfrom diffusers.models.attention import JointTransformerBlock\nfrom diffusers.models.attention_processor import Attention\nfrom diffusers.models.transformers.transformer_flux import (\n FluxSingleTransformerBlock,\n FluxTransformerBlock,\n)\nfrom omniconfig import configclass\n\nfrom deepcompressor.data.cache import (\n IOTensorsCache,\n ModuleForwardInput,\n TensorCache,\n TensorsCache,\n)\nfrom deepcompressor.data.utils.reshape import AttentionInputReshapeFn, LinearReshapeFn\nfrom deepcompressor.dataset.action import CacheAction, ConcatCacheAction\nfrom deepcompressor.dataset.cache import BaseCalibCacheLoader\nfrom deepcompressor.dataset.config import BaseDataLoaderConfig\n\nfrom ..nn.struct import DiffusionBlockStruct, DiffusionModelStruct\nfrom .base import DiffusionDataset\n\n__all__ = [\n \"DiffusionCalibCacheLoaderConfig\",\n \"DiffusionCalibDataset\",\n \"DiffusionConcatCacheAction\",\n \"DiffusionCalibCacheLoader\",\n]\n\n\n@configclass\n@dataclass(kw_only=True)\nclass DiffusionCalibCacheLoaderConfig(BaseDataLoaderConfig):\n \"\"\"Configuration for collecting calibration dataset for quantization.\n\n Args:\n data (`str`):\n Dataset name.\n num_samples (`int`):\n Number of dataset samples.\n batch_size (`int`):\n Batch size when loading dataset.\n path (`str`):\n Path to the dataset directory.\n num_workers (`int`):\n Number of workers for data loading.\n \"\"\"\n\n path: str\n num_workers: int = 8\n\n def build_dataset(self) -> \"DiffusionCalibDataset\":\n \"\"\"Build the calibration dataset.\"\"\"\n return DiffusionCalibDataset(self.path, num_samples=self.num_samples)\n\n def build_loader(self) -> \"DiffusionCalibCacheLoader\":\n \"\"\"Build the data loader.\"\"\"\n return DiffusionCalibCacheLoader(self)\n\n\nclass DiffusionCalibDataset(DiffusionDataset):\n data: list[dict[str, tp.Any]]\n\n def __init__(self, path: str, num_samples: int = -1, seed: int = 0) -> None:\n super().__init__(path, num_samples=num_samples, seed=seed, ext=\".pt\")\n data = [torch.load(path) for path in self.filepaths]\n random.Random(seed).shuffle(data)\n self.data = data\n\n def __len__(self) -> int:\n return len(self.data)\n\n def __getitem__(self, idx) -> dict[str, tp.Any]:\n return self.data[idx]\n\n\nclass DiffusionConcatCacheAction(ConcatCacheAction):\n def info(\n self,\n name: str,\n module: nn.Module,\n tensors: dict[int | str, torch.Tensor],\n cache: TensorsCache,\n ) -> None:\n \"\"\"Update cache information.\n\n Args:\n name (`str`):\n Module name.\n module (`nn.Module`):\n Module.\n tensors (`dict[int | str, torch.Tensor]`):\n Tensors to cache.\n cache (`TensorsCache`):\n Cache.\n \"\"\"\n if isinstance(module, Attention):\n encoder_hidden_states = tensors.get(\"encoder_hidden_states\", None)\n if encoder_hidden_states is None:\n tensors.pop(\"encoder_hidden_states\", None)\n cache.tensors.pop(\"encoder_hidden_states\", None)\n else:\n encoder_hidden_states_cache = cache.tensors[\"encoder_hidden_states\"]\n encoder_channels_dim = 1 if encoder_hidden_states.dim() == 4 else -1\n if encoder_hidden_states_cache.channels_dim is None:\n encoder_hidden_states_cache.channels_dim = encoder_channels_dim\n if encoder_channels_dim == -1:\n encoder_hidden_states_cache.reshape = LinearReshapeFn()\n else:\n encoder_hidden_states_cache.reshape = AttentionInputReshapeFn(encoder_channels_dim)\n else:\n assert encoder_hidden_states_cache.channels_dim == encoder_channels_dim\n hidden_states, hidden_states_cache = tensors[\"hidden_states\"], cache.tensors[\"hidden_states\"]\n channels_dim = 1 if hidden_states.dim() == 4 else -1\n if hidden_states_cache.channels_dim is None:\n hidden_states_cache.channels_dim = channels_dim\n if channels_dim == -1:\n hidden_states_cache.reshape = LinearReshapeFn()\n else:\n hidden_states_cache.reshape = AttentionInputReshapeFn(channels_dim)\n else:\n assert hidden_states_cache.channels_dim == channels_dim\n return super().info(name, module, tensors, cache)\n\n\nclass DiffusionCalibCacheLoader(BaseCalibCacheLoader):\n config: DiffusionCalibCacheLoaderConfig\n dataset: DiffusionCalibDataset\n\n def __init__(self, config: DiffusionCalibCacheLoaderConfig) -> None:\n \"\"\"Initialize the cache for the diffusion calibration dataset.\n\n Args:\n config (`DiffusionCalibCacheLoaderConfig`):\n Configuration for the calibration cache loader.\n \"\"\"\n super().__init__(dataset=config.build_dataset(), batch_size=config.batch_size)\n self.batch_size = min(config.batch_size, len(self.dataset))\n self.config = config\n\n def _init_cache(self, name: str, module: nn.Module) -> IOTensorsCache:\n \"\"\"Initialize cache.\n\n Args:\n name (`str`):\n Module name.\n module (`nn.Module`):\n Module.\n\n Returns:\n `IOTensorsCache`:\n Cache for inputs and outputs.\n \"\"\"\n if isinstance(module, FluxSingleTransformerBlock):\n return IOTensorsCache(\n inputs=TensorsCache(\n OrderedDict(\n hidden_states=TensorCache(channels_dim=-1, reshape=LinearReshapeFn()),\n temb=TensorCache(channels_dim=1, reshape=LinearReshapeFn()),\n )\n ),\n outputs=TensorCache(channels_dim=-1, reshape=LinearReshapeFn()),\n )\n elif isinstance(module, Attention):\n return IOTensorsCache(\n inputs=TensorsCache(\n OrderedDict(\n hidden_states=TensorCache(channels_dim=None, reshape=None),\n encoder_hidden_states=TensorCache(channels_dim=None, reshape=None),\n ),\n ),\n outputs=TensorCache(channels_dim=None, reshape=None),\n )\n else:\n return super()._init_cache(name, module)\n\n def iter_samples(self) -> tp.Generator[ModuleForwardInput, None, None]:\n dataloader = self.dataset.build_loader(\n batch_size=self.batch_size, shuffle=False, drop_last=True, num_workers=self.config.num_workers\n )\n for data in dataloader:\n yield ModuleForwardInput(args=data[\"input_args\"], kwargs=data[\"input_kwargs\"])\n\n def _convert_layer_inputs(\n self, m: nn.Module, args: tuple[tp.Any, ...], kwargs: dict[str, tp.Any], save_all: bool = False\n ) -> ModuleForwardInput:\n \"\"\"Convert layer inputs to module forward input.\n\n Args:\n m (`nn.Module`):\n Layer.\n args (`tuple[Any, ...]`):\n Layer input arguments.\n kwargs (`dict[str, Any]`):\n Layer input keyword arguments.\n save_all (`bool`, *optional*, defaults to `False`):\n Whether to save all inputs.\n\n Returns:\n `ModuleForwardInput`:\n Module forward input.\n \"\"\"\n kwargs = {k: v for k, v in kwargs.items()} # noqa: C416\n if \"res_hidden_states_tuple\" in kwargs:\n kwargs[\"res_hidden_states_tuple\"] = None\n if \"hidden_states\" in kwargs:\n hidden_states = kwargs.pop(\"hidden_states\")\n assert len(args) == 0, f\"Invalid args: {args}\"\n else:\n hidden_states = args[0]\n if isinstance(m, (FluxTransformerBlock, JointTransformerBlock)):\n if \"encoder_hidden_states\" in kwargs:\n encoder_hidden_states = kwargs.pop(\"encoder_hidden_states\")\n else:\n encoder_hidden_states = args[1]\n return ModuleForwardInput(\n args=[\n hidden_states.detach().cpu() if save_all else MISSING,\n encoder_hidden_states.detach().cpu() if save_all else MISSING,\n ],\n kwargs=kwargs,\n )\n else:\n return ModuleForwardInput(\n args=[hidden_states.detach().cpu() if save_all else MISSING, *args[1:]], kwargs=kwargs\n )\n\n def _convert_layer_outputs(self, m: nn.Module, outputs: tp.Any) -> dict[str | int, tp.Any]:\n \"\"\"Convert layer outputs to dictionary for updating the next layer inputs.\n\n Args:\n m (`nn.Module`):\n Layer.\n outputs (`Any`):\n Layer outputs.\n\n Returns:\n `dict[str | int, Any]`:\n Dictionary for updating the next layer inputs.\n \"\"\"\n if isinstance(m, (FluxTransformerBlock, JointTransformerBlock)):\n assert isinstance(outputs, tuple) and len(outputs) == 2\n encoder_hidden_states, hidden_states = outputs\n return {0: hidden_states.detach().cpu(), 1: encoder_hidden_states.detach().cpu()}\n else:\n return super()._convert_layer_outputs(m, outputs)\n\n def iter_layer_activations( # noqa: C901\n self,\n model: nn.Module | DiffusionModelStruct,\n *args,\n needs_inputs_fn: tp.Callable[[str, nn.Module], bool],\n needs_outputs_fn: tp.Callable[[str, nn.Module], bool] | None = None,\n action: CacheAction | None = None,\n skip_pre_modules: bool = True,\n skip_post_modules: bool = True,\n **kwargs,\n ) -> tp.Generator[\n tuple[\n str,\n tuple[\n DiffusionBlockStruct | nn.Module,\n dict[str, IOTensorsCache],\n dict[str, tp.Any],\n ],\n ],\n None,\n None,\n ]:\n \"\"\"Iterate over model activations in layers.\n\n Args:\n model (`nn.Module`):\n Model.\n action (`CacheAction`):\n Action for caching activations.\n needs_inputs_fn (`Callable[[str, nn.Module], bool]` or `bool` or `None`, *optional*, defaults to `True`):\n Function for determining whether to cache inputs for a module given its name and itself.\n needs_outputs_fn (`Callable[[str, nn.Module], bool]` or `bool` or `None`, *optional*, defaults to `None`):\n Function for determining whether to cache outputs for a module given its name and itself.\n *args: Arguments for ``iter_samples``.\n **kwargs: Keyword arguments for ``iter_samples``.\n\n Yields:\n Generator[\n tuple[str, tuple[DiffusionBlockStruct | nn.Module, dict[str, IOTensorsCache], dict[str, tp.Any]]],\n None,\n None\n ]:\n Generator of tuple of\n - layer name\n - a tuple of\n - layer itself\n - inputs and outputs cache of each module in the layer\n - layer input arguments\n \"\"\"\n if not isinstance(model, DiffusionModelStruct):\n model_struct = DiffusionModelStruct.construct(model)\n else:\n model_struct = model\n model = model_struct.module\n assert isinstance(model_struct, DiffusionModelStruct)\n assert isinstance(model, nn.Module)\n action = DiffusionConcatCacheAction(\"cpu\") if action is None else action\n layers, layer_structs, recomputes, use_prev_layer_outputs = model_struct.get_iter_layer_activations_args(\n skip_pre_modules=skip_pre_modules,\n skip_post_modules=skip_post_modules,\n **self.dataset[0][\"input_kwargs\"],\n )\n for layer_idx, (layer_name, (layer, layer_cache, layer_inputs)) in enumerate(\n self._iter_layer_activations(\n model,\n *args,\n action=action,\n layers=layers,\n needs_inputs_fn=needs_inputs_fn,\n needs_outputs_fn=needs_outputs_fn,\n recomputes=recomputes,\n use_prev_layer_outputs=use_prev_layer_outputs,\n **kwargs,\n )\n ):\n layer_kwargs = {k: v for k, v in layer_inputs[0].kwargs.items()} # noqa: C416\n layer_kwargs.pop(\"hidden_states\", None)\n layer_kwargs.pop(\"encoder_hidden_states\", None)\n layer_kwargs.pop(\"temb\", None)\n layer_struct = layer_structs[layer_idx]\n if isinstance(layer_struct, DiffusionBlockStruct):\n assert layer_struct.name == layer_name\n assert layer is layer_struct.module\n for transformer_block_struct in layer_struct.iter_transformer_block_structs():\n for attn_struct in transformer_block_struct.iter_attention_structs():\n if attn_struct.q_proj_name in layer_cache:\n if not attn_struct.is_cross_attn():\n cache = layer_cache[attn_struct.q_proj_name]\n layer_cache[attn_struct.k_proj_name] = cache\n layer_cache[attn_struct.v_proj_name] = cache\n if attn_struct.add_k_proj_name in layer_cache:\n assert not attn_struct.is_self_attn()\n cache = layer_cache[attn_struct.add_k_proj_name]\n layer_cache[attn_struct.add_v_proj_name] = cache\n if attn_struct.is_joint_attn():\n layer_cache[attn_struct.add_q_proj_name] = cache\n ffn_struct = transformer_block_struct.ffn_struct\n num_experts = ffn_struct.config.num_experts\n if ffn_struct is not None and num_experts > 1:\n for expert_idx in range(num_experts):\n if ffn_struct.up_proj_names[expert_idx] in layer_cache:\n cache = layer_cache[ffn_struct.up_proj_names[expert_idx]]\n for up_proj_name in ffn_struct.up_proj_names[expert_idx::num_experts]:\n layer_cache[up_proj_name] = cache\n if ffn_struct.down_proj_names[expert_idx] in layer_cache:\n cache = layer_cache[ffn_struct.down_proj_names[expert_idx]]\n for down_proj_name in ffn_struct.down_proj_names[expert_idx::num_experts]:\n layer_cache[down_proj_name] = cache\n yield layer_name, (layer_struct, layer_cache, layer_kwargs)\n"
},
{
"path": "deepcompressor/app/diffusion/dataset/collect/calib.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Collect calibration dataset.\"\"\"\n\nimport os\nfrom dataclasses import dataclass\n\nimport datasets\nimport torch\nfrom omniconfig import configclass\nfrom torch import nn\nfrom tqdm import tqdm\n\nfrom deepcompressor.app.diffusion.config import DiffusionPtqRunConfig\nfrom deepcompressor.utils.common import hash_str_to_int, tree_map\n\nfrom ...utils import get_control\nfrom ..data import get_dataset\nfrom .utils import CollectHook\n\n\ndef process(x: torch.Tensor) -> torch.Tensor:\n dtype = x.dtype\n return torch.from_numpy(x.float().numpy()).to(dtype)\n\n\ndef collect(config: DiffusionPtqRunConfig, dataset: datasets.Dataset):\n samples_dirpath = os.path.join(config.output.root, \"samples\")\n caches_dirpath = os.path.join(config.output.root, \"caches\")\n os.makedirs(samples_dirpath, exist_ok=True)\n os.makedirs(caches_dirpath, exist_ok=True)\n caches = []\n\n pipeline = config.pipeline.build()\n model = pipeline.unet if hasattr(pipeline, \"unet\") else pipeline.transformer\n assert isinstance(model, nn.Module)\n model.register_forward_hook(CollectHook(caches=caches), with_kwargs=True)\n\n batch_size = config.eval.batch_size\n print(f\"In total {len(dataset)} samples\")\n print(f\"Evaluating with batch size {batch_size}\")\n pipeline.set_progress_bar_config(desc=\"Sampling\", leave=False, dynamic_ncols=True, position=1)\n for batch in tqdm(\n dataset.iter(batch_size=batch_size, drop_last_batch=False),\n desc=\"Data\",\n leave=False,\n dynamic_ncols=True,\n total=(len(dataset) + batch_size - 1) // batch_size,\n ):\n filenames = batch[\"filename\"]\n prompts = batch[\"prompt\"]\n seeds = [hash_str_to_int(name) for name in filenames]\n generators = [torch.Generator(device=pipeline.device).manual_seed(seed) for seed in seeds]\n pipeline_kwargs = config.eval.get_pipeline_kwargs()\n\n task = config.pipeline.task\n control_root = config.eval.control_root\n if task in [\"canny-to-image\", \"depth-to-image\", \"inpainting\"]:\n controls = get_control(\n task,\n batch[\"image\"],\n names=batch[\"filename\"],\n data_root=os.path.join(\n control_root, collect_config.dataset_name, f\"{dataset.config_name}-{config.eval.num_samples}\"\n ),\n )\n if task == \"inpainting\":\n pipeline_kwargs[\"image\"] = controls[0]\n pipeline_kwargs[\"mask_image\"] = controls[1]\n else:\n pipeline_kwargs[\"control_image\"] = controls\n\n result_images = pipeline(prompts, generator=generators, **pipeline_kwargs).images\n num_guidances = (len(caches) // batch_size) // config.eval.num_steps\n num_steps = len(caches) // (batch_size * num_guidances)\n assert (\n len(caches) == batch_size * num_steps * num_guidances\n ), f\"Unexpected number of caches: {len(caches)} != {batch_size} * {config.eval.num_steps} * {num_guidances}\"\n for j, (filename, image) in enumerate(zip(filenames, result_images, strict=True)):\n image.save(os.path.join(samples_dirpath, f\"{filename}.png\"))\n for s in range(num_steps):\n for g in range(num_guidances):\n c = caches[s * batch_size * num_guidances + g * batch_size + j]\n c[\"filename\"] = filename\n c[\"step\"] = s\n c[\"guidance\"] = g\n c = tree_map(lambda x: process(x), c)\n torch.save(c, os.path.join(caches_dirpath, f\"{filename}-{s:05d}-{g}.pt\"))\n caches.clear()\n\n\n@configclass\n@dataclass\nclass CollectConfig:\n \"\"\"Configuration for collecting calibration dataset.\n\n Args:\n root (`str`, *optional*, defaults to `\"datasets\"`):\n Root directory to save the collected dataset.\n dataset_name (`str`, *optional*, defaults to `\"qdiff\"`):\n Name of the collected dataset.\n prompt_path (`str`, *optional*, defaults to `\"prompts/qdiff.yaml\"`):\n Path to the prompt file.\n num_samples (`int`, *optional*, defaults to `128`):\n Number of samples to collect.\n \"\"\"\n\n root: str = \"datasets\"\n dataset_name: str = \"qdiff\"\n data_path: str = \"prompts/qdiff.yaml\"\n num_samples: int = 128\n\n\nif __name__ == \"__main__\":\n parser = DiffusionPtqRunConfig.get_parser()\n parser.add_config(CollectConfig, scope=\"collect\", prefix=\"collect\")\n configs, _, unused_cfgs, unused_args, unknown_args = parser.parse_known_args()\n ptq_config, collect_config = configs[\"\"], configs[\"collect\"]\n assert isinstance(ptq_config, DiffusionPtqRunConfig)\n assert isinstance(collect_config, CollectConfig)\n if len(unused_cfgs) > 0:\n print(f\"Warning: unused configurations {unused_cfgs}\")\n if unused_args is not None:\n print(f\"Warning: unused arguments {unused_args}\")\n assert len(unknown_args) == 0, f\"Unknown arguments: {unknown_args}\"\n\n collect_dirpath = os.path.join(\n collect_config.root,\n str(ptq_config.pipeline.dtype),\n ptq_config.pipeline.name,\n ptq_config.eval.protocol,\n collect_config.dataset_name,\n f\"s{collect_config.num_samples}\",\n )\n print(f\"Saving caches to {collect_dirpath}\")\n\n dataset = get_dataset(\n collect_config.data_path,\n max_dataset_size=collect_config.num_samples,\n return_gt=ptq_config.pipeline.task in [\"canny-to-image\"],\n repeat=1,\n )\n\n ptq_config.output.root = collect_dirpath\n os.makedirs(ptq_config.output.root, exist_ok=True)\n collect(ptq_config, dataset=dataset)\n"
},
{
"path": "deepcompressor/app/diffusion/dataset/collect/utils.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Common utilities for collecting data.\"\"\"\n\nimport inspect\nimport typing as tp\n\nimport torch\nimport torch.nn as nn\nfrom diffusers.models.transformers import (\n FluxTransformer2DModel,\n PixArtTransformer2DModel,\n SanaTransformer2DModel,\n)\nfrom diffusers.models.unets.unet_2d_condition import UNet2DConditionModel\n\nfrom deepcompressor.utils.common import tree_map, tree_split\n\n__all__ = [\"CollectHook\"]\n\n\nclass CollectHook:\n def __init__(self, caches: list[dict[str, tp.Any]] = None, zero_redundancy: bool = False) -> None:\n self.caches = [] if caches is None else caches\n self.zero_redundancy = zero_redundancy\n\n def __call__(\n self,\n module: nn.Module,\n input_args: tuple[torch.Tensor, ...],\n input_kwargs: dict[str, tp.Any],\n output: tuple[torch.Tensor, ...],\n ) -> tp.Any:\n new_args = []\n signature = inspect.signature(module.forward)\n bound_arguments = signature.bind(*input_args, **input_kwargs)\n arguments = bound_arguments.arguments\n args_to_kwargs = {k: v for k, v in arguments.items() if k not in input_kwargs}\n input_kwargs.update(args_to_kwargs)\n\n if isinstance(module, UNet2DConditionModel):\n sample = input_kwargs.pop(\"sample\")\n new_args.append(sample)\n timestep = input_kwargs[\"timestep\"]\n timesteps = timestep\n if not torch.is_tensor(timesteps):\n is_mps = sample.device.type == \"mps\"\n if isinstance(timestep, float):\n dtype = torch.float32 if is_mps else torch.float64\n else:\n dtype = torch.int32 if is_mps else torch.int64\n timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)\n elif len(timesteps.shape) == 0:\n timesteps = timesteps[None].to(sample.device)\n # broadcast to batch dimension in a way that's compatible with ONNX/Core ML\n timesteps = timesteps.expand(sample.shape[0])\n input_kwargs[\"timestep\"] = timesteps\n elif isinstance(module, (PixArtTransformer2DModel, SanaTransformer2DModel)):\n new_args.append(input_kwargs.pop(\"hidden_states\"))\n elif isinstance(module, FluxTransformer2DModel):\n new_args.append(input_kwargs.pop(\"hidden_states\"))\n else:\n raise ValueError(f\"Unknown model: {module}\")\n cache = tree_map(lambda x: x.cpu(), {\"input_args\": new_args, \"input_kwargs\": input_kwargs, \"outputs\": output})\n split_cache = tree_split(cache)\n\n if isinstance(module, PixArtTransformer2DModel) and self.zero_redundancy:\n for cache in split_cache:\n cache_kwargs = cache[\"input_kwargs\"]\n encoder_hidden_states = cache_kwargs.pop(\"encoder_hidden_states\")\n assert encoder_hidden_states.shape[0] == 1\n encoder_attention_mask = cache_kwargs.get(\"encoder_attention_mask\", None)\n if encoder_attention_mask is not None:\n encoder_hidden_states = encoder_hidden_states[:, : max(encoder_attention_mask.sum(), 1)]\n cache_kwargs[\"encoder_hidden_states\"] = encoder_hidden_states\n\n self.caches.extend(split_cache)\n"
},
{
"path": "deepcompressor/app/diffusion/dataset/data/COCO/COCO.py",
"content": "# coding=utf-8\n# Copyright 2022 The HuggingFace Datasets Authors and the current dataset script contributor.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"COCO\"\"\"\n\nimport json\nimport os\nimport random\nfrom pathlib import Path\n\nimport datasets\nfrom PIL import Image\n\n_CITATION = \"\"\"\n@article{DBLP:journals/corr/LinMBHPRDZ14,\n author = {Tsung{-}Yi Lin and\n Michael Maire and\n Serge J. Belongie and\n Lubomir D. Bourdev and\n Ross B. Girshick and\n James Hays and\n Pietro Perona and\n Deva Ramanan and\n Piotr Doll{\\'{a}}r and\n C. Lawrence Zitnick},\n title = {Microsoft {COCO:} Common Objects in Context},\n journal = {CoRR},\n volume = {abs/1405.0312},\n year = {2014},\n url = {http://arxiv.org/abs/1405.0312},\n eprinttype = {arXiv},\n eprint = {1405.0312},\n timestamp = {Mon, 13 Aug 2018 16:48:13 +0200},\n biburl = {https://dblp.org/rec/journals/corr/LinMBHPRDZ14.bib},\n bibsource = {dblp computer science bibliography, https://dblp.org}\n}\n\"\"\"\n\n_DESCRIPTION = \"\"\"\nMS COCO is a large-scale object detection, segmentation, and captioning dataset.\n COCO has several features: Object segmentation, Recognition in context, Superpixel stuff segmentation,\n 330K images (>200K labeled), 1.5 million object instances, 80 object categories, 91 stuff categories,\n 5 captions per image, 250,000 people with keypoints.\n\"\"\"\n\n_HOMEPAGE = \"https://cocodataset.org/#home\"\n\n_LICENSE = \"CC BY 4.0\"\n\n\n_IMAGES_URLS = {\n \"train\": \"http://images.cocodataset.org/zips/train2014.zip\",\n \"validation\": \"http://images.cocodataset.org/zips/val2014.zip\",\n}\n\n_KARPATHY_FILES_URL = \"https://cs.stanford.edu/people/karpathy/deepimagesent/caption_datasets.zip\"\n\n_FEATURES = datasets.Features(\n {\n \"filepath\": datasets.Value(\"string\"),\n \"filename\": datasets.Value(\"string\"),\n \"image\": datasets.Image(),\n \"image_path\": datasets.Value(\"string\"),\n \"image_root\": datasets.Value(\"string\"),\n \"prompt\": datasets.Value(\"string\"),\n \"prompt_id\": datasets.Value(\"int32\"),\n \"imgid\": datasets.Value(\"int32\"),\n \"split\": datasets.Value(\"string\"),\n \"cocoid\": datasets.Value(\"int32\"),\n \"sentences_raw\": [datasets.Value(\"string\")],\n \"sentids\": [datasets.Value(\"int32\")],\n \"sentences_sentid\": [datasets.Value(\"int32\")],\n \"sentences_tokens\": [[datasets.Value(\"string\")]],\n }\n)\n\n\ndef hash_string_to_int(s: str) -> int:\n modulus = 10**9 + 7 # Large prime modulus\n hash_int = 0\n for char in s:\n hash_int = (hash_int * 31 + ord(char)) % modulus\n return hash_int\n\n\nclass COCOConfig(datasets.BuilderConfig):\n def __init__(self, max_dataset_size: int = -1, return_gt: bool = False, **kwargs):\n super(COCOConfig, self).__init__(\n name=kwargs.get(\"name\", \"default\"),\n version=kwargs.get(\"version\", \"0.0.0\"),\n data_dir=kwargs.get(\"data_dir\", None),\n data_files=kwargs.get(\"data_files\", None),\n description=kwargs.get(\"description\", None),\n )\n self.max_dataset_size = max_dataset_size\n self.return_gt = return_gt\n\n\nclass COCO(datasets.GeneratorBasedBuilder):\n \"\"\"COCO\"\"\"\n\n VERSION = datasets.Version(\"0.0.0\")\n\n BUILDER_CONFIG_CLASS = COCOConfig\n BUILDER_CONFIGS = [\n COCOConfig(name=\"COCO_val\", version=VERSION, description=\"COCO validation prompt set\"),\n COCOConfig(name=\"COCO_train\", version=VERSION, description=\"COCO train prompt set\"),\n COCOConfig(name=\"COCO_full\", version=VERSION, description=\"COCO full prompt set\"),\n ]\n DEFAULT_CONFIG_NAME = \"COCO_val\"\n\n def _info(self):\n return datasets.DatasetInfo(\n description=_DESCRIPTION,\n features=_FEATURES,\n homepage=_HOMEPAGE,\n license=_LICENSE,\n citation=_CITATION,\n )\n\n def _split_generators(self, dl_manager: datasets.download.DownloadManager):\n annotation_file = os.path.join(dl_manager.download_and_extract(_KARPATHY_FILES_URL), \"dataset_coco.json\")\n image_folders = {k: Path(v) for k, v in dl_manager.download_and_extract(_IMAGES_URLS).items()}\n\n if self.config.name == \"COCO_full\":\n split_keys = [\"validation\", \"train\"]\n else:\n split_keys = [self.config.name.split(\"_\")[-1]]\n\n return [\n datasets.SplitGenerator(\n name=datasets.Split.TRAIN,\n gen_kwargs={\n \"annotation_file\": annotation_file,\n \"image_folders\": image_folders,\n \"split_keys\": split_keys,\n },\n ),\n ]\n\n def _generate_examples(\n self, annotation_file: str, image_folders: dict[str, str], split_keys: list[str] | tuple[str, ...]\n ):\n with open(annotation_file, \"r\", encoding=\"utf-8\") as fi:\n annotations = json.load(fi)\n metas = []\n for split_key in split_keys:\n for image_metadata in annotations[\"images\"]:\n if split_key == \"train\":\n if image_metadata[\"split\"] != \"train\" and image_metadata[\"split\"] != \"restval\":\n continue\n elif split_key == \"val\":\n if image_metadata[\"split\"] != \"val\":\n continue\n elif split_key == \"test\":\n if image_metadata[\"split\"] != \"test\":\n continue\n\n metas.append(image_metadata)\n\n if self.config.max_dataset_size > 0:\n random.Random(0).shuffle(metas)\n metas = metas[: self.config.max_dataset_size]\n metas = sorted(metas, key=lambda x: x[\"filename\"])\n\n for i, meta in enumerate(metas):\n if \"val2014\" in meta[\"filename\"]:\n image_root = os.path.join(image_folders[\"validation\"], \"val2014\")\n else:\n image_root = os.path.join(image_folders[\"train\"], \"train2014\")\n filename = meta[\"filename\"].replace(\".jpg\", \"\").replace(\".png\", \"\")\n image_path = os.path.join(image_root, filename + \".jpg\")\n\n sentences_raw = [caption[\"raw\"] for caption in meta[\"sentences\"]]\n prompt_id = hash_string_to_int(filename) % len(sentences_raw)\n prompt = sentences_raw[prompt_id]\n\n yield (\n i,\n {\n \"filename\": filename,\n \"image\": Image.open(image_path) if self.config.return_gt else None,\n \"image_path\": image_path,\n \"image_root\": image_root,\n \"prompt\": prompt,\n \"prompt_id\": prompt_id,\n \"imgid\": meta[\"imgid\"],\n \"split\": self.config.name,\n \"coco_id\": meta[\"cocoid\"],\n \"sentences_raw\": sentences_raw,\n \"sentids\": meta[\"sentids\"],\n \"sentences_sentid\": [caption[\"sentid\"] for caption in meta[\"sentences\"]],\n \"sentences_tokens\": [caption[\"tokens\"] for caption in meta[\"sentences\"]],\n },\n )\n"
},
{
"path": "deepcompressor/app/diffusion/dataset/data/COCO/__init__.py",
"content": ""
},
{
"path": "deepcompressor/app/diffusion/dataset/data/DCI/DCI.py",
"content": "import os\nimport random\n\nimport datasets\nimport yaml\nfrom PIL import Image\n\n_CITATION = \"\"\"\\\n@InProceedings{Urbanek_2024_CVPR,\n author = {Urbanek, Jack and Bordes, Florian and Astolfi, Pietro and Williamson, Mary and Sharma, Vasu and Romero-Soriano, Adriana},\n title = {A Picture is Worth More Than 77 Text Tokens: Evaluating CLIP-Style Models on Dense Captions},\n booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},\n month = {June},\n year = {2024},\n pages = {26700-26709}\n}\n\"\"\" # noqa: E501\n\n_DESCRIPTION = \"\"\"\\\nThe Densely Captioned Images dataset, or DCI, consists of 7805 images from SA-1B,\n each with a complete description aiming to capture the full visual detail of what is present in the image.\n Much of the description is directly aligned to submasks of the image.\n\"\"\"\n\n_HOMEPAGE = \"https://github.com/facebookresearch/DCI\"\n\n_LICENSE = \"Attribution-NonCommercial 4.0 International (https://github.com/facebookresearch/DCI/blob/main/LICENSE)\"\n\nIMAGE_URL = \"https://huggingface.co/datasets/mit-han-lab/svdquant-datasets/resolve/main/sDCI.gz\"\n\nPROMPT_URLS = {\"sDCI\": \"https://huggingface.co/datasets/mit-han-lab/svdquant-datasets/resolve/main/sDCI.yaml\"}\n\n\nclass DCIConfig(datasets.BuilderConfig):\n def __init__(self, max_dataset_size: int = -1, return_gt: bool = False, **kwargs):\n super(DCIConfig, self).__init__(\n name=kwargs.get(\"name\", \"default\"),\n version=kwargs.get(\"version\", \"0.0.0\"),\n data_dir=kwargs.get(\"data_dir\", None),\n data_files=kwargs.get(\"data_files\", None),\n description=kwargs.get(\"description\", None),\n )\n self.max_dataset_size = max_dataset_size\n self.return_gt = return_gt\n\n\nclass DCI(datasets.GeneratorBasedBuilder):\n VERSION = datasets.Version(\"0.0.0\")\n\n BUILDER_CONFIG_CLASS = DCIConfig\n BUILDER_CONFIGS = [DCIConfig(name=\"sDCI\", version=VERSION, description=\"sDCI full prompt set\")]\n DEFAULT_CONFIG_NAME = \"sDCI\"\n\n def _info(self):\n features = datasets.Features(\n {\n \"filename\": datasets.Value(\"string\"),\n \"image\": datasets.Image(),\n \"prompt\": datasets.Value(\"string\"),\n \"meta_path\": datasets.Value(\"string\"),\n \"image_root\": datasets.Value(\"string\"),\n \"image_path\": datasets.Value(\"string\"),\n \"split\": datasets.Value(\"string\"),\n }\n )\n return datasets.DatasetInfo(\n description=_DESCRIPTION, features=features, homepage=_HOMEPAGE, license=_LICENSE, citation=_CITATION\n )\n\n def _split_generators(self, dl_manager: datasets.download.DownloadManager):\n image_url = IMAGE_URL\n meta_url = PROMPT_URLS[self.config.name]\n\n meta_path = dl_manager.download(meta_url)\n image_root = dl_manager.download_and_extract(image_url)\n\n return [\n datasets.SplitGenerator(\n name=datasets.Split.TRAIN, gen_kwargs={\"meta_path\": meta_path, \"image_root\": image_root}\n )\n ]\n\n def _generate_examples(self, meta_path: str, image_root: str):\n meta = yaml.safe_load(open(meta_path, \"r\"))\n names = list(meta.keys())\n if self.config.max_dataset_size > 0:\n random.Random(0).shuffle(names)\n names = names[: self.config.max_dataset_size]\n names = sorted(names)\n\n for i, name in enumerate(names):\n prompt = meta[name]\n image_path = os.path.join(image_root, f\"{name}.jpg\")\n yield (\n i,\n {\n \"filename\": name,\n \"image\": Image.open(image_path) if self.config.return_gt else None,\n \"prompt\": prompt,\n \"meta_path\": meta_path,\n \"image_root\": image_root,\n \"image_path\": image_path,\n \"split\": self.config.name,\n },\n )\n"
},
{
"path": "deepcompressor/app/diffusion/dataset/data/DCI/__init__.py",
"content": ""
},
{
"path": "deepcompressor/app/diffusion/dataset/data/MJHQ/MJHQ.py",
"content": "import json\nimport os\nimport random\n\nimport datasets\nfrom PIL import Image\n\n_CITATION = \"\"\"\\\n@misc{li2024playground,\n title={Playground v2.5: Three Insights towards Enhancing Aesthetic Quality in Text-to-Image Generation},\n author={Daiqing Li and Aleks Kamko and Ehsan Akhgari and Ali Sabet and Linmiao Xu and Suhail Doshi},\n year={2024},\n eprint={2402.17245},\n archivePrefix={arXiv},\n primaryClass={cs.CV}\n}\n\"\"\"\n\n_DESCRIPTION = \"\"\"\\\nWe introduce a new benchmark, MJHQ-30K, for automatic evaluation of a modelโs aesthetic quality.\n The benchmark computes FID on a high-quality dataset to gauge aesthetic quality.\n\"\"\"\n\n_HOMEPAGE = \"https://huggingface.co/datasets/playgroundai/MJHQ-30K\"\n\n_LICENSE = (\n \"Playground v2.5 Community License \"\n \"(https://huggingface.co/playgroundai/playground-v2.5-1024px-aesthetic/blob/main/LICENSE.md)\"\n)\n\nIMAGE_URL = \"https://huggingface.co/datasets/playgroundai/MJHQ-30K/resolve/main/mjhq30k_imgs.zip\"\n\nMETA_URL = \"https://huggingface.co/datasets/playgroundai/MJHQ-30K/resolve/main/meta_data.json\"\n\n\nclass MJHQConfig(datasets.BuilderConfig):\n def __init__(self, max_dataset_size: int = -1, return_gt: bool = False, **kwargs):\n super(MJHQConfig, self).__init__(\n name=kwargs.get(\"name\", \"default\"),\n version=kwargs.get(\"version\", \"0.0.0\"),\n data_dir=kwargs.get(\"data_dir\", None),\n data_files=kwargs.get(\"data_files\", None),\n description=kwargs.get(\"description\", None),\n )\n self.max_dataset_size = max_dataset_size\n self.return_gt = return_gt\n\n\nclass DCI(datasets.GeneratorBasedBuilder):\n VERSION = datasets.Version(\"0.0.0\")\n\n BUILDER_CONFIG_CLASS = MJHQConfig\n BUILDER_CONFIGS = [MJHQConfig(name=\"MJHQ\", version=VERSION, description=\"MJHQ-30K full dataset\")]\n DEFAULT_CONFIG_NAME = \"MJHQ\"\n\n def _info(self):\n features = datasets.Features(\n {\n \"filename\": datasets.Value(\"string\"),\n \"category\": datasets.Value(\"string\"),\n \"image\": datasets.Image(),\n \"prompt\": datasets.Value(\"string\"),\n \"prompt_path\": datasets.Value(\"string\"),\n \"image_root\": datasets.Value(\"string\"),\n \"image_path\": datasets.Value(\"string\"),\n \"split\": datasets.Value(\"string\"),\n }\n )\n return datasets.DatasetInfo(\n description=_DESCRIPTION, features=features, homepage=_HOMEPAGE, license=_LICENSE, citation=_CITATION\n )\n\n def _split_generators(self, dl_manager: datasets.download.DownloadManager):\n meta_path = dl_manager.download(META_URL)\n image_root = dl_manager.download_and_extract(IMAGE_URL)\n return [\n datasets.SplitGenerator(\n name=datasets.Split.TRAIN, gen_kwargs={\"meta_path\": meta_path, \"image_root\": image_root}\n ),\n ]\n\n def _generate_examples(self, meta_path: str, image_root: str):\n with open(meta_path, \"r\") as f:\n meta = json.load(f)\n\n names = list(meta.keys())\n if self.config.max_dataset_size > 0:\n random.Random(0).shuffle(names)\n names = names[: self.config.max_dataset_size]\n names = sorted(names)\n\n for i, name in enumerate(names):\n category = meta[name][\"category\"]\n prompt = meta[name][\"prompt\"]\n image_path = os.path.join(image_root, category, f\"{name}.jpg\")\n yield (\n i,\n {\n \"filename\": name,\n \"category\": category,\n \"image\": Image.open(image_path) if self.config.return_gt else None,\n \"prompt\": prompt,\n \"meta_path\": meta_path,\n \"image_root\": image_root,\n \"image_path\": image_path,\n \"split\": self.config.name,\n },\n )\n"
},
{
"path": "deepcompressor/app/diffusion/dataset/data/MJHQ/__init__.py",
"content": ""
},
{
"path": "deepcompressor/app/diffusion/dataset/data/__init__.py",
"content": "import os\nimport random\n\nimport datasets\nimport yaml\n\n__all__ = [\"get_dataset\"]\n\n\ndef load_dataset_yaml(meta_path: str, max_dataset_size: int = -1, repeat: int = 4) -> dict:\n meta = yaml.safe_load(open(meta_path, \"r\"))\n names = list(meta.keys())\n if max_dataset_size > 0:\n random.Random(0).shuffle(names)\n names = names[:max_dataset_size]\n names = sorted(names)\n\n ret = {\"filename\": [], \"prompt\": [], \"meta_path\": []}\n idx = 0\n for name in names:\n prompt = meta[name]\n for j in range(repeat):\n ret[\"filename\"].append(f\"{name}-{j}\")\n ret[\"prompt\"].append(prompt)\n ret[\"meta_path\"].append(meta_path)\n idx += 1\n return ret\n\n\ndef get_dataset(\n name: str,\n config_name: str | None = None,\n split: str = \"train\",\n max_dataset_size: int = -1,\n return_gt: bool = False,\n repeat: int = 4,\n chunk_start: int = 0,\n chunk_step: int = 1,\n) -> datasets.Dataset:\n prefix = os.path.dirname(__file__)\n kwargs = {\n \"name\": config_name,\n \"split\": split,\n \"trust_remote_code\": True,\n \"token\": True,\n \"max_dataset_size\": max_dataset_size,\n }\n if name.endswith((\".yaml\", \".yml\")):\n dataset = datasets.Dataset.from_dict(\n load_dataset_yaml(name, max_dataset_size=max_dataset_size, repeat=repeat),\n features=datasets.Features(\n {\n \"filename\": datasets.Value(\"string\"),\n \"prompt\": datasets.Value(\"string\"),\n \"meta_path\": datasets.Value(\"string\"),\n }\n ),\n )\n else:\n path = os.path.join(prefix, f\"{name}\")\n if name == \"COCO\":\n dataset = datasets.load_dataset(path, return_gt=return_gt, **kwargs)\n elif name == \"DCI\":\n dataset = datasets.load_dataset(path, return_gt=return_gt, **kwargs)\n elif name == \"MJHQ\":\n dataset = datasets.load_dataset(path, return_gt=return_gt, **kwargs)\n else:\n raise ValueError(f\"Unknown dataset name: {name}\")\n assert not hasattr(dataset, \"_unchunk_size\")\n assert not hasattr(dataset, \"_chunk_start\")\n assert not hasattr(dataset, \"_chunk_step\")\n unchunk_size = len(dataset)\n if chunk_step > 1 or chunk_start > 0:\n assert 0 <= chunk_start < chunk_step\n dataset = dataset.select(range(chunk_start, len(dataset), chunk_step))\n else:\n chunk_start, chunk_step = 0, 1\n dataset._unchunk_size = unchunk_size\n dataset._chunk_start = chunk_start\n dataset._chunk_step = chunk_step\n return dataset\n"
},
{
"path": "deepcompressor/app/diffusion/dataset/data/dump.py",
"content": "import argparse\nimport os\n\nimport yaml\nfrom tqdm import tqdm\n\nfrom ...utils import get_control\nfrom . import get_dataset\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--benchmarks\", type=str, nargs=\"*\", default=[\"COCO\", \"DCI\", \"MJHQ\"])\n parser.add_argument(\"--max-dataset-size\", type=int, default=-1)\n parser.add_argument(\"--dump-root\", type=str, default=\"benchmarks\")\n parser.add_argument(\"--copy-images\", action=\"store_true\")\n parser.add_argument(\"--prompts-only\", action=\"store_true\")\n parser.add_argument(\"--controls\", type=str, nargs=\"*\", default=[\"canny-to-image\", \"depth-to-image\", \"inpainting\"])\n parser.add_argument(\"--chunk-start\", type=int, default=0)\n parser.add_argument(\"--chunk-step\", type=int, default=1)\n args = parser.parse_args()\n\n if \"depth-to-image\" in args.controls:\n from image_gen_aux import DepthPreprocessor\n\n processor = DepthPreprocessor.from_pretrained(\"LiheYoung/depth-anything-large-hf\").to(\"cuda\")\n\n for benchmark in args.benchmarks:\n dataset = get_dataset(\n benchmark,\n max_dataset_size=args.max_dataset_size,\n return_gt=True,\n chunk_start=args.chunk_start,\n chunk_step=args.chunk_step,\n )\n prompts = {}\n benchmark_root = os.path.join(args.dump_root, benchmark, f\"{dataset.config_name}-{dataset._unchunk_size}\")\n for row in tqdm(dataset, desc=f\"Dumping {dataset.config_name}\"):\n prompts[row[\"filename\"]] = row[\"prompt\"]\n if not args.prompts_only:\n image = row.get(\"image\", None)\n if image is not None:\n image_root = os.path.join(benchmark_root, \"images\")\n os.makedirs(image_root, exist_ok=True)\n if args.copy_images:\n image.save(os.path.join(image_root, row[\"filename\"] + \".png\"))\n else:\n ext = os.path.basename(row[\"image_path\"]).split(\".\")[-1]\n os.symlink(\n os.path.abspath(os.path.expanduser(row[\"image_path\"])),\n os.path.abspath(os.path.expanduser(os.path.join(image_root, row[\"filename\"] + f\".{ext}\"))),\n )\n if \"canny-to-image\" in args.controls:\n canny_root = os.path.join(benchmark_root, \"canny_images\")\n os.makedirs(canny_root, exist_ok=True)\n canny = get_control(\"canny-to-image\", image)\n canny.save(os.path.join(canny_root, row[\"filename\"] + \".png\"))\n if \"depth-to-image\" in args.controls:\n depth_root = os.path.join(benchmark_root, \"depth_images\")\n os.makedirs(depth_root, exist_ok=True)\n depth = get_control(\"depth-to-image\", image, processor=processor)\n depth.save(os.path.join(depth_root, row[\"filename\"] + \".png\"))\n if \"inpainting\" in args.controls:\n mask_root = os.path.join(benchmark_root, \"mask_images\")\n cropped_image_root = os.path.join(benchmark_root, \"cropped_images\")\n os.makedirs(mask_root, exist_ok=True)\n os.makedirs(cropped_image_root, exist_ok=True)\n cropped_image, mask_image = get_control(\"inpainting\", image, names=row[\"filename\"])\n cropped_image.save(os.path.join(cropped_image_root, row[\"filename\"] + \".png\"))\n mask_image.save(os.path.join(mask_root, row[\"filename\"] + \".png\"))\n\n if args.chunk_step == 1:\n os.makedirs(benchmark_root, exist_ok=True)\n with open(os.path.join(benchmark_root, \"prompts.yaml\"), \"w\") as f:\n yaml.dump(prompts, f)\n"
},
{
"path": "deepcompressor/app/diffusion/eval/__init__.py",
"content": "# -*- coding: utf-8 -*-\n\nfrom .config import DiffusionEvalConfig\n"
},
{
"path": "deepcompressor/app/diffusion/eval/config.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Diffusion model evaluation.\"\"\"\n\nimport logging\nimport os\nimport typing as tp\nfrom dataclasses import dataclass, field\n\nimport datasets\nimport diffusers\nimport omniconfig\nimport torch\nfrom diffusers import DiffusionPipeline\nfrom omniconfig import configclass\nfrom torch import multiprocessing as mp\nfrom tqdm import tqdm\n\nfrom deepcompressor.app.diffusion.dataset.data import get_dataset\nfrom deepcompressor.utils.common import hash_str_to_int\n\nfrom ..utils import get_control\nfrom .metrics import compute_image_metrics\n\n__all__ = [\"DiffusionEvalConfig\"]\n\n\n@configclass\n@dataclass\nclass DiffusionEvalConfig:\n \"\"\"Diffusion model evaluation configuration.\n\n Args:\n protocol (`str`):\n The protocol of the evaluation pipeline.\n num_gpus (`int`, *optional*, defaults to `1`):\n The number of GPUs to use.\n batch_size (`int`, *optional*, defaults to `1`):\n The batch size used for inference.\n batch_size_per_gpu (`int`, *optional*, defaults to `None`):\n The batch size per GPU.\n height (`int`, *optional*, defaults to `None`):\n The height of the generated images.\n width (`int`, *optional*, defaults to `None`):\n The width of the generated images.\n clean_caption (`bool`, *optional*, defaults to `None`):\n Whether to clean the caption.\n num_steps (`int`, *optional*, defaults to `None`):\n The number of inference steps.\n guidance_scale (`float`, *optional*, defaults to `None`):\n The guidance scale.\n num_samples (`int`, *optional*, defaults to `1024`):\n The number of samples to generate.\n benchmarks (`list[str]`, *optional*, defaults to `[\"COCO\", \"DCI\", \"MJHQ\", \"GenEval\"]`):\n The benchmark datasets to evaluate on.\n gt_metrics (`list[str]`, *optional*, defaults to `[\"clip_iqa\", \"clip_score\", \"psnr\", \"lpips\", \"ssim\", \"fid\"]`):\n The ground truth metrics to compute.\n ref_metrics (`list[str]`, *optional*, defaults to `[\"psnr\", \"lpips\", \"ssim\", \"fid\"]`):\n The reference metrics to compute.\n ref_root (`str`, *optional*, defaults to `\"\"`):\n The root directory path to the reference images.\n gt_stats_root (`str`, *optional*, defaults to `\"\"`):\n The root directory path to the ground truth statistics.\n chunk_start (`int`, *optional*, defaults to `0`):\n The starting chunk index.\n chunk_step (`int`, *optional*, defaults to `1`):\n The chunk step size.\n \"\"\"\n\n protocol: str\n\n num_gpus: int = field(default=1, metadata={omniconfig.ARGPARSE_ARGS: (\"--num-gpus\", \"-n\")})\n batch_size: int = 1\n batch_size_per_gpu: int | None = None\n\n height: int | None = None\n width: int | None = None\n clean_caption: bool | None = None\n num_steps: int | None = None\n guidance_scale: float | None = None\n num_samples: int = 1024\n\n benchmarks: list[str] = field(\n default_factory=lambda: [\"COCO\", \"DCI\", \"MJHQ\", \"GenEval\"],\n metadata={omniconfig.ARGPARSE_KWARGS: {\"nargs\": \"+\", \"type\": str}},\n )\n gt_metrics: list[str] = field(\n default_factory=lambda: [\"clip_iqa\", \"clip_score\", \"image_reward\", \"fid\"],\n metadata={omniconfig.ARGPARSE_KWARGS: {\"nargs\": \"+\", \"type\": str}},\n )\n ref_metrics: list[str] = field(\n default_factory=lambda: [\"psnr\", \"lpips\", \"ssim\"],\n metadata={omniconfig.ARGPARSE_KWARGS: {\"nargs\": \"+\", \"type\": str}},\n )\n gen_root: str = \"\"\n ref_root: str = \"\"\n gt_stats_root: str = \"\"\n control_root: str | None = None\n\n chunk_start: int = 0\n chunk_step: int = 1\n chunk_only: bool = False\n\n def __post_init__(self):\n assert self.protocol\n self.protocol = self.protocol.lower().format(num_steps=self.num_steps, guidance_scale=self.guidance_scale)\n assert 0 <= self.chunk_start < self.chunk_step\n if self.chunk_start == 0 and self.chunk_step == 1:\n self.chunk_only = False\n\n def get_pipeline_kwargs(self) -> dict[str, tp.Any]:\n kwargs = {}\n if self.height is not None:\n kwargs[\"height\"] = self.height\n if self.width is not None:\n kwargs[\"width\"] = self.width\n if self.clean_caption is not None:\n kwargs[\"clean_caption\"] = self.clean_caption\n if self.num_steps is not None:\n kwargs[\"num_inference_steps\"] = self.num_steps\n if self.guidance_scale is not None:\n kwargs[\"guidance_scale\"] = self.guidance_scale\n return kwargs\n\n def _generate(\n self,\n rank: int,\n dataset: datasets.Dataset,\n pipeline: DiffusionPipeline,\n dirpath: str,\n logger: logging.Logger,\n dataset_name: str | None = None,\n task: str = \"text-to-image\",\n control_root: str | None = None,\n ) -> None:\n if self.num_gpus > 1:\n pipeline = pipeline.to(rank)\n if rank == 0:\n logger.info(\n f\" {dataset.config_name} has {len(dataset)} samples \"\n f\"(chunk_start={dataset._chunk_start}, chunk_step={dataset._chunk_step},\"\n f\" unchunk_size={dataset._unchunk_size})\"\n )\n pipeline.set_progress_bar_config(\n desc=\"Sampling\",\n leave=False,\n dynamic_ncols=True,\n position=1,\n disable=self.num_gpus > 1,\n )\n if dataset_name is None:\n dataset_name = dataset.config_name\n for batch in tqdm(\n dataset.iter(batch_size=self.batch_size, drop_last_batch=False),\n desc=dataset_name if self.num_gpus == 1 else f\"{dataset_name} (GPU {rank})\",\n leave=False,\n dynamic_ncols=True,\n position=rank,\n total=(len(dataset) + self.batch_size - 1) // self.batch_size,\n ):\n filenames = batch[\"filename\"][rank :: self.num_gpus]\n if len(filenames) == 0:\n continue\n if all(os.path.exists(os.path.join(dirpath, f\"{filename}.png\")) for filename in filenames):\n continue\n prompts = batch[\"prompt\"][rank :: self.num_gpus]\n seeds = [hash_str_to_int(name) for name in filenames]\n diffusers.training_utils.set_seed(seeds[0])\n generators = [torch.Generator().manual_seed(seed) for seed in seeds]\n\n pipeline_kwargs = self.get_pipeline_kwargs()\n\n if task in [\"canny-to-image\", \"depth-to-image\", \"inpainting\"]:\n controls = get_control(\n task,\n batch[\"image\"],\n names=batch[\"filename\"],\n data_root=os.path.join(control_root, f\"{dataset_name}-{dataset._unchunk_size}\"),\n )\n if task == \"inpainting\":\n pipeline_kwargs[\"image\"] = controls[0]\n pipeline_kwargs[\"mask_image\"] = controls[1]\n else:\n pipeline_kwargs[\"control_image\"] = controls\n\n output = pipeline(prompts, generator=generators, **pipeline_kwargs)\n images = output.images\n for filename, image in zip(filenames, images, strict=True):\n image.save(os.path.join(dirpath, f\"{filename}.png\"))\n\n def generate(\n self,\n pipeline: DiffusionPipeline,\n gen_root: str = \"\",\n task: str = \"text-to-image\",\n ) -> None:\n logger = logging.getLogger(f\"{__name__}.DiffusionEval\")\n gen_root = gen_root or self.gen_root\n for benchmark in self.benchmarks:\n dataset = get_dataset(\n benchmark,\n max_dataset_size=self.num_samples,\n chunk_start=self.chunk_start,\n chunk_step=self.chunk_step,\n return_gt=task in [\"canny-to-image\"],\n repeat=1,\n )\n if benchmark.endswith(\".yaml\") or benchmark.endswith(\".yml\"):\n dataset_name = os.path.splitext(os.path.basename(benchmark))[0]\n dirpath = os.path.join(\n gen_root,\n \"samples\",\n \"YAML\",\n f\"{dataset_name}-{dataset._unchunk_size}\",\n )\n else:\n dataset_name = dataset.config_name\n dirpath = os.path.join(\n gen_root,\n \"samples\",\n benchmark,\n f\"{dataset.config_name}-{dataset._unchunk_size}\",\n )\n if self.chunk_only:\n dirpath += f\".{dataset._chunk_start}.{dataset._chunk_step}\"\n os.makedirs(dirpath, exist_ok=True)\n args = (dataset, pipeline, dirpath, logger, dataset_name, task, os.path.join(self.control_root, benchmark))\n if self.num_gpus == 1:\n self._generate(0, *args)\n else:\n mp.spawn(self._generate, args=args, nprocs=self.num_gpus, join=True)\n\n def evaluate(\n self, pipeline: DiffusionPipeline, gen_root: str = \"\", skip_gen: bool = False, task: str = \"text-to-image\"\n ) -> dict[str, tp.Any] | None:\n gen_root = gen_root or self.gen_root\n if not skip_gen:\n self.generate(pipeline, gen_root=gen_root, task=task)\n if not self.chunk_only:\n return compute_image_metrics(\n gen_root=gen_root,\n benchmarks=self.benchmarks,\n max_dataset_size=self.num_samples,\n chunk_start=self.chunk_start,\n chunk_step=self.chunk_step,\n ref_root=self.ref_root,\n gt_stats_root=self.gt_stats_root,\n gt_metrics=self.gt_metrics,\n ref_metrics=self.ref_metrics,\n )\n else:\n return {}\n"
},
{
"path": "deepcompressor/app/diffusion/eval/metrics/__init__.py",
"content": "import logging\nimport os\n\nfrom deepcompressor.app.diffusion.dataset.data import get_dataset\n\nfrom .fid import compute_fid\nfrom .image_reward import compute_image_reward\nfrom .multimodal import compute_image_multimodal_metrics\nfrom .similarity import compute_image_similarity_metrics\n\nlogging.getLogger(\"PIL\").setLevel(logging.WARNING)\n\n__all__ = [\"compute_image_metrics\"]\n\n\ndef compute_image_metrics(\n gen_root: str,\n benchmarks: str | tuple[str, ...] = (\"DCI\", \"GenAIBench\", \"GenEval\", \"MJHQ\", \"T2ICompBench\"),\n max_dataset_size: int = -1,\n chunk_start: int = 0,\n chunk_step: int = 1,\n chunk_only: bool = False,\n ref_root: str = \"\",\n gt_stats_root: str = \"\",\n gt_metrics: tuple[str, ...] = (\"clip_iqa\", \"clip_score\", \"image_reward\", \"fid\"),\n ref_metrics: tuple[str, ...] = (\"psnr\", \"lpips\", \"ssim\", \"fid\"),\n) -> dict:\n if chunk_start == 0 and chunk_step == 1:\n chunk_only = False\n assert chunk_start == 0 and chunk_step == 1, \"Chunking is not supported for image data.\"\n os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n if isinstance(benchmarks, str):\n benchmarks = (benchmarks,)\n gt_multimodal_metrics, gt_similarity_metrics, gt_other_metrics = categorize_metrics(gt_metrics)\n _, ref_similarity_metrics, ref_other_metrics = categorize_metrics(ref_metrics)\n results = {}\n for benchmark in benchmarks:\n benchmark_results = {}\n dataset = get_dataset(benchmark, max_dataset_size=max_dataset_size, return_gt=True)\n dirname = f\"{dataset.config_name}-{dataset._unchunk_size}\"\n if dataset._chunk_start == 0 and dataset._chunk_step == 1:\n filename = f\"{dirname}.npz\"\n else:\n filename = os.path.join(dirname, f\"{dataset._chunk_start}-{dataset._chunk_step}.npz\")\n if chunk_only:\n dirname += f\".{dataset._chunk_start}.{dataset._chunk_step}\"\n gen_dirpath = os.path.join(gen_root, \"samples\", benchmark, dirname)\n if gt_metrics:\n gt_results = compute_image_multimodal_metrics(dataset, gen_dirpath, metrics=gt_multimodal_metrics)\n if \"image_reward\" in gt_other_metrics:\n gt_results.update(compute_image_reward(dataset, gen_dirpath))\n if benchmark in (\"COCO\", \"DCI\", \"MJHQ\"):\n gt_results.update(compute_image_similarity_metrics(dataset, gen_dirpath, metrics=gt_similarity_metrics))\n if \"fid\" in gt_other_metrics:\n gt_results[\"fid\"] = compute_fid(\n dataset,\n gen_dirpath,\n ref_cache_path=(os.path.join(gt_stats_root, benchmark, filename) if gt_stats_root else None),\n gen_cache_path=os.path.join(gen_root, \"fid_stats\", benchmark, filename),\n )\n benchmark_results[\"with_gt\"] = gt_results\n if ref_root and ref_metrics:\n assert os.path.exists(ref_root), f\"Reference root directory {ref_root} does not exist.\"\n ref_dirpath = os.path.join(ref_root, \"samples\", benchmark, dirname)\n ref_results = compute_image_similarity_metrics(ref_dirpath, gen_dirpath, metrics=ref_similarity_metrics)\n if \"fid\" in ref_other_metrics:\n ref_results[\"fid\"] = compute_fid(\n ref_dirpath,\n gen_dirpath,\n ref_cache_path=os.path.join(ref_root, \"fid_stats\", benchmark, filename),\n gen_cache_path=os.path.join(gen_root, \"fid_stats\", benchmark, filename),\n )\n benchmark_results[\"with_orig\"] = ref_results\n print(f\"{dirname} results:\")\n print(benchmark_results)\n results[dirname] = benchmark_results\n return results\n\n\ndef categorize_metrics(metrics: tuple[str, ...]) -> tuple[list[str], list[str], list[str]]:\n \"\"\"\n Categorize metrics into multimodal, similarity, and other metrics.\n\n Args:\n metrics (tuple[str, ...]): List of metrics.\n\n Returns:\n tuple[list[str], list[str], list[str]]: Tuple of multimodal, similarity, and other metrics.\n \"\"\"\n metrics = tuple(set(metrics))\n multimodal_metrics, similarity_metrics, other_metrics = [], [], []\n for metric in metrics:\n if metric in (\"clip_iqa\", \"clip_score\"):\n multimodal_metrics.append(metric)\n elif metric in (\"psnr\", \"lpips\", \"ssim\"):\n similarity_metrics.append(metric)\n else:\n other_metrics.append(metric)\n return multimodal_metrics, similarity_metrics, other_metrics\n"
},
{
"path": "deepcompressor/app/diffusion/eval/metrics/fid.py",
"content": "import os\nfrom datetime import datetime\n\nimport numpy as np\nimport torch\nimport torchvision\nfrom cleanfid import fid\nfrom cleanfid.resize import build_resizer\nfrom datasets import Dataset\nfrom tqdm import tqdm\n\n__all__ = [\"compute_fid\"]\n\n\ndef get_dataset_features(\n dataset: Dataset,\n model,\n mode: str = \"clean\",\n batch_size: int = 128,\n device: str | torch.device = \"cuda\",\n) -> np.ndarray:\n to_tensor = torchvision.transforms.ToTensor()\n fn_resize = build_resizer(mode)\n np_feats = []\n for batch in tqdm(\n dataset.iter(batch_size=batch_size, drop_last_batch=False),\n desc=f\"Extracting {dataset.config_name} features\",\n total=(len(dataset) + batch_size - 1) // batch_size,\n ):\n resized_images = [fn_resize(np.array(image.convert(\"RGB\"))) for image in batch[\"image\"]]\n image_tensors = []\n for resized_image in resized_images:\n if resized_image.dtype == \"uint8\":\n image_tensor = to_tensor(resized_image) * 255\n else:\n image_tensor = to_tensor(resized_image)\n image_tensors.append(image_tensor)\n image_tensors = torch.stack(image_tensors, dim=0)\n np_feats.append(fid.get_batch_features(image_tensors, model, device))\n np_feats = np.concatenate(np_feats, axis=0)\n return np_feats\n\n\ndef get_fid_features(\n dataset_or_folder: str | Dataset | None = None,\n cache_path: str | None = None,\n num: int | None = None,\n mode: str = \"clean\",\n num_workers: int = 8,\n batch_size: int = 64,\n device: str | torch.device = \"cuda\",\n force_overwrite: bool = False,\n verbose: bool = True,\n) -> tuple[np.ndarray, np.ndarray]:\n if cache_path is not None and os.path.exists(cache_path) and not force_overwrite:\n npz = np.load(cache_path)\n mu, sigma = npz[\"mu\"], npz[\"sigma\"]\n else:\n feat_model = fid.build_feature_extractor(mode, device)\n if isinstance(dataset_or_folder, str):\n np_feats = fid.get_folder_features(\n dataset_or_folder,\n feat_model,\n num_workers=num_workers,\n num=num,\n batch_size=batch_size,\n device=device,\n verbose=verbose,\n mode=mode,\n description=f\"Extracting {dataset_or_folder} features\",\n )\n else:\n assert isinstance(dataset_or_folder, Dataset)\n np_feats = get_dataset_features(\n dataset_or_folder, model=feat_model, mode=mode, batch_size=batch_size, device=device\n )\n\n mu = np.mean(np_feats, axis=0)\n sigma = np.cov(np_feats, rowvar=False)\n if cache_path is not None:\n os.makedirs(os.path.abspath(os.path.dirname(cache_path)), exist_ok=True)\n np.savez(cache_path, mu=mu, sigma=sigma)\n\n return mu, sigma\n\n\ndef compute_fid(\n ref_dirpath_or_dataset: str | Dataset,\n gen_dirpath: str,\n ref_cache_path: str | None = None,\n gen_cache_path: str | None = None,\n use_symlink: bool = True,\n timestamp: str | None = None,\n) -> float:\n sym_ref_dirpath, sym_gen_dirpath = None, None\n if use_symlink:\n if timestamp is None:\n timestamp = datetime.now().strftime(\"%y%m%d.%H%M%S\")\n\n os.makedirs(\".tmp\", exist_ok=True)\n\n if isinstance(ref_dirpath_or_dataset, str):\n sym_ref_dirpath = os.path.join(\".tmp\", f\"ref-{hash(str(ref_dirpath_or_dataset))}-{timestamp}\")\n os.symlink(os.path.abspath(ref_dirpath_or_dataset), os.path.abspath(sym_ref_dirpath))\n ref_dirpath_or_dataset = sym_ref_dirpath\n\n sym_gen_dirpath = os.path.join(\".tmp\", f\"gen-{hash(str(gen_dirpath))}-{timestamp}\")\n os.symlink(os.path.abspath(gen_dirpath), os.path.abspath(sym_gen_dirpath))\n gen_dirpath = sym_gen_dirpath\n mu1, sigma1 = get_fid_features(dataset_or_folder=ref_dirpath_or_dataset, cache_path=ref_cache_path)\n mu2, sigma2 = get_fid_features(dataset_or_folder=gen_dirpath, cache_path=gen_cache_path)\n fid_score = fid.frechet_distance(mu1, sigma1, mu2, sigma2)\n fid_score = float(fid_score)\n if use_symlink:\n if sym_ref_dirpath is not None:\n os.remove(sym_ref_dirpath)\n os.remove(sym_gen_dirpath)\n return fid_score\n"
},
{
"path": "deepcompressor/app/diffusion/eval/metrics/image_reward.py",
"content": "import os\n\nimport datasets\nimport torch\nfrom tqdm import tqdm\n\n__all__ = [\"compute_image_reward\"]\n\n\ndef compute_image_reward(\n ref_dataset: datasets.Dataset,\n gen_dirpath: str,\n) -> dict[str, float]:\n # import here to remove dependency on `ImageReward` git repo\n import ImageReward as RM\n\n scores = []\n model = RM.load(\"ImageReward-v1.0\")\n for batch in tqdm(\n ref_dataset.iter(batch_size=1, drop_last_batch=False),\n desc=f\"{ref_dataset.config_name} image reward\",\n total=len(ref_dataset),\n dynamic_ncols=True,\n ):\n filename = batch[\"filename\"][0]\n path = os.path.join(gen_dirpath, f\"{filename}.png\")\n prompt = batch[\"prompt\"][0]\n with torch.inference_mode():\n score = model.score(prompt, path)\n scores.append(score)\n result = {\"image_reward\": sum(scores) / len(scores)}\n return result\n"
},
{
"path": "deepcompressor/app/diffusion/eval/metrics/multimodal.py",
"content": "import os\n\nimport datasets\nimport numpy as np\nimport torch\nimport torchmetrics\nimport torchvision\nfrom PIL import Image\nfrom torch.utils import data\nfrom torchmetrics.multimodal import CLIPImageQualityAssessment, CLIPScore\nfrom tqdm import tqdm\n\n__all__ = [\"compute_image_multimodal_metrics\"]\n\n\nclass PromptImageDataset(data.Dataset):\n def __init__(self, ref_dataset: datasets.Dataset, gen_dirpath: str):\n super(data.Dataset, self).__init__()\n self.ref_dataset, self.gen_dirpath = ref_dataset, gen_dirpath\n self.transform = torchvision.transforms.ToTensor()\n\n def __len__(self):\n return len(self.ref_dataset)\n\n def __getitem__(self, idx: int):\n row = self.ref_dataset[idx]\n gen_image = Image.open(os.path.join(self.gen_dirpath, row[\"filename\"] + \".png\")).convert(\"RGB\")\n gen_tensor = torch.from_numpy(np.array(gen_image)).permute(2, 0, 1)\n prompt = row[\"prompt\"]\n return [gen_tensor, prompt]\n\n\ndef compute_image_multimodal_metrics(\n ref_dataset: datasets.Dataset,\n gen_dirpath: str,\n metrics: tuple[str, ...] = (\"clip_iqa\", \"clip_score\"),\n batch_size: int = 64,\n num_workers: int = 8,\n device: str | torch.device = \"cuda\",\n) -> dict[str, float]:\n if len(metrics) == 0:\n return {}\n metric_names = metrics\n metrics: dict[str, torchmetrics.Metric] = {}\n for metric_name in metric_names:\n if metric_name == \"clip_iqa\":\n metric = CLIPImageQualityAssessment(model_name_or_path=\"openai/clip-vit-large-patch14\").to(device)\n elif metric_name == \"clip_score\":\n metric = CLIPScore(model_name_or_path=\"openai/clip-vit-large-patch14\").to(device)\n else:\n raise NotImplementedError(f\"Metric {metric_name} is not implemented\")\n metrics[metric_name] = metric\n dataset = PromptImageDataset(ref_dataset, gen_dirpath)\n dataloader = data.DataLoader(\n dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False, drop_last=False\n )\n with torch.no_grad():\n for batch in tqdm(dataloader, desc=f\"{ref_dataset.config_name} multimodal metrics\"):\n batch[0] = batch[0].to(device)\n for metric_name, metric in metrics.items():\n if metric_name == \"clip_iqa\":\n metric.update(batch[0].to(torch.float32))\n else:\n prompts = list(batch[1])\n metric.update(batch[0], prompts)\n result = {metric_name: metric.compute().mean().item() for metric_name, metric in metrics.items()}\n return result\n"
},
{
"path": "deepcompressor/app/diffusion/eval/metrics/run.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Evaluate generated images or videos using the specified metrics.\"\"\"\n\nimport json\nimport os\n\nfrom ...config import DiffusionPtqRunConfig\n\nif __name__ == \"__main__\":\n config, _, unused_cfgs, unused_args, unknown_args = DiffusionPtqRunConfig.get_parser().parse_known_args()\n assert len(unknown_args) == 0, f\"Unknown arguments: {unknown_args}\"\n assert len(unused_cfgs) == 0, f\"Unused configurations: {unused_cfgs}\"\n assert unused_args is None, f\"Unused arguments: {unused_args}\"\n assert isinstance(config, DiffusionPtqRunConfig)\n results = config.eval.evaluate(pipeline=None, skip_gen=True, task=config.pipeline.task)\n save_path = os.path.join(config.eval.gen_root, f\"results-{config.output.timestamp}.json\")\n os.makedirs(os.path.abspath(os.path.dirname(save_path)), exist_ok=True)\n with open(save_path, \"w\") as f:\n json.dump(results, f, indent=2, sort_keys=True)\n print(results)\n"
},
{
"path": "deepcompressor/app/diffusion/eval/metrics/similarity.py",
"content": "import os\n\nimport datasets\nimport torch\nimport torchmetrics\nimport torchvision\nfrom PIL import Image\nfrom torch.utils import data\nfrom torchmetrics.image import (\n LearnedPerceptualImagePatchSimilarity,\n PeakSignalNoiseRatio,\n StructuralSimilarityIndexMeasure,\n)\nfrom tqdm import tqdm\n\n__all__ = [\"compute_image_similarity_metrics\"]\n\n\nclass MultiImageDataset(data.Dataset):\n def __init__(self, gen_dirpath: str, ref_dirpath_or_dataset: str | datasets.Dataset):\n super(data.Dataset, self).__init__()\n self.gen_names = sorted(\n [name for name in os.listdir(gen_dirpath) if name.endswith(\".png\") or name.endswith(\".jpg\")]\n )\n self.gen_dirpath, self.ref_dirpath_or_dataset = gen_dirpath, ref_dirpath_or_dataset\n if isinstance(ref_dirpath_or_dataset, str):\n self.ref_names = sorted(\n [name for name in os.listdir(ref_dirpath_or_dataset) if name.endswith(\".png\") or name.endswith(\".jpg\")]\n )\n assert len(self.ref_names) == len(self.gen_names)\n else:\n assert isinstance(ref_dirpath_or_dataset, datasets.Dataset)\n self.ref_names = self.gen_names\n assert len(ref_dirpath_or_dataset) == len(self.gen_names)\n self.transform = torchvision.transforms.ToTensor()\n\n def __len__(self):\n return len(self.ref_names)\n\n def __getitem__(self, idx: int):\n if isinstance(self.ref_dirpath_or_dataset, str):\n name = self.ref_names[idx]\n assert name == self.gen_names[idx]\n ref_image = Image.open(os.path.join(self.ref_dirpath_or_dataset, name)).convert(\"RGB\")\n else:\n row = self.ref_dirpath_or_dataset[idx]\n ref_image = row[\"image\"].convert(\"RGB\")\n name = row[\"filename\"] + \".png\"\n gen_image = Image.open(os.path.join(self.gen_dirpath, name)).convert(\"RGB\")\n gen_size = gen_image.size\n ref_size = ref_image.size\n if ref_size != gen_size:\n ref_image = ref_image.resize(gen_size, Image.Resampling.BICUBIC)\n gen_tensor = self.transform(gen_image)\n ref_tensor = self.transform(ref_image)\n return [gen_tensor, ref_tensor]\n\n\ndef compute_image_similarity_metrics(\n ref_dirpath_or_dataset: str | datasets.Dataset,\n gen_dirpath: str,\n metrics: tuple[str, ...] = (\"psnr\", \"lpips\", \"ssim\"),\n batch_size: int = 64,\n num_workers: int = 8,\n device: str | torch.device = \"cuda\",\n) -> dict[str, float]:\n if len(metrics) == 0:\n return {}\n metric_names = metrics\n metrics: dict[str, torchmetrics.Metric] = {}\n for metric_name in metric_names:\n if metric_name == \"psnr\":\n metric = PeakSignalNoiseRatio(data_range=(0, 1), reduction=\"elementwise_mean\", dim=(1, 2, 3)).to(device)\n elif metric_name == \"lpips\":\n metric = LearnedPerceptualImagePatchSimilarity(normalize=True).to(device)\n elif metric_name == \"ssim\":\n metric = StructuralSimilarityIndexMeasure(data_range=(0, 1)).to(device)\n else:\n raise NotImplementedError(f\"Metric {metric_name} is not implemented\")\n metrics[metric_name] = metric\n dataset = MultiImageDataset(gen_dirpath, ref_dirpath_or_dataset)\n dataloader = data.DataLoader(\n dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False, drop_last=False\n )\n with torch.no_grad():\n desc = (\n ref_dirpath_or_dataset.config_name\n if isinstance(ref_dirpath_or_dataset, datasets.Dataset)\n else os.path.basename(ref_dirpath_or_dataset)\n ) + \" similarity metrics\"\n for batch in tqdm(dataloader, desc=desc):\n batch = [tensor.to(device) for tensor in batch]\n for metric in metrics.values():\n metric.update(batch[0], batch[1])\n result = {metric_name: metric.compute().item() for metric_name, metric in metrics.items()}\n return result\n"
},
{
"path": "deepcompressor/app/diffusion/nn/__init__.py",
"content": "# -*- coding: utf-8 -*-\n"
},
{
"path": "deepcompressor/app/diffusion/nn/attention.py",
"content": "# -*- coding: utf-8 -*-\n\nimport typing as tp\n\nimport diffusers\nimport packaging.version\nimport torch\nimport torch.nn as nn\nfrom diffusers.models.attention_processor import (\n Attention,\n AttnProcessor2_0,\n FluxAttnProcessor2_0,\n JointAttnProcessor2_0,\n)\n\nfrom deepcompressor.nn.patch.sdpa import ScaleDotProductAttention\n\n__all__ = [\"DiffusionAttentionProcessor\"]\n\n\nif packaging.version.Version(diffusers.__version__) >= packaging.version.Version(\"0.31\"):\n from diffusers.models.embeddings import apply_rotary_emb\n\n def apply_flux_rope(query, key, image_rotary_emb):\n query = apply_rotary_emb(query, image_rotary_emb)\n key = apply_rotary_emb(key, image_rotary_emb)\n return query, key\n\nelse:\n from diffusers.models.attention_processor import apply_rope as apply_flux_rope\n\n\nclass DiffusionAttentionProcessor(nn.Module):\n def __init__(\n self,\n orig: AttnProcessor2_0 | FluxAttnProcessor2_0 | JointAttnProcessor2_0,\n sdpa: ScaleDotProductAttention | None = None,\n ) -> None:\n super().__init__()\n self.orig = orig\n if orig.__class__.__name__.startswith(\"Flux\"):\n self.rope = apply_flux_rope\n elif isinstance(orig, (AttnProcessor2_0, JointAttnProcessor2_0)):\n self.rope = None\n else:\n raise NotImplementedError(f\"Unsupported AttentionProcessor: {orig}\")\n self.sdpa = sdpa or ScaleDotProductAttention()\n\n def __call__( # noqa: C901\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: tp.Optional[torch.Tensor] = None,\n attention_mask: tp.Optional[torch.Tensor] = None,\n image_rotary_emb: tp.Optional[torch.Tensor] = None,\n *args,\n **kwargs,\n ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:\n assert len(args) == 0 and kwargs.get(\"scale\", None) is None\n assert attn.spatial_norm is None\n assert attn.group_norm is None\n assert attn.norm_cross is None\n assert not attn.residual_connection\n assert attn.rescale_output_factor == 1.0\n heads = attn.heads\n head_dim = attn.inner_dim // heads\n kv_heads = attn.inner_kv_dim // head_dim\n assert attn.scale == head_dim**-0.5\n\n input_ndim, input_shape = hidden_states.dim(), hidden_states.size()\n if input_ndim > 3:\n hidden_states = hidden_states.view(input_shape[0], input_shape[1], -1).transpose(1, 2)\n batch_size, input_length, _ = hidden_states.shape\n context_ndim, context_shape, context_length = None, None, None\n if encoder_hidden_states is not None:\n context_ndim, context_shape = encoder_hidden_states.ndim, encoder_hidden_states.shape\n assert context_shape[0] == batch_size\n if context_ndim > 3:\n encoder_hidden_states = encoder_hidden_states.view(batch_size, context_shape[1], -1).transpose(1, 2)\n context_length = encoder_hidden_states.shape[1]\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, context_length or input_length, batch_size)\n attention_mask = attention_mask.view(batch_size, heads, -1, attention_mask.shape[-1])\n\n query = attn.to_q(hidden_states)\n key, value, add_query, add_key, add_value = None, None, None, None, None\n if hasattr(attn, \"add_k_proj\"):\n if attn.to_k is not None:\n key = attn.to_k(hidden_states)\n value = attn.to_v(hidden_states)\n add_key = attn.add_k_proj(encoder_hidden_states)\n add_value = attn.add_v_proj(encoder_hidden_states)\n if hasattr(attn, \"add_q_proj\"):\n add_query = attn.add_q_proj(encoder_hidden_states)\n else:\n if attn.is_cross_attention:\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n else:\n assert encoder_hidden_states is None\n key = attn.to_k(hidden_states)\n value = attn.to_v(hidden_states)\n hidden_states, encoder_hidden_states = None, None\n\n query = query.view(batch_size, -1, heads, head_dim).transpose(1, 2)\n if key is not None:\n key = key.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2)\n if add_query is not None:\n add_query = add_query.view(batch_size, -1, heads, head_dim).transpose(1, 2)\n if add_key is not None:\n add_key = add_key.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2)\n add_value = add_value.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2)\n\n if kv_heads != heads:\n heads_per_kv_head = heads // kv_heads\n if key is not None:\n key = torch.repeat_interleave(key, heads_per_kv_head, dim=1)\n value = torch.repeat_interleave(value, heads_per_kv_head, dim=1)\n if add_key is not None:\n add_key = torch.repeat_interleave(add_key, heads_per_kv_head, dim=1)\n add_value = torch.repeat_interleave(add_value, heads_per_kv_head, dim=1)\n\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n key = attn.norm_k(key)\n if attn.norm_added_q is not None:\n add_query = attn.norm_added_q(add_query)\n add_key = attn.norm_added_k(add_key)\n\n if add_query is not None:\n query = torch.cat([add_query, query], dim=2)\n if add_key is not None:\n if key is None:\n key, value = add_key, add_value\n else:\n key = torch.cat([add_key, key], dim=2)\n value = torch.cat([add_value, value], dim=2)\n del add_query, add_key, add_value\n\n if image_rotary_emb is not None:\n query, key = self.rope(query, key, image_rotary_emb)\n\n hidden_states = self.sdpa(query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False)\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.inner_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n if hidden_states.shape[1] > input_length:\n encoder_hidden_states = hidden_states[:, :context_length]\n hidden_states = hidden_states[:, context_length:]\n\n if hasattr(attn, \"to_out\"):\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n if hasattr(attn, \"to_add_out\"):\n encoder_hidden_states = attn.to_add_out(encoder_hidden_states)\n\n if input_ndim > 3:\n hidden_states = hidden_states.transpose(-1, -2).reshape(input_shape)\n if encoder_hidden_states is not None and context_ndim > 3:\n assert encoder_hidden_states.ndim == 3\n encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(context_shape)\n\n if encoder_hidden_states is None:\n return hidden_states\n return hidden_states, encoder_hidden_states\n"
},
{
"path": "deepcompressor/app/diffusion/nn/patch.py",
"content": "import torch.nn as nn\nfrom diffusers.models.attention_processor import Attention\nfrom diffusers.models.transformers.transformer_flux import FluxSingleTransformerBlock\n\nfrom deepcompressor.nn.patch.conv import ConcatConv2d, ShiftedConv2d\nfrom deepcompressor.nn.patch.linear import ConcatLinear, ShiftedLinear\nfrom deepcompressor.utils import patch, tools\n\nfrom .attention import DiffusionAttentionProcessor\nfrom .struct import DiffusionFeedForwardStruct, DiffusionModelStruct, DiffusionResnetStruct, UNetStruct\n\n__all__ = [\n \"replace_up_block_conv_with_concat_conv\",\n \"replace_fused_linear_with_concat_linear\",\n \"replace_attn_processor\",\n \"shift_input_activations\",\n]\n\n\ndef replace_up_block_conv_with_concat_conv(model: nn.Module) -> None:\n \"\"\"Replace up_block convolutions in UNet with ConcatConv.\"\"\"\n model_struct = DiffusionModelStruct.construct(model)\n if not isinstance(model_struct, UNetStruct):\n return\n logger = tools.logging.getLogger(__name__)\n logger.info(\"Replacing up_block convolutions with ConcatConv.\")\n tools.logging.Formatter.indent_inc()\n parents_map = patch.get_module_parents_map(model)\n for up_block in model_struct.up_block_structs:\n logger.info(f\"+ Replacing convolutions in up_block {up_block.name}\")\n tools.logging.Formatter.indent_inc()\n for resnet in up_block.resnet_structs:\n assert len(resnet.convs[0]) == 1\n conv, conv_name = resnet.convs[0][0], resnet.conv_names[0][0]\n logger.info(f\"- Replacing {conv_name} in resnet {resnet.name}\")\n tools.logging.Formatter.indent_inc()\n if resnet.idx == 0:\n if up_block.idx == 0:\n prev_block = model_struct.mid_block_struct\n else:\n prev_block = model_struct.up_block_structs[up_block.idx - 1]\n logger.info(f\"+ using previous block {prev_block.name}\")\n prev_channels = prev_block.resnet_structs[-1].convs[-1][-1].out_channels\n else:\n prev_channels = up_block.resnet_structs[resnet.idx - 1].convs[-1][-1].out_channels\n logger.info(f\"+ conv_in_channels = {prev_channels}/{conv.in_channels}\")\n logger.info(f\"+ conv_out_channels = {conv.out_channels}\")\n concat_conv = ConcatConv2d.from_conv2d(conv, [prev_channels])\n for parent_name, parent_module, child_name in parents_map[conv]:\n logger.info(f\"+ replacing {child_name} in {parent_name}\")\n setattr(parent_module, child_name, concat_conv)\n tools.logging.Formatter.indent_dec()\n tools.logging.Formatter.indent_dec()\n tools.logging.Formatter.indent_dec()\n\n\ndef replace_fused_linear_with_concat_linear(model: nn.Module) -> None:\n \"\"\"Replace fused Linear in FluxSingleTransformerBlock with ConcatLinear.\"\"\"\n logger = tools.logging.getLogger(__name__)\n logger.info(\"Replacing fused Linear with ConcatLinear.\")\n tools.logging.Formatter.indent_inc()\n for name, module in model.named_modules():\n if isinstance(module, FluxSingleTransformerBlock):\n logger.info(f\"+ Replacing fused Linear in {name} with ConcatLinear.\")\n tools.logging.Formatter.indent_inc()\n logger.info(f\"- in_features = {module.proj_out.out_features}/{module.proj_out.in_features}\")\n logger.info(f\"- out_features = {module.proj_out.out_features}\")\n tools.logging.Formatter.indent_dec()\n module.proj_out = ConcatLinear.from_linear(module.proj_out, [module.proj_out.out_features])\n tools.logging.Formatter.indent_dec()\n\n\ndef shift_input_activations(model: nn.Module) -> None:\n \"\"\"Shift input activations of convolutions and linear layers if their lowerbound is negative.\n\n Args:\n model (nn.Module): model to shift input activations.\n \"\"\"\n logger = tools.logging.getLogger(__name__)\n model_struct = DiffusionModelStruct.construct(model)\n module_parents_map = patch.get_module_parents_map(model)\n logger.info(\"- Shifting input activations.\")\n tools.logging.Formatter.indent_inc()\n for _, module_name, module, parent, field_name in model_struct.named_key_modules():\n lowerbound = None\n if isinstance(parent, DiffusionResnetStruct) and field_name.startswith(\"conv\"):\n lowerbound = parent.config.intermediate_lowerbound\n elif isinstance(parent, DiffusionFeedForwardStruct) and field_name.startswith(\"down_proj\"):\n lowerbound = parent.config.intermediate_lowerbound\n if lowerbound is not None and lowerbound < 0:\n shift = -lowerbound\n logger.info(f\"+ Shifting input activations of {module_name} by {shift}\")\n tools.logging.Formatter.indent_inc()\n if isinstance(module, nn.Linear):\n shifted = ShiftedLinear.from_linear(module, shift=shift)\n shifted.linear.unsigned = True\n elif isinstance(module, nn.Conv2d):\n shifted = ShiftedConv2d.from_conv2d(module, shift=shift)\n shifted.conv.unsigned = True\n else:\n raise NotImplementedError(f\"Unsupported module type {type(module)}\")\n for parent_name, parent_module, child_name in module_parents_map[module]:\n logger.info(f\"+ Replacing {child_name} in {parent_name}\")\n setattr(parent_module, child_name, shifted)\n tools.logging.Formatter.indent_dec()\n tools.logging.Formatter.indent_dec()\n\n\ndef replace_attn_processor(model: nn.Module) -> None:\n \"\"\"Replace Attention processor with DiffusionAttentionProcessor.\"\"\"\n logger = tools.logging.getLogger(__name__)\n logger.info(\"Replacing Attention processors.\")\n tools.logging.Formatter.indent_inc()\n for name, module in model.named_modules():\n if isinstance(module, Attention):\n logger.info(f\"+ Replacing {name} processor with DiffusionAttentionProcessor.\")\n module.set_processor(DiffusionAttentionProcessor(module.processor))\n tools.logging.Formatter.indent_dec()\n"
},
{
"path": "deepcompressor/app/diffusion/nn/struct.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Utility functions for Diffusion Models.\"\"\"\n\nimport enum\nimport typing as tp\nfrom abc import abstractmethod\nfrom collections import OrderedDict, defaultdict\nfrom dataclasses import dataclass, field\n\n# region imports\nimport torch.nn as nn\nfrom diffusers.models.activations import GEGLU, GELU, ApproximateGELU, SwiGLU\nfrom diffusers.models.attention import BasicTransformerBlock, FeedForward, JointTransformerBlock\nfrom diffusers.models.attention_processor import Attention, SanaLinearAttnProcessor2_0\nfrom diffusers.models.embeddings import (\n CombinedTimestepGuidanceTextProjEmbeddings,\n CombinedTimestepTextProjEmbeddings,\n ImageHintTimeEmbedding,\n ImageProjection,\n ImageTimeEmbedding,\n PatchEmbed,\n PixArtAlphaTextProjection,\n TextImageProjection,\n TextImageTimeEmbedding,\n TextTimeEmbedding,\n TimestepEmbedding,\n)\nfrom diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormSingle, AdaLayerNormZero\nfrom diffusers.models.resnet import Downsample2D, ResnetBlock2D, Upsample2D\nfrom diffusers.models.transformers.pixart_transformer_2d import PixArtTransformer2DModel\nfrom diffusers.models.transformers.sana_transformer import GLUMBConv, SanaTransformer2DModel, SanaTransformerBlock\nfrom diffusers.models.transformers.transformer_2d import Transformer2DModel\nfrom diffusers.models.transformers.transformer_flux import (\n FluxSingleTransformerBlock,\n FluxTransformer2DModel,\n FluxTransformerBlock,\n)\nfrom diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel\nfrom diffusers.models.unets.unet_2d import UNet2DModel\nfrom diffusers.models.unets.unet_2d_blocks import (\n CrossAttnDownBlock2D,\n CrossAttnUpBlock2D,\n DownBlock2D,\n UNetMidBlock2D,\n UNetMidBlock2DCrossAttn,\n UpBlock2D,\n)\nfrom diffusers.models.unets.unet_2d_condition import UNet2DConditionModel\nfrom diffusers.pipelines import (\n FluxControlPipeline,\n FluxFillPipeline,\n FluxPipeline,\n PixArtAlphaPipeline,\n PixArtSigmaPipeline,\n SanaPipeline,\n StableDiffusion3Pipeline,\n StableDiffusionPipeline,\n StableDiffusionXLPipeline,\n)\n\nfrom deepcompressor.nn.patch.conv import ConcatConv2d, ShiftedConv2d\nfrom deepcompressor.nn.patch.linear import ConcatLinear, ShiftedLinear\nfrom deepcompressor.nn.struct.attn import (\n AttentionConfigStruct,\n AttentionStruct,\n BaseTransformerStruct,\n FeedForwardConfigStruct,\n FeedForwardStruct,\n TransformerBlockStruct,\n)\nfrom deepcompressor.nn.struct.base import BaseModuleStruct\nfrom deepcompressor.utils.common import join_name\n\nfrom .attention import DiffusionAttentionProcessor\n\n# endregion\n\n\n__all__ = [\"DiffusionModelStruct\", \"DiffusionBlockStruct\", \"DiffusionModelStruct\"]\n\n\nDIT_BLOCK_CLS = tp.Union[\n BasicTransformerBlock,\n JointTransformerBlock,\n FluxSingleTransformerBlock,\n FluxTransformerBlock,\n SanaTransformerBlock,\n]\nUNET_BLOCK_CLS = tp.Union[\n DownBlock2D,\n CrossAttnDownBlock2D,\n UNetMidBlock2D,\n UNetMidBlock2DCrossAttn,\n UpBlock2D,\n CrossAttnUpBlock2D,\n]\nDIT_CLS = tp.Union[\n Transformer2DModel,\n PixArtTransformer2DModel,\n SD3Transformer2DModel,\n FluxTransformer2DModel,\n SanaTransformer2DModel,\n]\nUNET_CLS = tp.Union[UNet2DModel, UNet2DConditionModel]\nMODEL_CLS = tp.Union[DIT_CLS, UNET_CLS]\nUNET_PIPELINE_CLS = tp.Union[StableDiffusionPipeline, StableDiffusionXLPipeline]\nDIT_PIPELINE_CLS = tp.Union[\n StableDiffusion3Pipeline,\n PixArtAlphaPipeline,\n PixArtSigmaPipeline,\n FluxPipeline,\n FluxControlPipeline,\n FluxFillPipeline,\n SanaPipeline,\n]\nPIPELINE_CLS = tp.Union[UNET_PIPELINE_CLS, DIT_PIPELINE_CLS]\n\n\n@dataclass(kw_only=True)\nclass DiffusionModuleStruct(BaseModuleStruct):\n def named_key_modules(self) -> tp.Generator[tuple[str, str, nn.Module, BaseModuleStruct, str], None, None]:\n if isinstance(self.module, (nn.Linear, nn.Conv2d)):\n yield self.key, self.name, self.module, self.parent, self.fname\n else:\n for name, module in self.module.named_modules():\n if name and isinstance(module, (nn.Linear, nn.Conv2d)):\n module_name = join_name(self.name, name)\n field_name = join_name(self.fname, name)\n yield self.key, module_name, module, self.parent, field_name\n\n\n@dataclass(kw_only=True)\nclass DiffusionBlockStruct(BaseModuleStruct):\n @abstractmethod\n def iter_attention_structs(self) -> tp.Generator[\"DiffusionAttentionStruct\", None, None]: ...\n\n @abstractmethod\n def iter_transformer_block_structs(self) -> tp.Generator[\"DiffusionTransformerBlockStruct\", None, None]: ...\n\n\n@dataclass(kw_only=True)\nclass DiffusionModelStruct(DiffusionBlockStruct):\n pre_module_structs: OrderedDict[str, DiffusionModuleStruct] = field(init=False, repr=False)\n post_module_structs: OrderedDict[str, DiffusionModuleStruct] = field(init=False, repr=False)\n\n @property\n @abstractmethod\n def num_blocks(self) -> int: ...\n\n @property\n @abstractmethod\n def block_structs(self) -> list[DiffusionBlockStruct]: ...\n\n @abstractmethod\n def get_prev_module_keys(self) -> tuple[str, ...]: ...\n\n @abstractmethod\n def get_post_module_keys(self) -> tuple[str, ...]: ...\n\n @abstractmethod\n def _get_iter_block_activations_args(\n self, **input_kwargs\n ) -> tuple[list[nn.Module], list[DiffusionModuleStruct | DiffusionBlockStruct], list[bool], list[bool]]: ...\n\n def _get_iter_pre_module_activations_args(\n self,\n ) -> tuple[list[nn.Module], list[DiffusionModuleStruct], list[bool], list[bool]]:\n layers, layer_structs, recomputes, use_prev_layer_outputs = [], [], [], []\n for layer_struct in self.pre_module_structs.values():\n layers.append(layer_struct.module)\n layer_structs.append(layer_struct)\n recomputes.append(False)\n use_prev_layer_outputs.append(False)\n return layers, layer_structs, recomputes, use_prev_layer_outputs\n\n def _get_iter_post_module_activations_args(\n self,\n ) -> tuple[list[nn.Module], list[DiffusionModuleStruct], list[bool], list[bool]]:\n layers, layer_structs, recomputes, use_prev_layer_outputs = [], [], [], []\n for layer_struct in self.post_module_structs.values():\n layers.append(layer_struct.module)\n layer_structs.append(layer_struct)\n recomputes.append(False)\n use_prev_layer_outputs.append(False)\n return layers, layer_structs, recomputes, use_prev_layer_outputs\n\n def get_iter_layer_activations_args(\n self, skip_pre_modules: bool, skip_post_modules: bool, **input_kwargs\n ) -> tuple[list[nn.Module], list[DiffusionModuleStruct | DiffusionBlockStruct], list[bool], list[bool]]:\n \"\"\"\n Get the arguments for iterating over the layers and their activations.\n\n Args:\n skip_pre_modules (`bool`):\n Whether to skip the pre-modules\n skip_post_modules (`bool`):\n Whether to skip the post-modules\n\n Returns:\n `tuple[list[nn.Module], list[DiffusionModuleStruct | DiffusionBlockStruct], list[bool], list[bool]]`:\n the layers, the layer structs, the recomputes, and the use_prev_layer_outputs\n \"\"\"\n layers, structs, recomputes, uses = [], [], [], []\n if not skip_pre_modules:\n layers, structs, recomputes, uses = self._get_iter_pre_module_activations_args()\n _layers, _structs, _recomputes, _uses = self._get_iter_block_activations_args(**input_kwargs)\n layers.extend(_layers)\n structs.extend(_structs)\n recomputes.extend(_recomputes)\n uses.extend(_uses)\n if not skip_post_modules:\n _layers, _structs, _recomputes, _uses = self._get_iter_post_module_activations_args()\n layers.extend(_layers)\n structs.extend(_structs)\n recomputes.extend(_recomputes)\n uses.extend(_uses)\n return layers, structs, recomputes, uses\n\n def named_key_modules(self) -> tp.Generator[tp.Tuple[str, str, nn.Module, BaseModuleStruct, str], None, None]:\n for module in self.pre_module_structs.values():\n yield from module.named_key_modules()\n for block in self.block_structs:\n yield from block.named_key_modules()\n for module in self.post_module_structs.values():\n yield from module.named_key_modules()\n\n def iter_attention_structs(self) -> tp.Generator[\"AttentionStruct\", None, None]:\n for block in self.block_structs:\n yield from block.iter_attention_structs()\n\n def iter_transformer_block_structs(self) -> tp.Generator[\"DiffusionTransformerBlockStruct\", None, None]:\n for block in self.block_structs:\n yield from block.iter_transformer_block_structs()\n\n def get_named_layers(\n self, skip_pre_modules: bool, skip_post_modules: bool, skip_blocks: bool = False\n ) -> OrderedDict[str, DiffusionBlockStruct | DiffusionModuleStruct]:\n named_layers = OrderedDict()\n if not skip_pre_modules:\n named_layers.update(self.pre_module_structs)\n if not skip_blocks:\n for block in self.block_structs:\n named_layers[block.name] = block\n if not skip_post_modules:\n named_layers.update(self.post_module_structs)\n return named_layers\n\n @staticmethod\n def _default_construct(\n module: tp.Union[PIPELINE_CLS, MODEL_CLS],\n /,\n parent: tp.Optional[BaseModuleStruct] = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"DiffusionModelStruct\":\n if isinstance(module, UNET_PIPELINE_CLS):\n module = module.unet\n elif isinstance(module, DIT_PIPELINE_CLS):\n module = module.transformer\n if isinstance(module, UNET_CLS):\n return UNetStruct.construct(module, parent=parent, fname=fname, rname=rname, rkey=rkey, idx=idx, **kwargs)\n elif isinstance(module, DIT_CLS):\n return DiTStruct.construct(module, parent=parent, fname=fname, rname=rname, rkey=rkey, idx=idx, **kwargs)\n raise NotImplementedError(f\"Unsupported module type: {type(module)}\")\n\n @classmethod\n def _get_default_key_map(cls) -> dict[str, set[str]]:\n unet_key_map = UNetStruct._get_default_key_map()\n dit_key_map = DiTStruct._get_default_key_map()\n flux_key_map = FluxStruct._get_default_key_map()\n key_map: dict[str, set[str]] = defaultdict(set)\n for rkey, keys in unet_key_map.items():\n key_map[rkey].update(keys)\n for rkey, keys in dit_key_map.items():\n key_map[rkey].update(keys)\n for rkey, keys in flux_key_map.items():\n key_map[rkey].update(keys)\n return {k: v for k, v in key_map.items() if v}\n\n @staticmethod\n def _simplify_keys(keys: tp.Iterable[str], *, key_map: dict[str, set[str]]) -> list[str]:\n \"\"\"Simplify the keys based on the key map.\n\n Args:\n keys (`Iterable[str]`):\n The keys to simplify.\n key_map (`dict[str, set[str]]`):\n The key map.\n\n Returns:\n `list[str]`:\n The simplified keys.\n \"\"\"\n # we first sort key_map by length of values in descending order\n key_map = dict(sorted(key_map.items(), key=lambda item: len(item[1]), reverse=True))\n ukeys, skeys = set(keys), set()\n for k, v in key_map.items():\n if k in ukeys:\n skeys.add(k)\n ukeys.discard(k)\n ukeys.difference_update(v)\n continue\n if ukeys.issuperset(v):\n skeys.add(k)\n ukeys.difference_update(v)\n assert not ukeys, f\"Unrecognized keys: {ukeys}\"\n return sorted(skeys)\n\n\n@dataclass(kw_only=True)\nclass DiffusionAttentionStruct(AttentionStruct):\n module: Attention = field(repr=False, kw_only=False)\n \"\"\"the module of AttentionBlock\"\"\"\n parent: tp.Optional[\"DiffusionTransformerBlockStruct\"] = field(repr=False)\n\n def filter_kwargs(self, kwargs: dict) -> dict:\n \"\"\"Filter layer kwargs to attn kwargs.\"\"\"\n if isinstance(self.parent.module, BasicTransformerBlock):\n if kwargs.get(\"cross_attention_kwargs\", None) is None:\n attn_kwargs = {}\n else:\n attn_kwargs = dict(kwargs[\"cross_attention_kwargs\"].items())\n attn_kwargs.pop(\"gligen\", None)\n if self.idx == 0:\n attn_kwargs[\"attention_mask\"] = kwargs.get(\"attention_mask\", None)\n else:\n attn_kwargs[\"attention_mask\"] = kwargs.get(\"encoder_attention_mask\", None)\n else:\n attn_kwargs = {}\n return attn_kwargs\n\n @staticmethod\n def _default_construct(\n module: Attention,\n /,\n parent: tp.Optional[\"DiffusionTransformerBlockStruct\"] = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"DiffusionAttentionStruct\":\n if module.is_cross_attention:\n q_proj, k_proj, v_proj = module.to_q, None, None\n add_q_proj, add_k_proj, add_v_proj, add_o_proj = None, module.to_k, module.to_v, None\n q_proj_rname, k_proj_rname, v_proj_rname = \"to_q\", \"\", \"\"\n add_q_proj_rname, add_k_proj_rname, add_v_proj_rname, add_o_proj_rname = \"\", \"to_k\", \"to_v\", \"\"\n else:\n q_proj, k_proj, v_proj = module.to_q, module.to_k, module.to_v\n add_q_proj = getattr(module, \"add_q_proj\", None)\n add_k_proj = getattr(module, \"add_k_proj\", None)\n add_v_proj = getattr(module, \"add_v_proj\", None)\n add_o_proj = getattr(module, \"to_add_out\", None)\n q_proj_rname, k_proj_rname, v_proj_rname = \"to_q\", \"to_k\", \"to_v\"\n add_q_proj_rname, add_k_proj_rname, add_v_proj_rname = \"add_q_proj\", \"add_k_proj\", \"add_v_proj\"\n add_o_proj_rname = \"to_add_out\"\n if getattr(module, \"to_out\", None) is not None:\n o_proj = module.to_out[0]\n o_proj_rname = \"to_out.0\"\n assert isinstance(o_proj, nn.Linear)\n elif parent is not None:\n assert isinstance(parent.module, FluxSingleTransformerBlock)\n assert isinstance(parent.module.proj_out, ConcatLinear)\n assert len(parent.module.proj_out.linears) == 2\n o_proj = parent.module.proj_out.linears[0]\n o_proj_rname = \".proj_out.linears.0\"\n else:\n raise RuntimeError(\"Cannot find the output projection.\")\n if isinstance(module.processor, DiffusionAttentionProcessor):\n with_rope = module.processor.rope is not None\n elif module.processor.__class__.__name__.startswith(\"Flux\"):\n with_rope = True\n else:\n with_rope = False # TODO: fix for other processors\n config = AttentionConfigStruct(\n hidden_size=q_proj.weight.shape[1],\n add_hidden_size=add_k_proj.weight.shape[1] if add_k_proj is not None else 0,\n inner_size=q_proj.weight.shape[0],\n num_query_heads=module.heads,\n num_key_value_heads=module.to_k.weight.shape[0] // (module.to_q.weight.shape[0] // module.heads),\n with_qk_norm=module.norm_q is not None,\n with_rope=with_rope,\n linear_attn=isinstance(module.processor, SanaLinearAttnProcessor2_0),\n )\n return DiffusionAttentionStruct(\n module=module,\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n config=config,\n q_proj=q_proj,\n k_proj=k_proj,\n v_proj=v_proj,\n o_proj=o_proj,\n add_q_proj=add_q_proj,\n add_k_proj=add_k_proj,\n add_v_proj=add_v_proj,\n add_o_proj=add_o_proj,\n q=None, # TODO: add q, k, v\n k=None,\n v=None,\n q_proj_rname=q_proj_rname,\n k_proj_rname=k_proj_rname,\n v_proj_rname=v_proj_rname,\n o_proj_rname=o_proj_rname,\n add_q_proj_rname=add_q_proj_rname,\n add_k_proj_rname=add_k_proj_rname,\n add_v_proj_rname=add_v_proj_rname,\n add_o_proj_rname=add_o_proj_rname,\n q_rname=\"\",\n k_rname=\"\",\n v_rname=\"\",\n )\n\n\n@dataclass(kw_only=True)\nclass DiffusionFeedForwardStruct(FeedForwardStruct):\n module: FeedForward = field(repr=False, kw_only=False)\n \"\"\"the module of FeedForward\"\"\"\n parent: tp.Optional[\"DiffusionTransformerBlockStruct\"] = field(repr=False)\n # region modules\n moe_gate: None = field(init=False, repr=False, default=None)\n experts: list[nn.Module] = field(init=False, repr=False)\n # endregion\n # region names\n moe_gate_rname: str = field(init=False, repr=False, default=\"\")\n experts_rname: str = field(init=False, repr=False, default=\"\")\n # endregion\n\n # region aliases\n\n @property\n def up_proj(self) -> nn.Linear:\n return self.up_projs[0]\n\n @property\n def down_proj(self) -> nn.Linear:\n return self.down_projs[0]\n\n @property\n def up_proj_rname(self) -> str:\n return self.up_proj_rnames[0]\n\n @property\n def down_proj_rname(self) -> str:\n return self.down_proj_rnames[0]\n\n @property\n def up_proj_name(self) -> str:\n return self.up_proj_names[0]\n\n @property\n def down_proj_name(self) -> str:\n return self.down_proj_names[0]\n\n # endregion\n\n def __post_init__(self) -> None:\n assert len(self.up_projs) == len(self.down_projs) == 1\n assert len(self.up_proj_rnames) == len(self.down_proj_rnames) == 1\n self.experts = [self.module]\n super().__post_init__()\n\n @staticmethod\n def _default_construct(\n module: FeedForward | FluxSingleTransformerBlock | GLUMBConv,\n /,\n parent: tp.Optional[\"DiffusionTransformerBlockStruct\"] = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"DiffusionFeedForwardStruct\":\n if isinstance(module, FeedForward):\n layer_1, layer_2 = module.net[0], module.net[2]\n assert isinstance(layer_1, (GEGLU, GELU, ApproximateGELU, SwiGLU))\n up_proj, up_proj_rname = layer_1.proj, \"net.0.proj\"\n assert isinstance(up_proj, nn.Linear)\n down_proj, down_proj_rname = layer_2, \"net.2\"\n if isinstance(layer_1, GEGLU):\n act_type = \"gelu_glu\"\n elif isinstance(layer_1, SwiGLU):\n act_type = \"swish_glu\"\n else:\n assert layer_1.__class__.__name__.lower().endswith(\"gelu\")\n act_type = \"gelu\"\n if isinstance(layer_2, ShiftedLinear):\n down_proj, down_proj_rname = layer_2.linear, \"net.2.linear\"\n act_type = \"gelu_shifted\"\n assert isinstance(down_proj, nn.Linear)\n ffn = module\n elif isinstance(module, FluxSingleTransformerBlock):\n up_proj, up_proj_rname = module.proj_mlp, \"proj_mlp\"\n act_type = \"gelu\"\n assert isinstance(module.proj_out, ConcatLinear)\n assert len(module.proj_out.linears) == 2\n layer_2 = module.proj_out.linears[1]\n if isinstance(layer_2, ShiftedLinear):\n down_proj, down_proj_rname = layer_2.linear, \"proj_out.linears.1.linear\"\n act_type = \"gelu_shifted\"\n else:\n down_proj, down_proj_rname = layer_2, \"proj_out.linears.1\"\n ffn = nn.Sequential(up_proj, module.act_mlp, layer_2)\n assert not rname, f\"Unsupported rname: {rname}\"\n elif isinstance(module, GLUMBConv):\n ffn = module\n up_proj, up_proj_rname = module.conv_inverted, \"conv_inverted\"\n down_proj, down_proj_rname = module.conv_point, \"conv_point\"\n act_type = \"silu_conv_silu_glu\"\n else:\n raise NotImplementedError(f\"Unsupported module type: {type(module)}\")\n config = FeedForwardConfigStruct(\n hidden_size=up_proj.weight.shape[1],\n intermediate_size=down_proj.weight.shape[1],\n intermediate_act_type=act_type,\n num_experts=1,\n )\n return DiffusionFeedForwardStruct(\n module=ffn, # this may be a virtual module\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n config=config,\n up_projs=[up_proj],\n down_projs=[down_proj],\n up_proj_rnames=[up_proj_rname],\n down_proj_rnames=[down_proj_rname],\n )\n\n\n@dataclass(kw_only=True)\nclass DiffusionTransformerBlockStruct(TransformerBlockStruct, DiffusionBlockStruct):\n # region relative keys\n norm_rkey: tp.ClassVar[str] = \"transformer_norm\"\n add_norm_rkey: tp.ClassVar[str] = \"transformer_add_norm\"\n attn_struct_cls: tp.ClassVar[type[DiffusionAttentionStruct]] = DiffusionAttentionStruct\n ffn_struct_cls: tp.ClassVar[type[DiffusionFeedForwardStruct]] = DiffusionFeedForwardStruct\n # endregion\n\n parent: tp.Optional[\"DiffusionTransformerStruct\"] = field(repr=False)\n # region child modules\n post_attn_norms: list[nn.LayerNorm] = field(init=False, repr=False, default_factory=list)\n post_attn_add_norms: list[nn.LayerNorm] = field(init=False, repr=False, default_factory=list)\n post_ffn_norm: None = field(init=False, repr=False, default=None)\n post_add_ffn_norm: None = field(init=False, repr=False, default=None)\n # endregion\n # region relative names\n post_attn_norm_rnames: list[str] = field(init=False, repr=False, default_factory=list)\n post_attn_add_norm_rnames: list[str] = field(init=False, repr=False, default_factory=list)\n post_ffn_norm_rname: str = field(init=False, repr=False, default=\"\")\n post_add_ffn_norm_rname: str = field(init=False, repr=False, default=\"\")\n # endregion\n # region attributes\n norm_type: str\n add_norm_type: str\n # endregion\n # region absolute keys\n norm_key: str = field(init=False, repr=False)\n add_norm_key: str = field(init=False, repr=False)\n # endregion\n # region child structs\n pre_attn_norm_structs: list[DiffusionModuleStruct | None] = field(init=False, repr=False)\n pre_attn_add_norm_structs: list[DiffusionModuleStruct | None] = field(init=False, repr=False)\n pre_ffn_norm_struct: DiffusionModuleStruct = field(init=False, repr=False, default=None)\n pre_add_ffn_norm_struct: DiffusionModuleStruct | None = field(init=False, repr=False, default=None)\n attn_structs: list[DiffusionAttentionStruct] = field(init=False, repr=False)\n ffn_struct: DiffusionFeedForwardStruct | None = field(init=False, repr=False)\n add_ffn_struct: DiffusionFeedForwardStruct | None = field(init=False, repr=False)\n # endregion\n\n def __post_init__(self) -> None:\n super().__post_init__()\n self.norm_key = join_name(self.key, self.norm_rkey, sep=\"_\")\n self.add_norm_key = join_name(self.key, self.add_norm_rkey, sep=\"_\")\n self.attn_norm_structs = [\n DiffusionModuleStruct(norm, parent=self, fname=\"pre_attn_norm\", rname=rname, rkey=self.norm_rkey, idx=idx)\n for idx, (norm, rname) in enumerate(zip(self.pre_attn_norms, self.pre_attn_norm_rnames, strict=True))\n ]\n self.add_attn_norm_structs = [\n DiffusionModuleStruct(\n norm, parent=self, fname=\"pre_attn_add_norm\", rname=rname, rkey=self.add_norm_rkey, idx=idx\n )\n for idx, (norm, rname) in enumerate(\n zip(self.pre_attn_add_norms, self.pre_attn_add_norm_rnames, strict=True)\n )\n ]\n if self.pre_ffn_norm is not None:\n self.pre_ffn_norm_struct = DiffusionModuleStruct(\n self.pre_ffn_norm, parent=self, fname=\"pre_ffn_norm\", rname=self.pre_ffn_norm_rname, rkey=self.norm_rkey\n )\n if self.pre_add_ffn_norm is not None:\n self.pre_add_ffn_norm_struct = DiffusionModuleStruct(\n self.pre_add_ffn_norm,\n parent=self,\n fname=\"pre_add_ffn_norm\",\n rname=self.pre_add_ffn_norm_rname,\n rkey=self.add_norm_rkey,\n )\n\n def named_key_modules(self) -> tp.Generator[tp.Tuple[str, str, nn.Module, BaseModuleStruct, str], None, None]:\n for attn_norm in self.attn_norm_structs:\n if attn_norm.module is not None:\n yield from attn_norm.named_key_modules()\n for add_attn_norm in self.add_attn_norm_structs:\n if add_attn_norm.module is not None:\n yield from add_attn_norm.named_key_modules()\n for attn_struct in self.attn_structs:\n yield from attn_struct.named_key_modules()\n if self.pre_ffn_norm_struct is not None:\n if self.pre_attn_norms and self.pre_attn_norms[0] is not self.pre_ffn_norm:\n yield from self.pre_ffn_norm_struct.named_key_modules()\n if self.ffn_struct is not None:\n yield from self.ffn_struct.named_key_modules()\n if self.pre_add_ffn_norm_struct is not None:\n if self.pre_attn_add_norms and self.pre_attn_add_norms[0] is not self.pre_add_ffn_norm:\n yield from self.pre_add_ffn_norm_struct.named_key_modules()\n if self.add_ffn_struct is not None:\n yield from self.add_ffn_struct.named_key_modules()\n\n @staticmethod\n def _default_construct(\n module: DIT_BLOCK_CLS,\n /,\n parent: tp.Optional[\"DiffusionTransformerStruct\"] = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"DiffusionTransformerBlockStruct\":\n if isinstance(module, (BasicTransformerBlock, SanaTransformerBlock)):\n parallel = False\n if isinstance(module, SanaTransformerBlock):\n norm_type = add_norm_type = \"ada_norm_single\"\n else:\n norm_type = add_norm_type = module.norm_type\n pre_attn_norms, pre_attn_norm_rnames = [], []\n attns, attn_rnames = [], []\n pre_attn_add_norms, pre_attn_add_norm_rnames = [], []\n assert module.norm1 is not None\n assert module.attn1 is not None\n pre_attn_norms.append(module.norm1)\n pre_attn_norm_rnames.append(\"norm1\")\n attns.append(module.attn1)\n attn_rnames.append(\"attn1\")\n pre_attn_add_norms.append(module.attn1.norm_cross)\n pre_attn_add_norm_rnames.append(\"attn1.norm_cross\")\n if module.attn2 is not None:\n if norm_type == \"ada_norm_single\":\n pre_attn_norms.append(None)\n pre_attn_norm_rnames.append(\"\")\n else:\n assert module.norm2 is not None\n pre_attn_norms.append(module.norm2)\n pre_attn_norm_rnames.append(\"norm2\")\n attns.append(module.attn2)\n attn_rnames.append(\"attn2\")\n pre_attn_add_norms.append(module.attn2.norm_cross)\n pre_attn_add_norm_rnames.append(\"attn2.norm_cross\")\n if norm_type == \"ada_norm_single\":\n assert module.norm2 is not None\n pre_ffn_norm, pre_ffn_norm_rname = module.norm2, \"norm2\"\n else:\n pre_ffn_norm, pre_ffn_norm_rname = module.norm3, \"\" if module.norm3 is None else \"norm3\"\n ffn, ffn_rname = module.ff, \"\" if module.ff is None else \"ff\"\n pre_add_ffn_norm, pre_add_ffn_norm_rname, add_ffn, add_ffn_rname = None, \"\", None, \"\"\n elif isinstance(module, JointTransformerBlock):\n parallel = False\n norm_type = \"ada_norm_zero\"\n pre_attn_norms, pre_attn_norm_rnames = [module.norm1], [\"norm1\"]\n if isinstance(module.norm1_context, AdaLayerNormZero):\n add_norm_type = \"ada_norm_zero\"\n else:\n add_norm_type = \"ada_norm_continous\"\n pre_attn_add_norms, pre_attn_add_norm_rnames = [module.norm1_context], [\"norm1_context\"]\n attns, attn_rnames = [module.attn], [\"attn\"]\n pre_ffn_norm, pre_ffn_norm_rname = module.norm2, \"norm2\"\n ffn, ffn_rname = module.ff, \"ff\"\n pre_add_ffn_norm, pre_add_ffn_norm_rname = module.norm2_context, \"norm2_context\"\n add_ffn, add_ffn_rname = module.ff_context, \"ff_context\"\n elif isinstance(module, FluxSingleTransformerBlock):\n parallel = True\n norm_type = add_norm_type = \"ada_norm_zero_single\"\n pre_attn_norms, pre_attn_norm_rnames = [module.norm], [\"norm\"]\n attns, attn_rnames = [module.attn], [\"attn\"]\n pre_attn_add_norms, pre_attn_add_norm_rnames = [], []\n pre_ffn_norm, pre_ffn_norm_rname = module.norm, \"norm\"\n ffn, ffn_rname = module, \"\"\n pre_add_ffn_norm, pre_add_ffn_norm_rname, add_ffn, add_ffn_rname = None, \"\", None, \"\"\n elif isinstance(module, FluxTransformerBlock):\n parallel = False\n norm_type = add_norm_type = \"ada_norm_zero\"\n pre_attn_norms, pre_attn_norm_rnames = [module.norm1], [\"norm1\"]\n attns, attn_rnames = [module.attn], [\"attn\"]\n pre_attn_add_norms, pre_attn_add_norm_rnames = [module.norm1_context], [\"norm1_context\"]\n pre_ffn_norm, pre_ffn_norm_rname = module.norm2, \"norm2\"\n ffn, ffn_rname = module.ff, \"ff\"\n pre_add_ffn_norm, pre_add_ffn_norm_rname = module.norm2_context, \"norm2_context\"\n add_ffn, add_ffn_rname = module.ff_context, \"ff_context\"\n else:\n raise NotImplementedError(f\"Unsupported module type: {type(module)}\")\n return DiffusionTransformerBlockStruct(\n module=module,\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n parallel=parallel,\n pre_attn_norms=pre_attn_norms,\n pre_attn_add_norms=pre_attn_add_norms,\n attns=attns,\n pre_ffn_norm=pre_ffn_norm,\n ffn=ffn,\n pre_add_ffn_norm=pre_add_ffn_norm,\n add_ffn=add_ffn,\n pre_attn_norm_rnames=pre_attn_norm_rnames,\n pre_attn_add_norm_rnames=pre_attn_add_norm_rnames,\n attn_rnames=attn_rnames,\n pre_ffn_norm_rname=pre_ffn_norm_rname,\n ffn_rname=ffn_rname,\n pre_add_ffn_norm_rname=pre_add_ffn_norm_rname,\n add_ffn_rname=add_ffn_rname,\n norm_type=norm_type,\n add_norm_type=add_norm_type,\n )\n\n @classmethod\n def _get_default_key_map(cls) -> dict[str, set[str]]:\n \"\"\"Get the default allowed keys.\"\"\"\n key_map: dict[str, set[str]] = defaultdict(set)\n norm_rkey = norm_key = cls.norm_rkey\n add_norm_rkey = add_norm_key = cls.add_norm_rkey\n key_map[norm_rkey].add(norm_key)\n key_map[add_norm_rkey].add(add_norm_key)\n attn_cls = cls.attn_struct_cls\n attn_key = attn_rkey = cls.attn_rkey\n qkv_proj_key = qkv_proj_rkey = join_name(attn_key, attn_cls.qkv_proj_rkey, sep=\"_\")\n out_proj_key = out_proj_rkey = join_name(attn_key, attn_cls.out_proj_rkey, sep=\"_\")\n add_qkv_proj_key = add_qkv_proj_rkey = join_name(attn_key, attn_cls.add_qkv_proj_rkey, sep=\"_\")\n add_out_proj_key = add_out_proj_rkey = join_name(attn_key, attn_cls.add_out_proj_rkey, sep=\"_\")\n key_map[attn_rkey].add(qkv_proj_key)\n key_map[attn_rkey].add(out_proj_key)\n if attn_cls.add_qkv_proj_rkey.startswith(\"add_\") and attn_cls.add_out_proj_rkey.startswith(\"add_\"):\n add_attn_rkey = join_name(attn_rkey, \"add\", sep=\"_\")\n key_map[add_attn_rkey].add(add_qkv_proj_key)\n key_map[add_attn_rkey].add(add_out_proj_key)\n key_map[qkv_proj_rkey].add(qkv_proj_key)\n key_map[out_proj_rkey].add(out_proj_key)\n key_map[add_qkv_proj_rkey].add(add_qkv_proj_key)\n key_map[add_out_proj_rkey].add(add_out_proj_key)\n ffn_cls = cls.ffn_struct_cls\n ffn_key = ffn_rkey = cls.ffn_rkey\n add_ffn_key = add_ffn_rkey = cls.add_ffn_rkey\n up_proj_key = up_proj_rkey = join_name(ffn_key, ffn_cls.up_proj_rkey, sep=\"_\")\n down_proj_key = down_proj_rkey = join_name(ffn_key, ffn_cls.down_proj_rkey, sep=\"_\")\n add_up_proj_key = add_up_proj_rkey = join_name(add_ffn_key, ffn_cls.up_proj_rkey, sep=\"_\")\n add_down_proj_key = add_down_proj_rkey = join_name(add_ffn_key, ffn_cls.down_proj_rkey, sep=\"_\")\n key_map[ffn_rkey].add(up_proj_key)\n key_map[ffn_rkey].add(down_proj_key)\n key_map[add_ffn_rkey].add(add_up_proj_key)\n key_map[add_ffn_rkey].add(add_down_proj_key)\n key_map[up_proj_rkey].add(up_proj_key)\n key_map[down_proj_rkey].add(down_proj_key)\n key_map[add_up_proj_rkey].add(add_up_proj_key)\n key_map[add_down_proj_rkey].add(add_down_proj_key)\n return {k: v for k, v in key_map.items() if v}\n\n\n@dataclass(kw_only=True)\nclass DiffusionTransformerStruct(BaseTransformerStruct, DiffusionBlockStruct):\n # region relative keys\n proj_in_rkey: tp.ClassVar[str] = \"transformer_proj_in\"\n proj_out_rkey: tp.ClassVar[str] = \"transformer_proj_out\"\n transformer_block_rkey: tp.ClassVar[str] = \"\"\n transformer_block_struct_cls: tp.ClassVar[type[DiffusionTransformerBlockStruct]] = DiffusionTransformerBlockStruct\n # endregion\n\n module: Transformer2DModel = field(repr=False, kw_only=False)\n # region modules\n norm_in: nn.GroupNorm | None\n \"\"\"Input normalization\"\"\"\n proj_in: nn.Linear | nn.Conv2d\n \"\"\"Input projection\"\"\"\n norm_out: nn.GroupNorm | None\n \"\"\"Output normalization\"\"\"\n proj_out: nn.Linear | nn.Conv2d\n \"\"\"Output projection\"\"\"\n transformer_blocks: nn.ModuleList = field(repr=False)\n \"\"\"Transformer blocks\"\"\"\n # endregion\n # region relative names\n transformer_blocks_rname: str\n # endregion\n # region absolute names\n transformer_blocks_name: str = field(init=False, repr=False)\n transformer_block_names: list[str] = field(init=False, repr=False)\n # endregion\n # region child structs\n transformer_block_structs: list[DiffusionTransformerBlockStruct] = field(init=False, repr=False)\n # endregion\n\n # region aliases\n\n @property\n def num_blocks(self) -> int:\n return len(self.transformer_blocks)\n\n @property\n def block_structs(self) -> list[DiffusionBlockStruct]:\n return self.transformer_block_structs\n\n @property\n def block_names(self) -> list[str]:\n return self.transformer_block_names\n\n # endregion\n\n def __post_init__(self):\n super().__post_init__()\n transformer_block_rnames = [\n f\"{self.transformer_blocks_rname}.{idx}\" for idx in range(len(self.transformer_blocks))\n ]\n self.transformer_blocks_name = join_name(self.name, self.transformer_blocks_rname)\n self.transformer_block_names = [join_name(self.name, rname) for rname in transformer_block_rnames]\n self.transformer_block_structs = [\n self.transformer_block_struct_cls.construct(\n layer,\n parent=self,\n fname=\"transformer_block\",\n rname=rname,\n rkey=self.transformer_block_rkey,\n idx=idx,\n )\n for idx, (layer, rname) in enumerate(zip(self.transformer_blocks, transformer_block_rnames, strict=True))\n ]\n\n @staticmethod\n def _default_construct(\n module: Transformer2DModel,\n /,\n parent: BaseModuleStruct = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"DiffusionTransformerStruct\":\n if isinstance(module, Transformer2DModel):\n assert module.is_input_continuous, \"input must be continuous\"\n transformer_blocks, transformer_blocks_rname = module.transformer_blocks, \"transformer_blocks\"\n norm_in, norm_in_rname = module.norm, \"norm\"\n proj_in, proj_in_rname = module.proj_in, \"proj_in\"\n proj_out, proj_out_rname = module.proj_out, \"proj_out\"\n norm_out, norm_out_rname = None, \"\"\n else:\n raise NotImplementedError(f\"Unsupported module type: {type(module)}\")\n return DiffusionTransformerStruct(\n module=module,\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n norm_in=norm_in,\n proj_in=proj_in,\n transformer_blocks=transformer_blocks,\n proj_out=proj_out,\n norm_out=norm_out,\n norm_in_rname=norm_in_rname,\n proj_in_rname=proj_in_rname,\n transformer_blocks_rname=transformer_blocks_rname,\n norm_out_rname=norm_out_rname,\n proj_out_rname=proj_out_rname,\n )\n\n @classmethod\n def _get_default_key_map(cls) -> dict[str, set[str]]:\n \"\"\"Get the default allowed keys.\"\"\"\n key_map: dict[str, set[str]] = defaultdict(set)\n proj_in_rkey = proj_in_key = cls.proj_in_rkey\n proj_out_rkey = proj_out_key = cls.proj_out_rkey\n key_map[proj_in_rkey].add(proj_in_key)\n key_map[proj_out_rkey].add(proj_out_key)\n block_cls = cls.transformer_block_struct_cls\n block_key = block_rkey = cls.transformer_block_rkey\n block_key_map = block_cls._get_default_key_map()\n for rkey, keys in block_key_map.items():\n rkey = join_name(block_rkey, rkey, sep=\"_\")\n for key in keys:\n key = join_name(block_key, key, sep=\"_\")\n key_map[rkey].add(key)\n return {k: v for k, v in key_map.items() if v}\n\n\n@dataclass(kw_only=True)\nclass DiffusionResnetStruct(BaseModuleStruct):\n # region relative keys\n conv_rkey: tp.ClassVar[str] = \"conv\"\n shortcut_rkey: tp.ClassVar[str] = \"shortcut\"\n time_proj_rkey: tp.ClassVar[str] = \"time_proj\"\n # endregion\n\n module: ResnetBlock2D = field(repr=False, kw_only=False)\n \"\"\"the module of Resnet\"\"\"\n config: FeedForwardConfigStruct\n # region child modules\n norms: list[nn.GroupNorm]\n convs: list[list[nn.Conv2d]]\n shortcut: nn.Conv2d | None\n time_proj: nn.Linear | None\n # endregion\n # region relative names\n norm_rnames: list[str]\n conv_rnames: list[list[str]]\n shortcut_rname: str\n time_proj_rname: str\n # endregion\n # region absolute names\n norm_names: list[str] = field(init=False, repr=False)\n conv_names: list[list[str]] = field(init=False, repr=False)\n shortcut_name: str = field(init=False, repr=False)\n time_proj_name: str = field(init=False, repr=False)\n # endregion\n # region absolute keys\n conv_key: str = field(init=False, repr=False)\n shortcut_key: str = field(init=False, repr=False)\n time_proj_key: str = field(init=False, repr=False)\n # endregion\n\n def __post_init__(self):\n super().__post_init__()\n self.norm_names = [join_name(self.name, rname) for rname in self.norm_rnames]\n self.conv_names = [[join_name(self.name, rname) for rname in rnames] for rnames in self.conv_rnames]\n self.shortcut_name = join_name(self.name, self.shortcut_rname)\n self.time_proj_name = join_name(self.name, self.time_proj_rname)\n self.conv_key = join_name(self.key, self.conv_rkey, sep=\"_\")\n self.shortcut_key = join_name(self.key, self.shortcut_rkey, sep=\"_\")\n self.time_proj_key = join_name(self.key, self.time_proj_rkey, sep=\"_\")\n\n def named_key_modules(self) -> tp.Generator[tp.Tuple[str, str, nn.Module, BaseModuleStruct, str], None, None]:\n for convs, names in zip(self.convs, self.conv_names, strict=True):\n for conv, name in zip(convs, names, strict=True):\n yield self.conv_key, name, conv, self, \"conv\"\n if self.shortcut is not None:\n yield self.shortcut_key, self.shortcut_name, self.shortcut, self, \"shortcut\"\n if self.time_proj is not None:\n yield self.time_proj_key, self.time_proj_name, self.time_proj, self, \"time_proj\"\n\n @staticmethod\n def construct(\n module: ResnetBlock2D,\n /,\n parent: BaseModuleStruct = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"DiffusionResnetStruct\":\n if isinstance(module, ResnetBlock2D):\n assert module.upsample is None, \"upsample must be None\"\n assert module.downsample is None, \"downsample must be None\"\n act_type = module.nonlinearity.__class__.__name__.lower()\n shifted = False\n if isinstance(module.conv1, ConcatConv2d):\n conv1_convs, conv1_names = [], []\n for conv_idx, conv in enumerate(module.conv1.convs):\n if isinstance(conv, ShiftedConv2d):\n shifted = True\n conv1_convs.append(conv.conv)\n conv1_names.append(f\"conv1.convs.{conv_idx}.conv\")\n else:\n assert isinstance(conv, nn.Conv2d)\n conv1_convs.append(conv)\n conv1_names.append(f\"conv1.convs.{conv_idx}\")\n elif isinstance(module.conv1, ShiftedConv2d):\n shifted = True\n conv1_convs = [module.conv1.conv]\n conv1_names = [\"conv1.conv\"]\n else:\n assert isinstance(module.conv1, nn.Conv2d)\n conv1_convs, conv1_names = [module.conv1], [\"conv1\"]\n if isinstance(module.conv2, ConcatConv2d):\n conv2_convs, conv2_names = [], []\n for conv_idx, conv in enumerate(module.conv2.convs):\n if isinstance(conv, ShiftedConv2d):\n shifted = True\n conv2_convs.append(conv.conv)\n conv2_names.append(f\"conv2.convs.{conv_idx}.conv\")\n else:\n assert isinstance(conv, nn.Conv2d)\n conv2_convs.append(conv)\n conv2_names.append(f\"conv2.convs.{conv_idx}\")\n elif isinstance(module.conv2, ShiftedConv2d):\n shifted = True\n conv2_convs = [module.conv2.conv]\n conv2_names = [\"conv2.conv\"]\n else:\n assert isinstance(module.conv2, nn.Conv2d)\n conv2_convs, conv2_names = [module.conv2], [\"conv2\"]\n convs, conv_rnames = [conv1_convs, conv2_convs], [conv1_names, conv2_names]\n norms, norm_rnames = [module.norm1, module.norm2], [\"norm1\", \"norm2\"]\n shortcut, shortcut_rname = module.conv_shortcut, \"\" if module.conv_shortcut is None else \"conv_shortcut\"\n time_proj, time_proj_rname = module.time_emb_proj, \"\" if module.time_emb_proj is None else \"time_emb_proj\"\n if shifted:\n assert all(hasattr(conv, \"shifted\") and conv.shifted for level_convs in convs for conv in level_convs)\n act_type += \"_shifted\"\n else:\n raise NotImplementedError(f\"Unsupported module type: {type(module)}\")\n config = FeedForwardConfigStruct(\n hidden_size=convs[0][0].weight.shape[1],\n intermediate_size=convs[0][0].weight.shape[0],\n intermediate_act_type=act_type,\n num_experts=1,\n )\n return DiffusionResnetStruct(\n module=module,\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n config=config,\n norms=norms,\n convs=convs,\n shortcut=shortcut,\n time_proj=time_proj,\n norm_rnames=norm_rnames,\n conv_rnames=conv_rnames,\n shortcut_rname=shortcut_rname,\n time_proj_rname=time_proj_rname,\n )\n\n @classmethod\n def _get_default_key_map(cls) -> dict[str, set[str]]:\n \"\"\"Get the default allowed keys.\"\"\"\n key_map: dict[str, set[str]] = defaultdict(set)\n conv_key = conv_rkey = cls.conv_rkey\n shortcut_key = shortcut_rkey = cls.shortcut_rkey\n time_proj_key = time_proj_rkey = cls.time_proj_rkey\n key_map[conv_rkey].add(conv_key)\n key_map[shortcut_rkey].add(shortcut_key)\n key_map[time_proj_rkey].add(time_proj_key)\n return {k: v for k, v in key_map.items() if v}\n\n\n@dataclass(kw_only=True)\nclass UNetBlockStruct(DiffusionBlockStruct):\n class BlockType(enum.StrEnum):\n DOWN = \"down\"\n MID = \"mid\"\n UP = \"up\"\n\n # region relative keys\n resnet_rkey: tp.ClassVar[str] = \"resblock\"\n sampler_rkey: tp.ClassVar[str] = \"sample\"\n transformer_rkey: tp.ClassVar[str] = \"\"\n resnet_struct_cls: tp.ClassVar[type[DiffusionResnetStruct]] = DiffusionResnetStruct\n transformer_struct_cls: tp.ClassVar[type[DiffusionTransformerStruct]] = DiffusionTransformerStruct\n # endregion\n\n parent: tp.Optional[\"UNetStruct\"] = field(repr=False)\n # region attributes\n block_type: BlockType\n # endregion\n # region modules\n resnets: nn.ModuleList = field(repr=False)\n transformers: nn.ModuleList = field(repr=False)\n sampler: nn.Conv2d | None\n # endregion\n # region relative names\n resnets_rname: str\n transformers_rname: str\n sampler_rname: str\n # endregion\n # region absolute names\n resnets_name: str = field(init=False, repr=False)\n transformers_name: str = field(init=False, repr=False)\n sampler_name: str = field(init=False, repr=False)\n resnet_names: list[str] = field(init=False, repr=False)\n transformer_names: list[str] = field(init=False, repr=False)\n # endregion\n # region absolute keys\n sampler_key: str = field(init=False, repr=False)\n # endregion\n # region child structs\n resnet_structs: list[DiffusionResnetStruct] = field(init=False, repr=False)\n transformer_structs: list[DiffusionTransformerStruct] = field(init=False, repr=False)\n # endregion\n\n @property\n def downsample(self) -> nn.Conv2d | None:\n return self.sampler if self.is_downsample_block() else None\n\n @property\n def upsample(self) -> nn.Conv2d | None:\n return self.sampler if self.is_upsample_block() else None\n\n def __post_init__(self) -> None:\n super().__post_init__()\n if self.is_downsample_block():\n assert len(self.resnets) == len(self.transformers) or len(self.transformers) == 0\n if self.parent is not None and isinstance(self.parent, UNetStruct):\n assert self.rname == f\"{self.parent.down_blocks_rname}.{self.idx}\"\n elif self.is_mid_block():\n assert len(self.resnets) == len(self.transformers) + 1 or len(self.transformers) == 0\n if self.parent is not None and isinstance(self.parent, UNetStruct):\n assert self.rname == self.parent.mid_block_name\n assert self.idx == 0\n else:\n assert self.is_upsample_block(), f\"Unsupported block type: {self.block_type}\"\n assert len(self.resnets) == len(self.transformers) or len(self.transformers) == 0\n if self.parent is not None and isinstance(self.parent, UNetStruct):\n assert self.rname == f\"{self.parent.up_blocks_rname}.{self.idx}\"\n resnet_rnames = [f\"{self.resnets_rname}.{idx}\" for idx in range(len(self.resnets))]\n transformer_rnames = [f\"{self.transformers_rname}.{idx}\" for idx in range(len(self.transformers))]\n self.resnets_name = join_name(self.name, self.resnets_rname)\n self.transformers_name = join_name(self.name, self.transformers_rname)\n self.resnet_names = [join_name(self.name, rname) for rname in resnet_rnames]\n self.transformer_names = [join_name(self.name, rname) for rname in transformer_rnames]\n self.sampler_name = join_name(self.name, self.sampler_rname)\n self.sampler_key = join_name(self.key, self.sampler_rkey, sep=\"_\")\n self.resnet_structs = [\n self.resnet_struct_cls.construct(\n resnet, parent=self, fname=\"resnet\", rname=rname, rkey=self.resnet_rkey, idx=idx\n )\n for idx, (resnet, rname) in enumerate(zip(self.resnets, resnet_rnames, strict=True))\n ]\n self.transformer_structs = [\n self.transformer_struct_cls.construct(\n transformer, parent=self, fname=\"transformer\", rname=rname, rkey=self.transformer_rkey, idx=idx\n )\n for idx, (transformer, rname) in enumerate(zip(self.transformers, transformer_rnames, strict=True))\n ]\n\n def is_downsample_block(self) -> bool:\n return self.block_type == self.BlockType.DOWN\n\n def is_mid_block(self) -> bool:\n return self.block_type == self.BlockType.MID\n\n def is_upsample_block(self) -> bool:\n return self.block_type == self.BlockType.UP\n\n def has_downsample(self) -> bool:\n return self.is_downsample_block() and self.sampler is not None\n\n def has_upsample(self) -> bool:\n return self.is_upsample_block() and self.sampler is not None\n\n def named_key_modules(self) -> tp.Generator[tp.Tuple[str, str, nn.Module, BaseModuleStruct, str], None, None]:\n for resnet in self.resnet_structs:\n yield from resnet.named_key_modules()\n for transformer in self.transformer_structs:\n yield from transformer.named_key_modules()\n if self.sampler is not None:\n yield self.sampler_key, self.sampler_name, self.sampler, self, \"sampler\"\n\n def iter_attention_structs(self) -> tp.Generator[DiffusionAttentionStruct, None, None]:\n for transformer in self.transformer_structs:\n yield from transformer.iter_attention_structs()\n\n def iter_transformer_block_structs(self) -> tp.Generator[DiffusionTransformerBlockStruct, None, None]:\n for transformer in self.transformer_structs:\n yield from transformer.iter_transformer_block_structs()\n\n @staticmethod\n def _default_construct(\n module: UNET_BLOCK_CLS,\n /,\n parent: tp.Optional[\"UNetStruct\"] = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"UNetBlockStruct\":\n resnets, resnets_rname = module.resnets, \"resnets\"\n if isinstance(module, (DownBlock2D, CrossAttnDownBlock2D)):\n block_type = UNetBlockStruct.BlockType.DOWN\n if isinstance(module, CrossAttnDownBlock2D) and module.attentions is not None:\n transformers, transformers_rname = module.attentions, \"attentions\"\n else:\n transformers, transformers_rname = [], \"\"\n if module.downsamplers is None:\n sampler, sampler_rname = None, \"\"\n else:\n assert len(module.downsamplers) == 1\n downsampler = module.downsamplers[0]\n assert isinstance(downsampler, Downsample2D)\n sampler, sampler_rname = downsampler.conv, \"downsamplers.0.conv\"\n assert isinstance(sampler, nn.Conv2d)\n elif isinstance(module, (UNetMidBlock2D, UNetMidBlock2DCrossAttn)):\n block_type = UNetBlockStruct.BlockType.MID\n if (isinstance(module, UNetMidBlock2DCrossAttn) or module.add_attention) and module.attentions is not None:\n transformers, transformers_rname = module.attentions, \"attentions\"\n else:\n transformers, transformers_rname = [], \"\"\n sampler, sampler_rname = None, \"\"\n elif isinstance(module, (UpBlock2D, CrossAttnUpBlock2D)):\n block_type = UNetBlockStruct.BlockType.UP\n if isinstance(module, CrossAttnUpBlock2D) and module.attentions is not None:\n transformers, transformers_rname = module.attentions, \"attentions\"\n else:\n transformers, transformers_rname = [], \"\"\n if module.upsamplers is None:\n sampler, sampler_rname = None, \"\"\n else:\n assert len(module.upsamplers) == 1\n upsampler = module.upsamplers[0]\n assert isinstance(upsampler, Upsample2D)\n sampler, sampler_rname = upsampler.conv, \"upsamplers.0.conv\"\n assert isinstance(sampler, nn.Conv2d)\n else:\n raise NotImplementedError(f\"Unsupported module type: {type(module)}\")\n return UNetBlockStruct(\n module=module,\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n block_type=block_type,\n resnets=resnets,\n transformers=transformers,\n sampler=sampler,\n resnets_rname=resnets_rname,\n transformers_rname=transformers_rname,\n sampler_rname=sampler_rname,\n )\n\n @classmethod\n def _get_default_key_map(cls) -> dict[str, set[str]]:\n \"\"\"Get the default allowed keys.\"\"\"\n key_map: dict[str, set[str]] = defaultdict(set)\n resnet_cls = cls.resnet_struct_cls\n resnet_key = resnet_rkey = cls.resnet_rkey\n resnet_key_map = resnet_cls._get_default_key_map()\n for rkey, keys in resnet_key_map.items():\n rkey = join_name(resnet_rkey, rkey, sep=\"_\")\n for key in keys:\n key = join_name(resnet_key, key, sep=\"_\")\n key_map[rkey].add(key)\n key_map[resnet_rkey].add(key)\n transformer_cls = cls.transformer_struct_cls\n transformer_key = transformer_rkey = cls.transformer_rkey\n transformer_key_map = transformer_cls._get_default_key_map()\n for rkey, keys in transformer_key_map.items():\n trkey = join_name(transformer_rkey, rkey, sep=\"_\")\n for key in keys:\n key = join_name(transformer_key, key, sep=\"_\")\n key_map[rkey].add(key)\n key_map[trkey].add(key)\n return {k: v for k, v in key_map.items() if v}\n\n\n@dataclass(kw_only=True)\nclass UNetStruct(DiffusionModelStruct):\n # region relative keys\n input_embed_rkey: tp.ClassVar[str] = \"input_embed\"\n \"\"\"hidden_states = input_embed(hidden_states), e.g., conv_in\"\"\"\n time_embed_rkey: tp.ClassVar[str] = \"time_embed\"\n \"\"\"temb = time_embed(timesteps, hidden_states)\"\"\"\n add_time_embed_rkey: tp.ClassVar[str] = \"time_embed\"\n \"\"\"add_temb = add_time_embed(timesteps, encoder_hidden_states)\"\"\"\n text_embed_rkey: tp.ClassVar[str] = \"text_embed\"\n \"\"\"encoder_hidden_states = text_embed(encoder_hidden_states)\"\"\"\n norm_out_rkey: tp.ClassVar[str] = \"output_embed\"\n \"\"\"hidden_states = norm_out(hidden_states), e.g., conv_norm_out\"\"\"\n proj_out_rkey: tp.ClassVar[str] = \"output_embed\"\n \"\"\"hidden_states = output_embed(hidden_states), e.g., conv_out\"\"\"\n down_block_rkey: tp.ClassVar[str] = \"down\"\n mid_block_rkey: tp.ClassVar[str] = \"mid\"\n up_block_rkey: tp.ClassVar[str] = \"up\"\n down_block_struct_cls: tp.ClassVar[type[UNetBlockStruct]] = UNetBlockStruct\n mid_block_struct_cls: tp.ClassVar[type[UNetBlockStruct]] = UNetBlockStruct\n up_block_struct_cls: tp.ClassVar[type[UNetBlockStruct]] = UNetBlockStruct\n # endregion\n\n # region child modules\n # region pre-block modules\n input_embed: nn.Conv2d\n time_embed: TimestepEmbedding\n \"\"\"Time embedding\"\"\"\n add_time_embed: (\n TextTimeEmbedding\n | TextImageTimeEmbedding\n | TimestepEmbedding\n | ImageTimeEmbedding\n | ImageHintTimeEmbedding\n | None\n )\n \"\"\"Additional time embedding\"\"\"\n text_embed: nn.Linear | ImageProjection | TextImageProjection | None\n \"\"\"Text embedding\"\"\"\n # region post-block modules\n norm_out: nn.GroupNorm | None\n proj_out: nn.Conv2d\n # endregion\n # endregion\n down_blocks: nn.ModuleList = field(repr=False)\n mid_block: nn.Module = field(repr=False)\n up_blocks: nn.ModuleList = field(repr=False)\n # endregion\n # region relative names\n input_embed_rname: str\n time_embed_rname: str\n add_time_embed_rname: str\n text_embed_rname: str\n norm_out_rname: str\n proj_out_rname: str\n down_blocks_rname: str\n mid_block_rname: str\n up_blocks_rname: str\n # endregion\n # region absolute names\n input_embed_name: str = field(init=False, repr=False)\n time_embed_name: str = field(init=False, repr=False)\n add_time_embed_name: str = field(init=False, repr=False)\n text_embed_name: str = field(init=False, repr=False)\n norm_out_name: str = field(init=False, repr=False)\n proj_out_name: str = field(init=False, repr=False)\n down_blocks_name: str = field(init=False, repr=False)\n mid_block_name: str = field(init=False, repr=False)\n up_blocks_name: str = field(init=False, repr=False)\n down_block_names: list[str] = field(init=False, repr=False)\n up_block_names: list[str] = field(init=False, repr=False)\n # endregion\n # region absolute keys\n input_embed_key: str = field(init=False, repr=False)\n time_embed_key: str = field(init=False, repr=False)\n add_time_embed_key: str = field(init=False, repr=False)\n text_embed_key: str = field(init=False, repr=False)\n norm_out_key: str = field(init=False, repr=False)\n proj_out_key: str = field(init=False, repr=False)\n # endregion\n # region child structs\n down_block_structs: list[UNetBlockStruct] = field(init=False, repr=False)\n mid_block_struct: UNetBlockStruct = field(init=False, repr=False)\n up_block_structs: list[UNetBlockStruct] = field(init=False, repr=False)\n # endregion\n\n @property\n def num_down_blocks(self) -> int:\n return len(self.down_blocks)\n\n @property\n def num_up_blocks(self) -> int:\n return len(self.up_blocks)\n\n @property\n def num_blocks(self) -> int:\n return self.num_down_blocks + 1 + self.num_up_blocks\n\n @property\n def block_structs(self) -> list[UNetBlockStruct]:\n return [*self.down_block_structs, self.mid_block_struct, *self.up_block_structs]\n\n def __post_init__(self) -> None:\n super().__post_init__()\n down_block_rnames = [f\"{self.down_blocks_rname}.{idx}\" for idx in range(len(self.down_blocks))]\n up_block_rnames = [f\"{self.up_blocks_rname}.{idx}\" for idx in range(len(self.up_blocks))]\n self.down_blocks_name = join_name(self.name, self.down_blocks_rname)\n self.mid_block_name = join_name(self.name, self.mid_block_rname)\n self.up_blocks_name = join_name(self.name, self.up_blocks_rname)\n self.down_block_names = [join_name(self.name, rname) for rname in down_block_rnames]\n self.up_block_names = [join_name(self.name, rname) for rname in up_block_rnames]\n self.pre_module_structs = {}\n for fname in (\"time_embed\", \"add_time_embed\", \"text_embed\", \"input_embed\"):\n module, rname, rkey = getattr(self, fname), getattr(self, f\"{fname}_rname\"), getattr(self, f\"{fname}_rkey\")\n setattr(self, f\"{fname}_key\", join_name(self.key, rkey, sep=\"_\"))\n if module is not None or rname:\n setattr(self, f\"{fname}_name\", join_name(self.name, rname))\n else:\n setattr(self, f\"{fname}_name\", \"\")\n if module is not None:\n assert rname, f\"rname of {fname} must not be empty\"\n self.pre_module_structs[getattr(self, f\"{fname}_name\")] = DiffusionModuleStruct(\n module=module, parent=self, fname=fname, rname=rname, rkey=rkey\n )\n self.post_module_structs = {}\n for fname in (\"norm_out\", \"proj_out\"):\n module, rname, rkey = getattr(self, fname), getattr(self, f\"{fname}_rname\"), getattr(self, f\"{fname}_rkey\")\n setattr(self, f\"{fname}_key\", join_name(self.key, rkey, sep=\"_\"))\n if module is not None or rname:\n setattr(self, f\"{fname}_name\", join_name(self.name, rname))\n else:\n setattr(self, f\"{fname}_name\", \"\")\n if module is not None:\n self.post_module_structs[getattr(self, f\"{fname}_name\")] = DiffusionModuleStruct(\n module=module, parent=self, fname=fname, rname=rname, rkey=rkey\n )\n self.down_block_structs = [\n self.down_block_struct_cls.construct(\n block, parent=self, fname=\"down_block\", rname=rname, rkey=self.down_block_rkey, idx=idx\n )\n for idx, (block, rname) in enumerate(zip(self.down_blocks, down_block_rnames, strict=True))\n ]\n self.mid_block_struct = self.mid_block_struct_cls.construct(\n self.mid_block, parent=self, fname=\"mid_block\", rname=self.mid_block_name, rkey=self.mid_block_rkey\n )\n self.up_block_structs = [\n self.up_block_struct_cls.construct(\n block, parent=self, fname=\"up_block\", rname=rname, rkey=self.up_block_rkey, idx=idx\n )\n for idx, (block, rname) in enumerate(zip(self.up_blocks, up_block_rnames, strict=True))\n ]\n\n def get_prev_module_keys(self) -> tuple[str, ...]:\n return tuple({self.input_embed_key, self.time_embed_key, self.add_time_embed_key, self.text_embed_key})\n\n def get_post_module_keys(self) -> tuple[str, ...]:\n return tuple({self.norm_out_key, self.proj_out_key})\n\n def _get_iter_block_activations_args(\n self, **input_kwargs\n ) -> tuple[list[nn.Module], list[DiffusionModuleStruct | DiffusionBlockStruct], list[bool], list[bool]]:\n layers, layer_structs, recomputes, use_prev_layer_outputs = [], [], [], []\n num_down_blocks = len(self.down_blocks)\n num_up_blocks = len(self.up_blocks)\n layers.extend(self.down_blocks)\n layer_structs.extend(self.down_block_structs)\n use_prev_layer_outputs.append(False)\n use_prev_layer_outputs.extend([True] * (num_down_blocks - 1))\n recomputes.append(False)\n # region check whether down block's outputs are changed\n _mid_block_additional_residual = input_kwargs.get(\"mid_block_additional_residual\", None)\n _down_block_additional_residuals = input_kwargs.get(\"down_block_additional_residuals\", None)\n _is_adapter = input_kwargs.get(\"down_intrablock_additional_residuals\", None) is not None\n if not _is_adapter and _mid_block_additional_residual is None and _down_block_additional_residuals is not None:\n _is_adapter = True\n for down_block in self.down_blocks:\n if hasattr(down_block, \"has_cross_attention\") and down_block.has_cross_attention:\n # outputs unchanged\n recomputes.append(False)\n elif _is_adapter:\n # outputs changed\n recomputes.append(True)\n else:\n # outputs unchanged\n recomputes.append(False)\n # endregion\n layers.append(self.mid_block)\n layer_structs.append(self.mid_block_struct)\n use_prev_layer_outputs.append(False)\n # recomputes is already appened in the previous down blocks\n layers.extend(self.up_blocks)\n layer_structs.extend(self.up_block_structs)\n use_prev_layer_outputs.append(False)\n use_prev_layer_outputs.extend([True] * (num_up_blocks - 1))\n recomputes += [True] * num_up_blocks\n return layers, layer_structs, recomputes, use_prev_layer_outputs\n\n @staticmethod\n def _default_construct(\n module: tp.Union[UNET_PIPELINE_CLS, UNET_CLS],\n /,\n parent: tp.Optional[BaseModuleStruct] = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"UNetStruct\":\n if isinstance(module, UNET_PIPELINE_CLS):\n module = module.unet\n if isinstance(module, (UNet2DConditionModel, UNet2DModel)):\n input_embed, time_embed = module.conv_in, module.time_embedding\n input_embed_rname, time_embed_rname = \"conv_in\", \"time_embedding\"\n text_embed, text_embed_rname = None, \"\"\n add_time_embed, add_time_embed_rname = None, \"\"\n if hasattr(module, \"encoder_hid_proj\"):\n text_embed, text_embed_rname = module.encoder_hid_proj, \"encoder_hid_proj\"\n if hasattr(module, \"add_embedding\"):\n add_time_embed, add_time_embed_rname = module.add_embedding, \"add_embedding\"\n norm_out, norm_out_rname = module.conv_norm_out, \"conv_norm_out\"\n proj_out, proj_out_rname = module.conv_out, \"conv_out\"\n down_blocks, down_blocks_rname = module.down_blocks, \"down_blocks\"\n mid_block, mid_block_rname = module.mid_block, \"mid_block\"\n up_blocks, up_blocks_rname = module.up_blocks, \"up_blocks\"\n return UNetStruct(\n module=module,\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n input_embed=input_embed,\n time_embed=time_embed,\n add_time_embed=add_time_embed,\n text_embed=text_embed,\n norm_out=norm_out,\n proj_out=proj_out,\n down_blocks=down_blocks,\n mid_block=mid_block,\n up_blocks=up_blocks,\n input_embed_rname=input_embed_rname,\n time_embed_rname=time_embed_rname,\n add_time_embed_rname=add_time_embed_rname,\n text_embed_rname=text_embed_rname,\n norm_out_rname=norm_out_rname,\n proj_out_rname=proj_out_rname,\n down_blocks_rname=down_blocks_rname,\n mid_block_rname=mid_block_rname,\n up_blocks_rname=up_blocks_rname,\n )\n raise NotImplementedError(f\"Unsupported module type: {type(module)}\")\n\n @classmethod\n def _get_default_key_map(cls) -> dict[str, set[str]]:\n \"\"\"Get the default allowed keys.\"\"\"\n key_map: dict[str, set[str]] = defaultdict(set)\n for idx, (block_key, block_cls) in enumerate(\n (\n (cls.down_block_rkey, cls.down_block_struct_cls),\n (cls.mid_block_rkey, cls.mid_block_struct_cls),\n (cls.up_block_rkey, cls.up_block_struct_cls),\n )\n ):\n block_key_map: dict[str, set[str]] = defaultdict(set)\n if idx != 1:\n sampler_key = join_name(block_key, block_cls.sampler_rkey, sep=\"_\")\n sampler_rkey = block_cls.sampler_rkey\n block_key_map[sampler_rkey].add(sampler_key)\n _block_key_map = block_cls._get_default_key_map()\n for rkey, keys in _block_key_map.items():\n for key in keys:\n key = join_name(block_key, key, sep=\"_\")\n block_key_map[rkey].add(key)\n for rkey, keys in block_key_map.items():\n key_map[rkey].update(keys)\n if block_key:\n key_map[block_key].update(keys)\n keys: set[str] = set()\n keys.add(cls.input_embed_rkey)\n keys.add(cls.time_embed_rkey)\n keys.add(cls.add_time_embed_rkey)\n keys.add(cls.text_embed_rkey)\n keys.add(cls.norm_out_rkey)\n keys.add(cls.proj_out_rkey)\n for mapped_keys in key_map.values():\n for key in mapped_keys:\n keys.add(key)\n if \"embed\" not in keys and \"embed\" not in key_map:\n key_map[\"embed\"].add(cls.input_embed_rkey)\n key_map[\"embed\"].add(cls.time_embed_rkey)\n key_map[\"embed\"].add(cls.add_time_embed_rkey)\n key_map[\"embed\"].add(cls.text_embed_rkey)\n key_map[\"embed\"].add(cls.norm_out_rkey)\n key_map[\"embed\"].add(cls.proj_out_rkey)\n for key in keys:\n if key in key_map:\n key_map[key].clear()\n key_map[key].add(key)\n return {k: v for k, v in key_map.items() if v}\n\n\n@dataclass(kw_only=True)\nclass DiTStruct(DiffusionModelStruct, DiffusionTransformerStruct):\n # region relative keys\n input_embed_rkey: tp.ClassVar[str] = \"input_embed\"\n \"\"\"hidden_states = input_embed(hidden_states), e.g., conv_in\"\"\"\n time_embed_rkey: tp.ClassVar[str] = \"time_embed\"\n \"\"\"temb = time_embed(timesteps)\"\"\"\n text_embed_rkey: tp.ClassVar[str] = \"text_embed\"\n \"\"\"encoder_hidden_states = text_embed(encoder_hidden_states)\"\"\"\n norm_in_rkey: tp.ClassVar[str] = \"input_embed\"\n \"\"\"hidden_states = norm_in(hidden_states)\"\"\"\n proj_in_rkey: tp.ClassVar[str] = \"input_embed\"\n \"\"\"hidden_states = proj_in(hidden_states)\"\"\"\n norm_out_rkey: tp.ClassVar[str] = \"output_embed\"\n \"\"\"hidden_states = norm_out(hidden_states)\"\"\"\n proj_out_rkey: tp.ClassVar[str] = \"output_embed\"\n \"\"\"hidden_states = proj_out(hidden_states)\"\"\"\n transformer_block_rkey: tp.ClassVar[str] = \"\"\n # endregion\n\n # region child modules\n input_embed: PatchEmbed\n time_embed: AdaLayerNormSingle | CombinedTimestepTextProjEmbeddings | TimestepEmbedding\n text_embed: PixArtAlphaTextProjection | nn.Linear\n norm_in: None = field(init=False, repr=False, default=None)\n proj_in: None = field(init=False, repr=False, default=None)\n norm_out: nn.LayerNorm | AdaLayerNormContinuous | None\n proj_out: nn.Linear\n # endregion\n # region relative names\n input_embed_rname: str\n time_embed_rname: str\n text_embed_rname: str\n norm_in_rname: str = field(init=False, repr=False, default=\"\")\n proj_in_rname: str = field(init=False, repr=False, default=\"\")\n norm_out_rname: str\n proj_out_rname: str\n # endregion\n # region absolute names\n input_embed_name: str = field(init=False, repr=False)\n time_embed_name: str = field(init=False, repr=False)\n text_embed_name: str = field(init=False, repr=False)\n # endregion\n # region absolute keys\n input_embed_key: str = field(init=False, repr=False)\n time_embed_key: str = field(init=False, repr=False)\n text_embed_key: str = field(init=False, repr=False)\n norm_out_key: str = field(init=False, repr=False)\n # endregion\n\n @property\n def num_blocks(self) -> int:\n return len(self.transformer_blocks)\n\n @property\n def block_structs(self) -> list[DiffusionTransformerBlockStruct]:\n return self.transformer_block_structs\n\n @property\n def block_names(self) -> list[str]:\n return self.transformer_block_names\n\n def __post_init__(self) -> None:\n super().__post_init__()\n self.pre_module_structs = {}\n for fname in (\"input_embed\", \"time_embed\", \"text_embed\"):\n module, rname, rkey = getattr(self, fname), getattr(self, f\"{fname}_rname\"), getattr(self, f\"{fname}_rkey\")\n setattr(self, f\"{fname}_key\", join_name(self.key, rkey, sep=\"_\"))\n if module is not None or rname:\n setattr(self, f\"{fname}_name\", join_name(self.name, rname))\n else:\n setattr(self, f\"{fname}_name\", \"\")\n if module is not None:\n self.pre_module_structs.setdefault(\n getattr(self, f\"{fname}_name\"),\n DiffusionModuleStruct(module=module, parent=self, fname=fname, rname=rname, rkey=rkey),\n )\n self.post_module_structs = {}\n self.norm_out_key = join_name(self.key, self.norm_out_rkey, sep=\"_\")\n for fname in (\"norm_out\", \"proj_out\"):\n module, rname, rkey = getattr(self, fname), getattr(self, f\"{fname}_rname\"), getattr(self, f\"{fname}_rkey\")\n if module is not None:\n self.post_module_structs.setdefault(\n getattr(self, f\"{fname}_name\"),\n DiffusionModuleStruct(module=module, parent=self, fname=fname, rname=rname, rkey=rkey),\n )\n\n def get_prev_module_keys(self) -> tuple[str, ...]:\n return tuple({self.input_embed_key, self.time_embed_key, self.text_embed_key})\n\n def get_post_module_keys(self) -> tuple[str, ...]:\n return tuple({self.norm_out_key, self.proj_out_key})\n\n def _get_iter_block_activations_args(\n self, **input_kwargs\n ) -> tuple[list[nn.Module], list[DiffusionModuleStruct | DiffusionBlockStruct], list[bool], list[bool]]:\n \"\"\"\n Get the arguments for iterating over the layers and their activations.\n\n Args:\n skip_pre_modules (`bool`):\n Whether to skip the pre-modules\n skip_post_modules (`bool`):\n Whether to skip the post-modules\n\n Returns:\n `tuple[list[nn.Module], list[DiffusionModuleStruct | DiffusionBlockStruct], list[bool], list[bool]]`:\n the layers, the layer structs, the recomputes, and the use_prev_layer_outputs\n \"\"\"\n layers, layer_structs, recomputes, use_prev_layer_outputs = [], [], [], []\n layers.extend(self.transformer_blocks)\n layer_structs.extend(self.transformer_block_structs)\n use_prev_layer_outputs.append(False)\n use_prev_layer_outputs.extend([True] * (len(self.transformer_blocks) - 1))\n recomputes.extend([False] * len(self.transformer_blocks))\n return layers, layer_structs, recomputes, use_prev_layer_outputs\n\n @staticmethod\n def _default_construct(\n module: tp.Union[DIT_PIPELINE_CLS, DIT_CLS],\n /,\n parent: tp.Optional[BaseModuleStruct] = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"DiTStruct\":\n if isinstance(module, DIT_PIPELINE_CLS):\n module = module.transformer\n if isinstance(module, FluxTransformer2DModel):\n return FluxStruct.construct(module, parent=parent, fname=fname, rname=rname, rkey=rkey, idx=idx, **kwargs)\n else:\n if isinstance(module, PixArtTransformer2DModel):\n input_embed, input_embed_rname = module.pos_embed, \"pos_embed\"\n time_embed, time_embed_rname = module.adaln_single, \"adaln_single\"\n text_embed, text_embed_rname = module.caption_projection, \"caption_projection\"\n norm_out, norm_out_rname = module.norm_out, \"norm_out\"\n proj_out, proj_out_rname = module.proj_out, \"proj_out\"\n transformer_blocks, transformer_blocks_rname = module.transformer_blocks, \"transformer_blocks\"\n # ! in fact, `module.adaln_single.emb` is `time_embed`,\n # ! `module.adaln_single.linear` is `transformer_norm`\n # ! but since PixArt shares the `transformer_norm`, we categorize it as `time_embed`\n elif isinstance(module, SanaTransformer2DModel):\n input_embed, input_embed_rname = module.patch_embed, \"patch_embed\"\n time_embed, time_embed_rname = module.time_embed, \"time_embed\"\n text_embed, text_embed_rname = module.caption_projection, \"caption_projection\"\n norm_out, norm_out_rname = module.norm_out, \"norm_out\"\n proj_out, proj_out_rname = module.proj_out, \"proj_out\"\n transformer_blocks, transformer_blocks_rname = module.transformer_blocks, \"transformer_blocks\"\n elif isinstance(module, SD3Transformer2DModel):\n input_embed, input_embed_rname = module.pos_embed, \"pos_embed\"\n time_embed, time_embed_rname = module.time_text_embed, \"time_text_embed\"\n text_embed, text_embed_rname = module.context_embedder, \"context_embedder\"\n norm_out, norm_out_rname = module.norm_out, \"norm_out\"\n proj_out, proj_out_rname = module.proj_out, \"proj_out\"\n transformer_blocks, transformer_blocks_rname = module.transformer_blocks, \"transformer_blocks\"\n else:\n raise NotImplementedError(f\"Unsupported module type: {type(module)}\")\n return DiTStruct(\n module=module,\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n input_embed=input_embed,\n time_embed=time_embed,\n text_embed=text_embed,\n transformer_blocks=transformer_blocks,\n norm_out=norm_out,\n proj_out=proj_out,\n input_embed_rname=input_embed_rname,\n time_embed_rname=time_embed_rname,\n text_embed_rname=text_embed_rname,\n norm_out_rname=norm_out_rname,\n proj_out_rname=proj_out_rname,\n transformer_blocks_rname=transformer_blocks_rname,\n )\n\n @classmethod\n def _get_default_key_map(cls) -> dict[str, set[str]]:\n \"\"\"Get the default allowed keys.\"\"\"\n key_map: dict[str, set[str]] = defaultdict(set)\n block_cls = cls.transformer_block_struct_cls\n block_key = block_rkey = cls.transformer_block_rkey\n block_key_map = block_cls._get_default_key_map()\n for rkey, keys in block_key_map.items():\n brkey = join_name(block_rkey, rkey, sep=\"_\")\n for key in keys:\n key = join_name(block_key, key, sep=\"_\")\n key_map[rkey].add(key)\n key_map[brkey].add(key)\n if block_rkey:\n key_map[block_rkey].add(key)\n keys: set[str] = set()\n keys.add(cls.input_embed_rkey)\n keys.add(cls.time_embed_rkey)\n keys.add(cls.text_embed_rkey)\n keys.add(cls.norm_in_rkey)\n keys.add(cls.proj_in_rkey)\n keys.add(cls.norm_out_rkey)\n keys.add(cls.proj_out_rkey)\n for mapped_keys in key_map.values():\n for key in mapped_keys:\n keys.add(key)\n if \"embed\" not in keys and \"embed\" not in key_map:\n key_map[\"embed\"].add(cls.input_embed_rkey)\n key_map[\"embed\"].add(cls.time_embed_rkey)\n key_map[\"embed\"].add(cls.text_embed_rkey)\n key_map[\"embed\"].add(cls.norm_in_rkey)\n key_map[\"embed\"].add(cls.proj_in_rkey)\n key_map[\"embed\"].add(cls.norm_out_rkey)\n key_map[\"embed\"].add(cls.proj_out_rkey)\n for key in keys:\n if key in key_map:\n key_map[key].clear()\n key_map[key].add(key)\n return {k: v for k, v in key_map.items() if v}\n\n\n@dataclass(kw_only=True)\nclass FluxStruct(DiTStruct):\n # region relative keys\n single_transformer_block_rkey: tp.ClassVar[str] = \"\"\n single_transformer_block_struct_cls: tp.ClassVar[type[DiffusionTransformerBlockStruct]] = (\n DiffusionTransformerBlockStruct\n )\n # endregion\n\n module: FluxTransformer2DModel = field(repr=False, kw_only=False)\n \"\"\"the module of FluxTransformer2DModel\"\"\"\n # region child modules\n input_embed: nn.Linear\n time_embed: CombinedTimestepGuidanceTextProjEmbeddings | CombinedTimestepTextProjEmbeddings\n text_embed: nn.Linear\n single_transformer_blocks: nn.ModuleList = field(repr=False)\n # endregion\n # region relative names\n single_transformer_blocks_rname: str\n # endregion\n # region absolute names\n single_transformer_blocks_name: str = field(init=False, repr=False)\n single_transformer_block_names: list[str] = field(init=False, repr=False)\n # endregion\n # region child structs\n single_transformer_block_structs: list[DiffusionTransformerBlockStruct] = field(init=False)\n # endregion\n\n @property\n def num_blocks(self) -> int:\n return len(self.transformer_block_structs) + len(self.single_transformer_block_structs)\n\n @property\n def block_structs(self) -> list[DiffusionTransformerBlockStruct]:\n return [*self.transformer_block_structs, *self.single_transformer_block_structs]\n\n @property\n def block_names(self) -> list[str]:\n return [*self.transformer_block_names, *self.single_transformer_block_names]\n\n def __post_init__(self) -> None:\n super().__post_init__()\n single_transformer_block_rnames = [\n f\"{self.single_transformer_blocks_rname}.{idx}\" for idx in range(len(self.single_transformer_blocks))\n ]\n self.single_transformer_blocks_name = join_name(self.name, self.single_transformer_blocks_rname)\n self.single_transformer_block_names = [join_name(self.name, rname) for rname in single_transformer_block_rnames]\n self.single_transformer_block_structs = [\n self.single_transformer_block_struct_cls.construct(\n block,\n parent=self,\n fname=\"single_transformer_block\",\n rname=rname,\n rkey=self.single_transformer_block_rkey,\n idx=idx,\n )\n for idx, (block, rname) in enumerate(\n zip(self.single_transformer_blocks, single_transformer_block_rnames, strict=True)\n )\n ]\n\n def _get_iter_block_activations_args(\n self, **input_kwargs\n ) -> tuple[list[nn.Module], list[DiffusionModuleStruct | DiffusionBlockStruct], list[bool], list[bool]]:\n layers, layer_structs, recomputes, use_prev_layer_outputs = super()._get_iter_block_activations_args()\n layers.extend(self.single_transformer_blocks)\n layer_structs.extend(self.single_transformer_block_structs)\n use_prev_layer_outputs.append(False)\n use_prev_layer_outputs.extend([True] * (len(self.single_transformer_blocks) - 1))\n recomputes.extend([False] * len(self.single_transformer_blocks))\n return layers, layer_structs, recomputes, use_prev_layer_outputs\n\n @staticmethod\n def _default_construct(\n module: tp.Union[FluxPipeline, FluxControlPipeline, FluxTransformer2DModel],\n /,\n parent: tp.Optional[BaseModuleStruct] = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"FluxStruct\":\n if isinstance(module, (FluxPipeline, FluxControlPipeline)):\n module = module.transformer\n if isinstance(module, FluxTransformer2DModel):\n input_embed, time_embed, text_embed = module.x_embedder, module.time_text_embed, module.context_embedder\n input_embed_rname, time_embed_rname, text_embed_rname = \"x_embedder\", \"time_text_embed\", \"context_embedder\"\n norm_out, norm_out_rname = module.norm_out, \"norm_out\"\n proj_out, proj_out_rname = module.proj_out, \"proj_out\"\n transformer_blocks, transformer_blocks_rname = module.transformer_blocks, \"transformer_blocks\"\n single_transformer_blocks = module.single_transformer_blocks\n single_transformer_blocks_rname = \"single_transformer_blocks\"\n return FluxStruct(\n module=module,\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n input_embed=input_embed,\n time_embed=time_embed,\n text_embed=text_embed,\n transformer_blocks=transformer_blocks,\n single_transformer_blocks=single_transformer_blocks,\n norm_out=norm_out,\n proj_out=proj_out,\n input_embed_rname=input_embed_rname,\n time_embed_rname=time_embed_rname,\n text_embed_rname=text_embed_rname,\n norm_out_rname=norm_out_rname,\n proj_out_rname=proj_out_rname,\n transformer_blocks_rname=transformer_blocks_rname,\n single_transformer_blocks_rname=single_transformer_blocks_rname,\n )\n raise NotImplementedError(f\"Unsupported module type: {type(module)}\")\n\n @classmethod\n def _get_default_key_map(cls) -> dict[str, set[str]]:\n \"\"\"Get the default allowed keys.\"\"\"\n key_map: dict[str, set[str]] = defaultdict(set)\n for block_rkey, block_cls in (\n (cls.transformer_block_rkey, cls.transformer_block_struct_cls),\n (cls.single_transformer_block_rkey, cls.single_transformer_block_struct_cls),\n ):\n block_key = block_rkey\n block_key_map = block_cls._get_default_key_map()\n for rkey, keys in block_key_map.items():\n brkey = join_name(block_rkey, rkey, sep=\"_\")\n for key in keys:\n key = join_name(block_key, key, sep=\"_\")\n key_map[rkey].add(key)\n key_map[brkey].add(key)\n if block_rkey:\n key_map[block_rkey].add(key)\n keys: set[str] = set()\n keys.add(cls.input_embed_rkey)\n keys.add(cls.time_embed_rkey)\n keys.add(cls.text_embed_rkey)\n keys.add(cls.norm_in_rkey)\n keys.add(cls.proj_in_rkey)\n keys.add(cls.norm_out_rkey)\n keys.add(cls.proj_out_rkey)\n for mapped_keys in key_map.values():\n for key in mapped_keys:\n keys.add(key)\n if \"embed\" not in keys and \"embed\" not in key_map:\n key_map[\"embed\"].add(cls.input_embed_rkey)\n key_map[\"embed\"].add(cls.time_embed_rkey)\n key_map[\"embed\"].add(cls.text_embed_rkey)\n key_map[\"embed\"].add(cls.norm_in_rkey)\n key_map[\"embed\"].add(cls.proj_in_rkey)\n key_map[\"embed\"].add(cls.norm_out_rkey)\n key_map[\"embed\"].add(cls.proj_out_rkey)\n for key in keys:\n if key in key_map:\n key_map[key].clear()\n key_map[key].add(key)\n return {k: v for k, v in key_map.items() if v}\n\n\nDiffusionAttentionStruct.register_factory(Attention, DiffusionAttentionStruct._default_construct)\n\nDiffusionFeedForwardStruct.register_factory(\n (FeedForward, FluxSingleTransformerBlock, GLUMBConv), DiffusionFeedForwardStruct._default_construct\n)\n\nDiffusionTransformerBlockStruct.register_factory(DIT_BLOCK_CLS, DiffusionTransformerBlockStruct._default_construct)\n\nUNetBlockStruct.register_factory(UNET_BLOCK_CLS, UNetBlockStruct._default_construct)\n\nUNetStruct.register_factory(tp.Union[UNET_PIPELINE_CLS, UNET_CLS], UNetStruct._default_construct)\n\nFluxStruct.register_factory(\n tp.Union[FluxPipeline, FluxControlPipeline, FluxTransformer2DModel], FluxStruct._default_construct\n)\n\nDiTStruct.register_factory(tp.Union[DIT_PIPELINE_CLS, DIT_CLS], DiTStruct._default_construct)\n\nDiffusionTransformerStruct.register_factory(Transformer2DModel, DiffusionTransformerStruct._default_construct)\n\nDiffusionModelStruct.register_factory(tp.Union[PIPELINE_CLS, MODEL_CLS], DiffusionModelStruct._default_construct)\n"
},
{
"path": "deepcompressor/app/diffusion/pipeline/__init__.py",
"content": "# -*- coding: utf-8 -*-\n\nfrom .config import DiffusionPipelineConfig\n"
},
{
"path": "deepcompressor/app/diffusion/pipeline/config.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Diffusion pipeline configuration module.\"\"\"\n\nimport gc\nimport typing as tp\nfrom dataclasses import dataclass, field\n\nimport torch\nfrom diffusers.pipelines import (\n AutoPipelineForText2Image,\n DiffusionPipeline,\n FluxControlPipeline,\n FluxFillPipeline,\n SanaPipeline,\n)\nfrom omniconfig import configclass\nfrom torch import nn\nfrom transformers import PreTrainedModel, PreTrainedTokenizer, T5EncoderModel\n\nfrom deepcompressor.data.utils.dtype import eval_dtype\nfrom deepcompressor.quantizer.processor import Quantizer\nfrom deepcompressor.utils import tools\nfrom deepcompressor.utils.hooks import AccumBranchHook, ProcessHook\n\nfrom ....nn.patch.linear import ConcatLinear, ShiftedLinear\nfrom ....nn.patch.lowrank import LowRankBranch\nfrom ..nn.patch import (\n replace_fused_linear_with_concat_linear,\n replace_up_block_conv_with_concat_conv,\n shift_input_activations,\n)\n\n__all__ = [\"DiffusionPipelineConfig\"]\n\n\n@configclass\n@dataclass\nclass LoRAConfig:\n \"\"\"LoRA configuration.\n\n Args:\n path (`str`):\n The path of the LoRA branch.\n weight_name (`str`):\n The weight name of the LoRA branch.\n alpha (`float`):\n The alpha value of the LoRA branch.\n \"\"\"\n\n path: str\n weight_name: str\n alpha: float = 1.0\n\n\n@configclass\n@dataclass\nclass DiffusionPipelineConfig:\n \"\"\"Diffusion pipeline configuration.\n\n Args:\n name (`str`):\n The name of the pipeline.\n dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The data type of the pipeline.\n device (`str`, *optional*, defaults to `\"cuda\"`):\n The device of the pipeline.\n shift_activations (`bool`, *optional*, defaults to `False`):\n Whether to shift activations.\n \"\"\"\n\n _pipeline_factories: tp.ClassVar[\n dict[str, tp.Callable[[str, str, torch.dtype, torch.device, bool], DiffusionPipeline]]\n ] = {}\n _text_extractors: tp.ClassVar[\n dict[\n str,\n tp.Callable[\n [DiffusionPipeline, tuple[type[PreTrainedModel], ...]],\n list[tuple[str, PreTrainedModel, PreTrainedTokenizer]],\n ],\n ]\n ] = {}\n\n name: str\n path: str = \"\"\n dtype: torch.dtype = field(\n default_factory=lambda s=torch.float32: eval_dtype(s, with_quant_dtype=False, with_none=False)\n )\n device: str = \"cuda\"\n shift_activations: bool = False\n lora: LoRAConfig | None = None\n family: str = field(init=False)\n task: str = \"text-to-image\"\n\n def __post_init__(self):\n self.family = self.name.split(\"-\")[0]\n\n if self.name == \"flux.1-canny-dev\":\n self.task = \"canny-to-image\"\n elif self.name == \"flux.1-depth-dev\":\n self.task = \"depth-to-image\"\n elif self.name == \"flux.1-fill-dev\":\n self.task = \"inpainting\"\n\n def build(\n self, *, dtype: str | torch.dtype | None = None, device: str | torch.device | None = None\n ) -> DiffusionPipeline:\n \"\"\"Build the diffusion pipeline.\n\n Args:\n dtype (`str` or `torch.dtype`, *optional*):\n The data type of the pipeline.\n device (`str` or `torch.device`, *optional*):\n The device of the pipeline.\n\n Returns:\n `DiffusionPipeline`:\n The diffusion pipeline.\n \"\"\"\n if dtype is None:\n dtype = self.dtype\n if device is None:\n device = self.device\n _factory = self._pipeline_factories.get(self.name, self._default_build)\n return _factory(\n name=self.name, path=self.path, dtype=dtype, device=device, shift_activations=self.shift_activations\n )\n\n def extract_text_encoders(\n self, pipeline: DiffusionPipeline, supported: tuple[type[PreTrainedModel], ...] = (T5EncoderModel,)\n ) -> list[tuple[str, PreTrainedModel, PreTrainedTokenizer]]:\n \"\"\"Extract the text encoders and tokenizers from the pipeline.\n\n Args:\n pipeline (`DiffusionPipeline`):\n The diffusion pipeline.\n supported (`tuple[type[PreTrainedModel], ...]`, *optional*, defaults to `(T5EncoderModel,)`):\n The supported text encoder types. If not specified, all text encoders will be extracted.\n\n Returns:\n `list[tuple[str, PreTrainedModel, PreTrainedTokenizer]]`:\n The list of text encoder name, model, and tokenizer.\n \"\"\"\n _extractor = self._text_extractors.get(self.name, self._default_extract_text_encoders)\n return _extractor(pipeline, supported)\n\n @classmethod\n def register_pipeline_factory(\n cls,\n names: str | tuple[str, ...],\n /,\n factory: tp.Callable[[str, str, torch.dtype, torch.device, bool], DiffusionPipeline],\n *,\n overwrite: bool = False,\n ) -> None:\n \"\"\"Register a pipeline factory.\n\n Args:\n names (`str` or `tuple[str, ...]`):\n The name of the pipeline.\n factory (`Callable[[str, str,torch.dtype, torch.device, bool], DiffusionPipeline]`):\n The pipeline factory function.\n overwrite (`bool`, *optional*, defaults to `False`):\n Whether to overwrite the existing factory for the pipeline.\n \"\"\"\n if isinstance(names, str):\n names = [names]\n for name in names:\n if name in cls._pipeline_factories and not overwrite:\n raise ValueError(f\"Pipeline factory {name} already exists.\")\n cls._pipeline_factories[name] = factory\n\n @classmethod\n def register_text_extractor(\n cls,\n names: str | tuple[str, ...],\n /,\n extractor: tp.Callable[\n [DiffusionPipeline, tuple[type[PreTrainedModel], ...]],\n list[tuple[str, PreTrainedModel, PreTrainedTokenizer]],\n ],\n *,\n overwrite: bool = False,\n ) -> None:\n \"\"\"Register a text extractor.\n\n Args:\n names (`str` or `tuple[str, ...]`):\n The name of the pipeline.\n extractor (`Callable[[DiffusionPipeline], list[tuple[str, PreTrainedModel, PreTrainedTokenizer]]`):\n The text extractor function.\n overwrite (`bool`, *optional*, defaults to `False`):\n Whether to overwrite the existing extractor for the pipeline.\n \"\"\"\n if isinstance(names, str):\n names = [names]\n for name in names:\n if name in cls._text_extractors and not overwrite:\n raise ValueError(f\"Text extractor {name} already exists.\")\n cls._text_extractors[name] = extractor\n\n def load_lora( # noqa: C901\n self, pipeline: DiffusionPipeline, smooth_cache: dict[str, torch.Tensor] | None = None\n ) -> DiffusionPipeline:\n smooth_cache = smooth_cache or {}\n model = pipeline.unet if hasattr(pipeline, \"unet\") else pipeline.transformer\n assert isinstance(model, nn.Module)\n if self.lora is not None:\n logger = tools.logging.getLogger(__name__)\n logger.info(f\"Load LoRA branches from {self.lora.path}\")\n lora_state_dict, alphas = pipeline.lora_state_dict(\n self.lora.path, return_alphas=True, weight_name=self.lora.weight_name\n )\n tools.logging.Formatter.indent_inc()\n for name, module in model.named_modules():\n if isinstance(module, (nn.Linear, ConcatLinear, ShiftedLinear)):\n lora_a_key, lora_b_key = f\"transformer.{name}.lora_A.weight\", f\"transformer.{name}.lora_B.weight\"\n if lora_a_key in lora_state_dict:\n assert lora_b_key in lora_state_dict\n logger.info(f\"+ Load LoRA branch for {name}\")\n tools.logging.Formatter.indent_inc()\n a = lora_state_dict.pop(lora_a_key)\n b = lora_state_dict.pop(lora_b_key)\n assert isinstance(a, torch.Tensor)\n assert isinstance(b, torch.Tensor)\n assert a.shape[1] == module.in_features\n assert b.shape[0] == module.out_features\n if isinstance(module, ConcatLinear):\n logger.debug(\n f\"- split LoRA branch into {len(module.linears)} parts ({module.in_features_list})\"\n )\n m_splits = module.linears\n a_splits = a.split(module.in_features_list, dim=1)\n b_splits = [b] * len(a_splits)\n else:\n m_splits, a_splits, b_splits = [module], [a], [b]\n for m, a, b in zip(m_splits, a_splits, b_splits, strict=True):\n assert a.shape[0] == b.shape[1]\n if isinstance(m, ShiftedLinear):\n s, m = m.shift, m.linear\n logger.debug(f\"- shift LoRA input by {s.item() if s.numel() == 1 else s}\")\n else:\n s = None\n assert isinstance(m, nn.Linear)\n device, dtype = m.weight.device, m.weight.dtype\n a, b = a.to(device=device, dtype=torch.float64), b.to(device=device, dtype=torch.float64)\n if s is not None:\n if s.numel() == 1:\n s = torch.matmul(b, a.sum(dim=1).mul_(s.double())).mul_(self.lora.alpha)\n else:\n s = torch.matmul(b, torch.matmul(a, s.view(1, -1).double())).mul_(self.lora.alpha)\n if hasattr(m, \"in_smooth_cache_key\"):\n logger.debug(f\"- smooth LoRA input using {m.in_smooth_cache_key} smooth scale\")\n ss = smooth_cache[m.in_smooth_cache_key].to(device=device, dtype=torch.float64)\n a = a.mul_(ss.view(1, -1))\n del ss\n if hasattr(m, \"out_smooth_cache_key\"):\n logger.debug(f\"- smooth LoRA output using {m.out_smooth_cache_key} smooth scale\")\n ss = smooth_cache[m.out_smooth_cache_key].to(device=device, dtype=torch.float64)\n b = b.div_(ss.view(-1, 1))\n if s is not None:\n s = s.div_(ss.view(-1))\n del ss\n branch_hook, quant_hook = None, None\n for hook in m._forward_pre_hooks.values():\n if isinstance(hook, AccumBranchHook) and isinstance(hook.branch, LowRankBranch):\n branch_hook = hook\n if isinstance(hook, ProcessHook) and isinstance(hook.processor, Quantizer):\n quant_hook = hook\n if branch_hook is not None:\n logger.debug(\"- fuse with existing LoRA branch\")\n assert isinstance(branch_hook.branch, LowRankBranch)\n _a = branch_hook.branch.a.weight.data\n _b = branch_hook.branch.b.weight.data\n if branch_hook.branch.alpha != self.lora.alpha:\n a, b = a.to(dtype=dtype), b.mul_(self.lora.alpha).to(dtype=dtype)\n _b = _b.to(dtype=torch.float64).mul_(branch_hook.branch.alpha).to(dtype=dtype)\n alpha = 1\n else:\n a, b = a.to(dtype=dtype), b.to(dtype=dtype)\n alpha = self.lora.alpha\n branch_hook.branch = LowRankBranch(\n m.in_features,\n m.out_features,\n rank=a.shape[0] + branch_hook.branch.rank,\n alpha=alpha,\n ).to(device=device, dtype=dtype)\n branch_hook.branch.a.weight.data[: a.shape[0], :] = a\n branch_hook.branch.b.weight.data[:, : b.shape[1]] = b\n branch_hook.branch.a.weight.data[a.shape[0] :, :] = _a\n branch_hook.branch.b.weight.data[:, b.shape[1] :] = _b\n else:\n logger.debug(\"- create a new LoRA branch\")\n branch = LowRankBranch(\n m.in_features, m.out_features, rank=a.shape[0], alpha=self.lora.alpha\n )\n branch = branch.to(device=device, dtype=dtype)\n branch.a.weight.data.copy_(a.to(dtype=dtype))\n branch.b.weight.data.copy_(b.to(dtype=dtype))\n # low rank branch hook should be registered before the quantization hook\n if quant_hook is not None:\n logger.debug(f\"- remove quantization hook from {name}\")\n quant_hook.remove(m)\n logger.debug(f\"- register LoRA branch to {name}\")\n branch.as_hook().register(m)\n if quant_hook is not None:\n logger.debug(f\"- re-register quantization hook to {name}\")\n quant_hook.register(m)\n if s is not None:\n assert m.bias is not None\n m.bias.data.copy_((m.bias.double().sub_(s)).to(dtype))\n del m_splits, a_splits, b_splits, a, b, s\n gc.collect()\n torch.cuda.empty_cache()\n tools.logging.Formatter.indent_dec()\n tools.logging.Formatter.indent_dec()\n if len(lora_state_dict) > 0:\n logger.warning(f\"Unused LoRA weights: {lora_state_dict.keys()}\")\n branches = nn.ModuleList()\n for _, module in model.named_modules():\n for hook in module._forward_hooks.values():\n if isinstance(hook, AccumBranchHook) and isinstance(hook.branch, LowRankBranch):\n branches.append(hook.branch)\n model.register_module(\"_low_rank_branches\", branches)\n\n @staticmethod\n def _default_build(\n name: str, path: str, dtype: str | torch.dtype, device: str | torch.device, shift_activations: bool\n ) -> DiffusionPipeline:\n if not path:\n if name == \"sdxl\":\n path = \"stabilityai/stable-diffusion-xl-base-1.0\"\n elif name == \"sdxl-turbo\":\n path = \"stabilityai/sdxl-turbo\"\n elif name == \"pixart-sigma\":\n path = \"PixArt-alpha/PixArt-Sigma-XL-2-1024-MS\"\n elif name == \"flux.1-dev\":\n path = \"black-forest-labs/FLUX.1-dev\"\n elif name == \"flux.1-canny-dev\":\n path = \"black-forest-labs/FLUX.1-Canny-dev\"\n elif name == \"flux.1-depth-dev\":\n path = \"black-forest-labs/FLUX.1-Depth-dev\"\n elif name == \"flux.1-fill-dev\":\n path = \"black-forest-labs/FLUX.1-Fill-dev\"\n elif name == \"flux.1-schnell\":\n path = \"black-forest-labs/FLUX.1-schnell\"\n else:\n raise ValueError(f\"Path for {name} is not specified.\")\n if name in [\"flux.1-canny-dev\", \"flux.1-depth-dev\"]:\n pipeline = FluxControlPipeline.from_pretrained(path, torch_dtype=dtype)\n elif name == \"flux.1-fill-dev\":\n pipeline = FluxFillPipeline.from_pretrained(path, torch_dtype=dtype)\n elif name.startswith(\"sana-\"):\n if dtype == torch.bfloat16:\n pipeline = SanaPipeline.from_pretrained(path, variant=\"bf16\", torch_dtype=dtype, use_safetensors=True)\n pipeline.vae.to(dtype)\n pipeline.text_encoder.to(dtype)\n else:\n pipeline = SanaPipeline.from_pretrained(path, torch_dtype=dtype)\n else:\n pipeline = AutoPipelineForText2Image.from_pretrained(path, torch_dtype=dtype)\n pipeline = pipeline.to(device)\n model = pipeline.unet if hasattr(pipeline, \"unet\") else pipeline.transformer\n replace_fused_linear_with_concat_linear(model)\n replace_up_block_conv_with_concat_conv(model)\n if shift_activations:\n shift_input_activations(model)\n return pipeline\n\n @staticmethod\n def _default_extract_text_encoders(\n pipeline: DiffusionPipeline, supported: tuple[type[PreTrainedModel], ...]\n ) -> list[tuple[str, PreTrainedModel, PreTrainedTokenizer]]:\n \"\"\"Extract the text encoders and tokenizers from the pipeline.\n\n Args:\n pipeline (`DiffusionPipeline`):\n The diffusion pipeline.\n supported (`tuple[type[PreTrainedModel], ...]`, *optional*, defaults to `(T5EncoderModel,)`):\n The supported text encoder types. If not specified, all text encoders will be extracted.\n\n Returns:\n `list[tuple[str, PreTrainedModel, PreTrainedTokenizer]]`:\n The list of text encoder name, model, and tokenizer.\n \"\"\"\n results: list[tuple[str, PreTrainedModel, PreTrainedTokenizer]] = []\n for key in vars.__dict__.keys():\n if key.startswith(\"text_encoder\"):\n suffix = key[len(\"text_encoder\") :]\n encoder, tokenizer = getattr(pipeline, f\"text_encoder{suffix}\"), getattr(pipeline, f\"tokenizer{suffix}\")\n if not supported or isinstance(encoder, supported):\n results.append((key, encoder, tokenizer))\n return results\n"
},
{
"path": "deepcompressor/app/diffusion/ptq.py",
"content": "import gc\nimport json\nimport os\nimport pprint\nimport traceback\n\nimport torch\nfrom diffusers import DiffusionPipeline\n\nfrom deepcompressor.app.llm.nn.patch import patch_attention, patch_gemma_rms_norm\nfrom deepcompressor.app.llm.ptq import ptq as llm_ptq\nfrom deepcompressor.utils import tools\n\nfrom .config import DiffusionPtqCacheConfig, DiffusionPtqRunConfig, DiffusionQuantCacheConfig, DiffusionQuantConfig\nfrom .nn.struct import DiffusionModelStruct\nfrom .quant import (\n load_diffusion_weights_state_dict,\n quantize_diffusion_activations,\n quantize_diffusion_weights,\n rotate_diffusion,\n smooth_diffusion,\n)\n\n__all__ = [\"ptq\"]\n\n\ndef ptq( # noqa: C901\n model: DiffusionModelStruct,\n config: DiffusionQuantConfig,\n cache: DiffusionPtqCacheConfig | None = None,\n load_dirpath: str = \"\",\n save_dirpath: str = \"\",\n copy_on_save: bool = False,\n save_model: bool = False,\n) -> DiffusionModelStruct:\n \"\"\"Post-training quantization of a diffusion model.\n\n Args:\n model (`DiffusionModelStruct`):\n The diffusion model.\n config (`DiffusionQuantConfig`):\n The diffusion model post-training quantization configuration.\n cache (`DiffusionPtqCacheConfig`, *optional*, defaults to `None`):\n The diffusion model quantization cache path configuration.\n load_dirpath (`str`, *optional*, defaults to `\"\"`):\n The directory path to load the quantization checkpoint.\n save_dirpath (`str`, *optional*, defaults to `\"\"`):\n The directory path to save the quantization checkpoint.\n copy_on_save (`bool`, *optional*, defaults to `False`):\n Whether to copy the cache to the save directory.\n save_model (`bool`, *optional*, defaults to `False`):\n Whether to save the quantized model checkpoint.\n\n Returns:\n `DiffusionModelStruct`:\n The quantized diffusion model.\n \"\"\"\n logger = tools.logging.getLogger(__name__)\n if not isinstance(model, DiffusionModelStruct):\n model = DiffusionModelStruct.construct(model)\n assert isinstance(model, DiffusionModelStruct)\n\n quant_wgts = config.enabled_wgts\n quant_ipts = config.enabled_ipts\n quant_opts = config.enabled_opts\n quant_acts = quant_ipts or quant_opts\n quant = quant_wgts or quant_acts\n\n load_model_path, load_path, save_path = \"\", None, None\n if load_dirpath:\n load_path = DiffusionQuantCacheConfig(\n smooth=os.path.join(load_dirpath, \"smooth.pt\"),\n branch=os.path.join(load_dirpath, \"branch.pt\"),\n wgts=os.path.join(load_dirpath, \"wgts.pt\"),\n acts=os.path.join(load_dirpath, \"acts.pt\"),\n )\n load_model_path = os.path.join(load_dirpath, \"model.pt\")\n if os.path.exists(load_model_path):\n if config.enabled_wgts and config.wgts.enabled_low_rank:\n if os.path.exists(load_path.branch):\n load_model = True\n else:\n logger.warning(f\"Model low-rank branch checkpoint {load_path.branch} does not exist\")\n load_model = False\n else:\n load_model = True\n if load_model:\n logger.info(f\"* Loading model from {load_model_path}\")\n save_dirpath = \"\" # do not save the model if loading\n else:\n logger.warning(f\"Model checkpoint {load_model_path} does not exist\")\n load_model = False\n else:\n load_model = False\n if save_dirpath:\n os.makedirs(save_dirpath, exist_ok=True)\n save_path = DiffusionQuantCacheConfig(\n smooth=os.path.join(save_dirpath, \"smooth.pt\"),\n branch=os.path.join(save_dirpath, \"branch.pt\"),\n wgts=os.path.join(save_dirpath, \"wgts.pt\"),\n acts=os.path.join(save_dirpath, \"acts.pt\"),\n )\n else:\n save_model = False\n\n if quant and config.enabled_rotation:\n logger.info(\"* Rotating model for quantization\")\n tools.logging.Formatter.indent_inc()\n rotate_diffusion(model, config=config)\n tools.logging.Formatter.indent_dec()\n gc.collect()\n torch.cuda.empty_cache()\n\n # region smooth quantization\n if quant and config.enabled_smooth:\n logger.info(\"* Smoothing model for quantization\")\n tools.logging.Formatter.indent_inc()\n load_from = \"\"\n if load_path and os.path.exists(load_path.smooth):\n load_from = load_path.smooth\n elif cache and cache.path.smooth and os.path.exists(cache.path.smooth):\n load_from = cache.path.smooth\n if load_from:\n logger.info(f\"- Loading smooth scales from {load_from}\")\n smooth_cache = torch.load(load_from)\n smooth_diffusion(model, config, smooth_cache=smooth_cache)\n else:\n logger.info(\"- Generating smooth scales\")\n smooth_cache = smooth_diffusion(model, config)\n if cache and cache.path.smooth:\n logger.info(f\"- Saving smooth scales to {cache.path.smooth}\")\n os.makedirs(cache.dirpath.smooth, exist_ok=True)\n torch.save(smooth_cache, cache.path.smooth)\n load_from = cache.path.smooth\n if save_path:\n if not copy_on_save and load_from:\n logger.info(f\"- Linking smooth scales to {save_path.smooth}\")\n os.symlink(os.path.relpath(load_from, save_dirpath), save_path.smooth)\n else:\n logger.info(f\"- Saving smooth scales to {save_path.smooth}\")\n torch.save(smooth_cache, save_path.smooth)\n del smooth_cache\n tools.logging.Formatter.indent_dec()\n gc.collect()\n torch.cuda.empty_cache()\n # endregion\n # region collect original state dict\n if config.needs_acts_quantizer_cache:\n if load_path and os.path.exists(load_path.acts):\n orig_state_dict = None\n elif cache and cache.path.acts and os.path.exists(cache.path.acts):\n orig_state_dict = None\n else:\n orig_state_dict: dict[str, torch.Tensor] = {\n name: param.detach().clone() for name, param in model.module.named_parameters() if param.ndim > 1\n }\n else:\n orig_state_dict = None\n # endregion\n if load_model:\n logger.info(f\"* Loading model checkpoint from {load_model_path}\")\n load_diffusion_weights_state_dict(\n model,\n config,\n state_dict=torch.load(load_model_path),\n branch_state_dict=torch.load(load_path.branch) if os.path.exists(load_path.branch) else None,\n )\n gc.collect()\n torch.cuda.empty_cache()\n elif quant_wgts:\n logger.info(\"* Quantizing weights\")\n tools.logging.Formatter.indent_inc()\n quantizer_state_dict, quantizer_load_from = None, \"\"\n if load_path and os.path.exists(load_path.wgts):\n quantizer_load_from = load_path.wgts\n elif cache and cache.path.wgts and os.path.exists(cache.path.wgts):\n quantizer_load_from = cache.path.wgts\n if quantizer_load_from:\n logger.info(f\"- Loading weight settings from {quantizer_load_from}\")\n quantizer_state_dict = torch.load(quantizer_load_from)\n branch_state_dict, branch_load_from = None, \"\"\n if load_path and os.path.exists(load_path.branch):\n branch_load_from = load_path.branch\n elif cache and cache.path.branch and os.path.exists(cache.path.branch):\n branch_load_from = cache.path.branch\n if branch_load_from:\n logger.info(f\"- Loading branch settings from {branch_load_from}\")\n branch_state_dict = torch.load(branch_load_from)\n if not quantizer_load_from:\n logger.info(\"- Generating weight settings\")\n if not branch_load_from:\n logger.info(\"- Generating branch settings\")\n quantizer_state_dict, branch_state_dict, scale_state_dict = quantize_diffusion_weights(\n model,\n config,\n quantizer_state_dict=quantizer_state_dict,\n branch_state_dict=branch_state_dict,\n return_with_scale_state_dict=bool(save_dirpath),\n )\n if not quantizer_load_from and cache and cache.dirpath.wgts:\n logger.info(f\"- Saving weight settings to {cache.path.wgts}\")\n os.makedirs(cache.dirpath.wgts, exist_ok=True)\n torch.save(quantizer_state_dict, cache.path.wgts)\n quantizer_load_from = cache.path.wgts\n if not branch_load_from and cache and cache.dirpath.branch:\n logger.info(f\"- Saving branch settings to {cache.path.branch}\")\n os.makedirs(cache.dirpath.branch, exist_ok=True)\n torch.save(branch_state_dict, cache.path.branch)\n branch_load_from = cache.path.branch\n if save_path:\n if not copy_on_save and quantizer_load_from:\n logger.info(f\"- Linking weight settings to {save_path.wgts}\")\n os.symlink(os.path.relpath(quantizer_load_from, save_dirpath), save_path.wgts)\n else:\n logger.info(f\"- Saving weight settings to {save_path.wgts}\")\n torch.save(quantizer_state_dict, save_path.wgts)\n if not copy_on_save and branch_load_from:\n logger.info(f\"- Linking branch settings to {save_path.branch}\")\n os.symlink(os.path.relpath(branch_load_from, save_dirpath), save_path.branch)\n else:\n logger.info(f\"- Saving branch settings to {save_path.branch}\")\n torch.save(branch_state_dict, save_path.branch)\n if save_model:\n logger.info(f\"- Saving model to {save_dirpath}\")\n torch.save(scale_state_dict, os.path.join(save_dirpath, \"scale.pt\"))\n torch.save(model.module.state_dict(), os.path.join(save_dirpath, \"model.pt\"))\n del quantizer_state_dict, branch_state_dict, scale_state_dict\n tools.logging.Formatter.indent_dec()\n gc.collect()\n torch.cuda.empty_cache()\n if quant_acts:\n logger.info(\" * Quantizing activations\")\n tools.logging.Formatter.indent_inc()\n if config.needs_acts_quantizer_cache:\n load_from = \"\"\n if load_path and os.path.exists(load_path.acts):\n load_from = load_path.acts\n elif cache and cache.path.acts and os.path.exists(cache.path.acts):\n load_from = cache.path.acts\n if load_from:\n logger.info(f\"- Loading activation settings from {load_from}\")\n quantizer_state_dict = torch.load(load_from)\n quantize_diffusion_activations(\n model, config, quantizer_state_dict=quantizer_state_dict, orig_state_dict=orig_state_dict\n )\n else:\n logger.info(\"- Generating activation settings\")\n quantizer_state_dict = quantize_diffusion_activations(model, config, orig_state_dict=orig_state_dict)\n if cache and cache.dirpath.acts and quantizer_state_dict is not None:\n logger.info(f\"- Saving activation settings to {cache.path.acts}\")\n os.makedirs(cache.dirpath.acts, exist_ok=True)\n torch.save(quantizer_state_dict, cache.path.acts)\n load_from = cache.path.acts\n if save_dirpath:\n if not copy_on_save and load_from:\n logger.info(f\"- Linking activation quantizer settings to {save_path.acts}\")\n os.symlink(os.path.relpath(load_from, save_dirpath), save_path.acts)\n else:\n logger.info(f\"- Saving activation quantizer settings to {save_path.acts}\")\n torch.save(quantizer_state_dict, save_path.acts)\n del quantizer_state_dict\n else:\n logger.info(\"- No need to generate/load activation quantizer settings\")\n quantize_diffusion_activations(model, config, orig_state_dict=orig_state_dict)\n tools.logging.Formatter.indent_dec()\n del orig_state_dict\n gc.collect()\n torch.cuda.empty_cache()\n return model\n\n\ndef main(config: DiffusionPtqRunConfig, logging_level: int = tools.logging.DEBUG) -> DiffusionPipeline:\n \"\"\"Post-training quantization of a diffusion model.\n\n Args:\n config (`DiffusionPtqRunConfig`):\n The diffusion model post-training quantization configuration.\n logging_level (`int`, *optional*, defaults to `logging.DEBUG`):\n The logging level.\n\n Returns:\n `DiffusionPipeline`:\n The diffusion pipeline with quantized model.\n \"\"\"\n config.output.lock()\n config.dump(path=config.output.get_running_job_path(\"config.yaml\"))\n tools.logging.setup(path=config.output.get_running_job_path(\"run.log\"), level=logging_level)\n logger = tools.logging.getLogger(__name__)\n\n logger.info(\"=== Configurations ===\")\n tools.logging.info(config.formatted_str(), logger=logger)\n logger.info(\"=== Dumped Configurations ===\")\n tools.logging.info(pprint.pformat(config.dump(), indent=2, width=120), logger=logger)\n logger.info(\"=== Output Directory ===\")\n logger.info(config.output.job_dirpath)\n\n logger.info(\"=== Start Evaluating ===\")\n logger.info(\"* Building diffusion model pipeline\")\n tools.logging.Formatter.indent_inc()\n pipeline = config.pipeline.build()\n if \"nf4\" not in config.pipeline.name and \"gguf\" not in config.pipeline.name:\n model = DiffusionModelStruct.construct(pipeline)\n tools.logging.Formatter.indent_dec()\n save_dirpath = os.path.join(config.output.running_job_dirpath, \"cache\")\n if config.save_model:\n if config.save_model.lower() in (\"false\", \"none\", \"null\", \"nil\"):\n save_model = False\n elif config.save_model.lower() in (\"true\", \"default\"):\n save_dirpath, save_model = os.path.join(config.output.running_job_dirpath, \"model\"), True\n else:\n save_dirpath, save_model = config.save_model, True\n else:\n save_model = False\n model = ptq(\n model,\n config.quant,\n cache=config.cache,\n load_dirpath=config.load_from,\n save_dirpath=save_dirpath,\n copy_on_save=config.copy_on_save,\n save_model=save_model,\n )\n if config.pipeline.lora is not None:\n load_from = \"\"\n if config.quant.enabled_smooth:\n if config.load_from and os.path.exists(os.path.join(config.load_from, \"smooth.pt\")):\n load_from = os.path.join(config.load_from, \"smooth.pt\")\n elif config.cache.path and os.path.exists(config.cache.path.smooth):\n load_from = config.cache.path.smooth\n elif os.path.exists(os.path.join(save_dirpath, \"smooth.pt\")):\n load_from = os.path.join(save_dirpath, \"smooth.pt\")\n logger.info(f\"* Loading smooth scales from {load_from}\")\n config.pipeline.load_lora(pipeline, smooth_cache=torch.load(load_from) if load_from else None)\n if config.text is not None and config.text.is_enabled():\n for encoder_name, encoder, tokenizer in config.pipeline.extract_text_encoders(pipeline):\n logger.info(f\"* Post-training quantizing the text encoder {encoder_name}\")\n patch_attention(encoder)\n patch_gemma_rms_norm(encoder)\n save_dirpath = os.path.join(save_dirpath, \"encoder\")\n setattr(\n pipeline,\n encoder_name,\n llm_ptq(\n encoder,\n tokenizer,\n config.text,\n cache=config.text_cache,\n load_dirpath=os.path.join(config.load_from, \"encoder\") if config.load_from else \"\",\n save_dirpath=save_dirpath,\n copy_on_save=config.copy_on_save,\n save_model=save_model,\n ),\n )\n config.eval.gen_root = config.eval.gen_root.format(\n output=config.output.running_dirpath, job=config.output.running_job_dirname\n )\n if config.skip_eval:\n if not config.skip_gen:\n logger.info(\"* Generating image\")\n tools.logging.Formatter.indent_inc()\n config.eval.generate(pipeline, task=config.pipeline.task)\n tools.logging.Formatter.indent_dec()\n else:\n logger.info(f\"* Evaluating model {'(skipping generation)' if config.skip_gen else ''}\")\n tools.logging.Formatter.indent_inc()\n results = config.eval.evaluate(pipeline, skip_gen=config.skip_gen, task=config.pipeline.task)\n tools.logging.Formatter.indent_dec()\n if results is not None:\n logger.info(f\"* Saving results to {config.output.job_dirpath}\")\n with open(config.output.get_running_job_path(\"results.json\"), \"w\") as f:\n json.dump(results, f, indent=2, sort_keys=True)\n config.output.unlock()\n\n\nif __name__ == \"__main__\":\n config, _, unused_cfgs, unused_args, unknown_args = DiffusionPtqRunConfig.get_parser().parse_known_args()\n assert isinstance(config, DiffusionPtqRunConfig)\n if len(unused_cfgs) > 0:\n tools.logging.warning(f\"Unused configurations: {unused_cfgs}\")\n if unused_args is not None:\n tools.logging.warning(f\"Unused arguments: {unused_args}\")\n assert len(unknown_args) == 0, f\"Unknown arguments: {unknown_args}\"\n try:\n main(config, logging_level=tools.logging.DEBUG)\n except Exception as e:\n tools.logging.Formatter.indent_reset()\n tools.logging.error(\"=== Error ===\")\n tools.logging.error(traceback.format_exc())\n tools.logging.shutdown()\n traceback.print_exc()\n config.output.unlock(error=True)\n raise e\n"
},
{
"path": "deepcompressor/app/diffusion/quant/__init__.py",
"content": "# -*- coding: utf-8 -*-\n\nfrom .activation import quantize_diffusion_activations\nfrom .config import DiffusionQuantCacheConfig, DiffusionQuantConfig\nfrom .quantizer import DiffusionActivationQuantizer, DiffusionWeightQuantizer\nfrom .rotate import rotate_diffusion\nfrom .smooth import smooth_diffusion\nfrom .weight import load_diffusion_weights_state_dict, quantize_diffusion_weights\n"
},
{
"path": "deepcompressor/app/diffusion/quant/activation.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Diffusion model activation quantization calibration module.\"\"\"\n\nimport gc\nimport typing as tp\n\nimport torch\nimport torch.nn as nn\nfrom tqdm import tqdm\n\nfrom deepcompressor.data.cache import IOTensorsCache\nfrom deepcompressor.data.common import TensorType\nfrom deepcompressor.utils import tools\n\nfrom ..nn.struct import (\n DiffusionAttentionStruct,\n DiffusionBlockStruct,\n DiffusionModelStruct,\n DiffusionModuleStruct,\n DiffusionTransformerBlockStruct,\n)\nfrom .config import DiffusionQuantConfig\nfrom .quantizer import DiffusionActivationQuantizer\nfrom .utils import get_needs_inputs_fn, get_needs_outputs_fn\n\n__all__ = [\"quantize_diffusion_activations\"]\n\n\n@torch.inference_mode()\ndef quantize_diffusion_block_activations( # noqa: C901\n layer: DiffusionBlockStruct | DiffusionModuleStruct,\n config: DiffusionQuantConfig,\n quantizer_state_dict: dict[str, dict[str, torch.Tensor | float | None]],\n layer_cache: dict[str, IOTensorsCache] | None = None,\n layer_kwargs: dict[str, tp.Any] | None = None,\n orig_state_dict: dict[str, torch.Tensor] | None = None,\n) -> dict[str, DiffusionActivationQuantizer]:\n \"\"\"Quantize the activations of a diffusion model block.\n\n Args:\n layer (`DiffusionBlockStruct` or `DiffusionModuleStruct`):\n The diffusion model block.\n config (`DiffusionQuantConfig`):\n The quantization configuration.\n quantizer_state_dict (`dict[str, dict[str, torch.Tensor | float | None]]`):\n The activation quantizers state dict cache.\n layer_cache (`dict[str, IOTensorsCache]`, *optional*):\n The layer cache.\n layer_kwargs (`dict[str, Any]`, *optional*):\n The layer keyword arguments.\n orig_state_dict (`dict[str, torch.Tensor]`, *optional*):\n The original state dictionary.\n\n Returns:\n `dict[str, DiffusionActivationQuantizer]`:\n The activation quantizers.\n \"\"\"\n logger = tools.logging.getLogger(f\"{__name__}.ActivationQuant\")\n logger.debug(\"- Quantizing layer %s\", layer.name)\n layer_cache = layer_cache or {}\n layer_kwargs = layer_kwargs or {}\n orig_state_dict = orig_state_dict or {}\n args_caches: list[\n tuple[\n str, # key\n TensorType,\n list[nn.Linear], # modules\n list[str], # module names\n nn.Module, # eval module\n str, # eval name\n dict[str, tp.Any], # eval kwargs\n list[tuple[nn.Parameter, torch.Tensor]], # original wgts\n ]\n ] = []\n In, Out = TensorType.Inputs, TensorType.Outputs # noqa: F841\n\n used_modules: set[nn.Module] = set()\n for module_key, module_name, module, parent, field_name in layer.named_key_modules():\n modules, orig_struct_wgts = None, {}\n if field_name in (\"k_proj\", \"v_proj\", \"add_q_proj\", \"add_v_proj\"):\n continue\n if field_name in (\"q_proj\", \"add_k_proj\", \"up_proj\"):\n grandparent = parent.parent\n assert isinstance(grandparent, DiffusionTransformerBlockStruct)\n if grandparent.parallel and parent.idx == 0:\n if orig_state_dict:\n orig_struct_wgts = {\n proj_module: (proj_module.weight, orig_state_dict[f\"{proj_name}.weight\"])\n for _, proj_name, proj_module, _, _ in grandparent.named_key_modules()\n }\n if field_name == \"q_proj\":\n assert isinstance(parent, DiffusionAttentionStruct)\n assert module_name == parent.q_proj_name\n modules, module_names = parent.qkv_proj, parent.qkv_proj_names\n if grandparent.ffn_struct is not None:\n modules.append(grandparent.ffn_struct.up_proj)\n module_names.append(grandparent.ffn_struct.up_proj_name)\n elif field_name == \"add_k_proj\":\n assert isinstance(parent, DiffusionAttentionStruct)\n assert module_name == parent.add_k_proj_name\n modules, module_names = parent.add_qkv_proj, parent.add_qkv_proj_names\n if grandparent.add_ffn_struct is not None:\n modules.append(grandparent.add_ffn_struct.up_proj)\n module_names.append(grandparent.add_ffn_struct.up_proj_name)\n else:\n assert field_name == \"up_proj\"\n if module in used_modules:\n continue\n assert module_name == grandparent.add_ffn_struct.up_proj_name\n assert grandparent.attn_structs[0].is_self_attn()\n eval_module, eval_name, eval_kwargs = grandparent.module, grandparent.name, layer_kwargs\n elif isinstance(parent, DiffusionAttentionStruct):\n eval_module, eval_name = parent.module, parent.name\n eval_kwargs = parent.filter_kwargs(layer_kwargs) if layer_kwargs else {}\n if orig_state_dict:\n orig_struct_wgts = {\n proj_module: (proj_module.weight, orig_state_dict[f\"{proj_name}.weight\"])\n for _, proj_name, proj_module, _, _ in parent.named_key_modules()\n }\n if field_name == \"q_proj\":\n assert module_name == parent.q_proj_name\n modules, module_names = parent.qkv_proj, parent.qkv_proj_names\n else:\n assert field_name == \"add_k_proj\"\n assert module_name == parent.add_k_proj_name\n modules, module_names = parent.add_qkv_proj, parent.add_qkv_proj_names\n if modules is None:\n assert module not in used_modules\n used_modules.add(module)\n orig_wgts = [(module.weight, orig_state_dict[f\"{module_name}.weight\"])] if orig_state_dict else None\n args_caches.append((module_key, In, [module], [module_name], module, module_name, None, orig_wgts))\n else:\n orig_wgts = []\n for proj_module in modules:\n assert proj_module not in used_modules\n used_modules.add(proj_module)\n if orig_state_dict:\n orig_wgts.append(orig_struct_wgts.pop(proj_module))\n orig_wgts.extend(orig_struct_wgts.values())\n orig_wgts = None if not orig_wgts else orig_wgts\n args_caches.append((module_key, In, modules, module_names, eval_module, eval_name, eval_kwargs, orig_wgts))\n # endregion\n quantizers: dict[str, DiffusionActivationQuantizer] = {}\n tools.logging.Formatter.indent_inc()\n for module_key, tensor_type, modules, module_names, eval_module, eval_name, eval_kwargs, orig_wgts in args_caches:\n if isinstance(modules[0], nn.Linear):\n channels_dim = -1\n assert all(isinstance(m, nn.Linear) for m in modules)\n elif isinstance(modules[0], nn.Conv2d):\n channels_dim = 1\n assert all(isinstance(m, nn.Conv2d) for m in modules)\n else:\n raise ValueError(f\"Unknown module type: {type(modules[0])}\")\n if tensor_type == TensorType.Inputs:\n cache_keys = [f\"{name}.input\" for name in module_names]\n quantizer_config = config.unsigned_ipts if getattr(modules[0], \"unsigned\", False) else config.ipts\n activations = layer_cache.get(module_names[0], IOTensorsCache()).inputs\n else:\n cache_keys = [f\"{name}.output\" for name in module_names]\n quantizer_config = config.opts\n activations = layer_cache.get(module_names[0], IOTensorsCache()).outputs\n quantizer = DiffusionActivationQuantizer(\n quantizer_config,\n channels_dim=channels_dim,\n develop_dtype=config.develop_dtype,\n key=module_key,\n tensor_type=tensor_type,\n )\n if quantizer.is_enabled():\n if cache_keys[0] not in quantizer_state_dict:\n logger.debug(\"- Calibrating %s\", \", \".join(cache_keys))\n quantizer.calibrate_dynamic_range(\n modules=modules,\n activations=activations,\n eval_module=eval_module,\n eval_inputs=layer_cache[eval_name].inputs if layer_cache else None,\n eval_kwargs=eval_kwargs,\n orig_weights=orig_wgts,\n )\n quantizer_state_dict[cache_keys[0]] = quantizer.state_dict()\n gc.collect()\n torch.cuda.empty_cache()\n else:\n quantizer.load_state_dict(quantizer_state_dict[cache_keys[0]], device=modules[0].weight.device)\n for cache_key in cache_keys:\n quantizers[cache_key] = quantizer\n del quantizer\n tools.logging.Formatter.indent_dec()\n return quantizers\n\n\n@torch.inference_mode()\ndef quantize_diffusion_activations(\n model: nn.Module | DiffusionModelStruct,\n config: DiffusionQuantConfig,\n quantizer_state_dict: dict[str, dict[str, torch.Tensor | float | None]] | None = None,\n orig_state_dict: dict[str, torch.Tensor] | None = None,\n) -> dict[str, dict[str, torch.Tensor | float | None]]:\n \"\"\"Quantize the activations of a diffusion model.\n\n Args:\n model (`nn.Module` or `DiffusionModelStruct`):\n The diffusion model.\n config (`DiffusionQuantConfig`):\n The quantization configuration.\n quantizer_state_dict (`dict[str, dict[str, torch.Tensor | float | None]]`, *optional*, defaults to `None`):\n The activation quantizers state dict cache.\n orig_state_dict (`dict[str, torch.Tensor]`, *optional*, defaults to `None`):\n The original state dictionary.\n\n Returns:\n `dict[str, dict[str, torch.Tensor | float | None]]`:\n The activation quantizers state dict cache.\n \"\"\"\n logger = tools.logging.getLogger(f\"{__name__}.ActivationQuant\")\n if not isinstance(model, DiffusionModelStruct):\n model = DiffusionModelStruct.construct(model)\n assert isinstance(model, DiffusionModelStruct)\n quantizer_state_dict = quantizer_state_dict or {}\n quantizers: dict[str, DiffusionActivationQuantizer] = {}\n skip_pre_modules = all(key in config.ipts.skips for key in model.get_prev_module_keys())\n skip_post_modules = all(key in config.ipts.skips for key in model.get_post_module_keys())\n if not quantizer_state_dict and config.needs_acts_quantizer_cache:\n with tools.logging.redirect_tqdm():\n for _, (layer, layer_cache, layer_kwargs) in tqdm(\n config.calib.build_loader().iter_layer_activations(\n model,\n needs_inputs_fn=get_needs_inputs_fn(model, config=config),\n needs_outputs_fn=get_needs_outputs_fn(model, config=config),\n skip_pre_modules=skip_pre_modules,\n skip_post_modules=skip_post_modules,\n ),\n desc=\"quantizing activations\",\n leave=False,\n total=model.num_blocks + int(not skip_post_modules) + int(not skip_pre_modules) * 3,\n dynamic_ncols=True,\n ):\n block_quantizers = quantize_diffusion_block_activations(\n layer=layer,\n config=config,\n quantizer_state_dict=quantizer_state_dict,\n layer_cache=layer_cache,\n layer_kwargs=layer_kwargs,\n orig_state_dict=orig_state_dict,\n )\n quantizers.update(block_quantizers)\n else:\n for _, layer in model.get_named_layers(\n skip_pre_modules=skip_pre_modules, skip_post_modules=skip_post_modules\n ).items():\n block_quantizers = quantize_diffusion_block_activations(\n layer=layer,\n config=config,\n quantizer_state_dict=quantizer_state_dict,\n orig_state_dict=orig_state_dict,\n )\n quantizers.update(block_quantizers)\n for _, module_name, module, _, _ in model.named_key_modules():\n ipts_quantizer = quantizers.get(f\"{module_name}.input\", None)\n opts_quantizer = quantizers.get(f\"{module_name}.output\", None)\n needs_quant_ipts = ipts_quantizer is not None and ipts_quantizer.is_enabled()\n needs_quant_opts = opts_quantizer is not None and opts_quantizer.is_enabled()\n if needs_quant_ipts or needs_quant_opts:\n logger.debug(\n \"- Quantizing %s (%s)\",\n module_name,\n (\"inputs\" if needs_quant_ipts else \"\")\n + (\" and \" if needs_quant_ipts and needs_quant_opts else \"\")\n + (\"outputs\" if needs_quant_opts else \"\"),\n )\n if needs_quant_ipts:\n ipts_quantizer.as_hook(is_output=False).register(module)\n if needs_quant_opts:\n opts_quantizer.as_hook(is_output=True).register(module)\n return quantizer_state_dict\n"
},
{
"path": "deepcompressor/app/diffusion/quant/config.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Quantization config.\"\"\"\n\nimport os\nfrom dataclasses import dataclass, field\n\nimport torch\nfrom omniconfig import configclass\n\nfrom deepcompressor.calib.config import (\n QuantRotationConfig,\n SearchBasedCalibGranularity,\n SearchBasedCalibObjective,\n SearchBasedCalibStrategy,\n SmoothTransfomerConfig,\n)\nfrom deepcompressor.data.utils.dtype import eval_dtype\nfrom deepcompressor.quantizer.config import QuantLowRankConfig\nfrom deepcompressor.utils.common import num2str\n\nfrom ..cache.config import DiffusionQuantCacheConfig\nfrom ..dataset.calib import DiffusionCalibCacheLoaderConfig\nfrom .quantizer.config import DiffusionModuleQuantizerConfig\n\n__all__ = [\"DiffusionQuantConfig\"]\n\n\n@configclass\n@dataclass(kw_only=True)\nclass DiffusionQuantConfig(DiffusionModuleQuantizerConfig):\n \"\"\"Diffusion model quantization configuration.\n\n Args:\n wgts (`DiffusionWeightQuantizerConfig`):\n The weight quantization configuration.\n ipts (`DiffusionActivationQuantizerConfig`):\n The input activation quantization configuration.\n opts (`DiffusionActivationQuantizerConfig`):\n The output activation quantization configuration.\n calib (`DiffusionCalibDatasetConfig`):\n The calibration dataset configuration.\n smooth (`TransfomerQuantSmoothConfig` or `None`, *optional*, defaults to `None`):\n The smooth quantization configuration.\n develop_dtype (`torch.dtype`, *optional*, defaults to `None`):\n The development data type.\n \"\"\"\n\n calib: DiffusionCalibCacheLoaderConfig\n rotation: QuantRotationConfig | None = None\n smooth: SmoothTransfomerConfig | None = None\n develop_dtype: torch.dtype = field(default_factory=lambda s=torch.float32: eval_dtype(s, with_quant_dtype=False))\n\n def __post_init__(self) -> None: # noqa: C901\n super().__post_init__()\n if self.rotation is not None and not self.rotation.transforms:\n self.rotation = None\n if self.smooth is not None:\n if not self.smooth.enabled_proj and not self.smooth.enabled_attn:\n self.smooth = None\n if self.enabled_smooth and self.smooth.enabled_proj and self.smooth.proj.allow_low_rank:\n if self.enabled_wgts:\n self.smooth.proj.allow_low_rank = self.wgts.enabled_low_rank\n if self.smooth.proj.allow_low_rank:\n self.smooth.proj.granularity = SearchBasedCalibGranularity.Layer\n else:\n self.smooth.proj.allow_low_rank = False\n if self.enabled_ipts:\n if self.ipts.enabled_calib_range and self.ipts.calib_range.granularity == SearchBasedCalibGranularity.Group:\n self.ipts.calib_range.granularity = SearchBasedCalibGranularity.ChannelGroup\n if self.ipts.static:\n assert self.ipts.smallest_group_shape[0] == -1, \"static quantization requires batch group size to be -1\"\n if self.enabled_opts:\n if self.opts.enabled_calib_range and self.opts.calib_range.granularity == SearchBasedCalibGranularity.Group:\n self.opts.calib_range.granularity = SearchBasedCalibGranularity.ChannelGroup\n if self.opts.static:\n assert self.opts.smallest_group_shape[0] == -1, \"static quantization requires batch group size to be -1\"\n self.organize()\n self.unsigned_ipts = self.ipts.for_unsigned()\n\n @property\n def enabled_rotation(self) -> bool:\n \"\"\"Whether to enable rotation.\"\"\"\n return self.rotation is not None and bool(self.rotation.transforms)\n\n @property\n def enabled_smooth(self) -> bool:\n \"\"\"Whether to enable smooth quantization.\"\"\"\n return self.smooth is not None\n\n @property\n def enabled_smooth_proj(self) -> bool:\n \"\"\"Whether to enable smooth quantization for projections.\"\"\"\n return self.enabled_smooth and self.smooth.enabled_proj\n\n @property\n def enabled_smooth_attn(self) -> bool:\n \"\"\"Whether to enable smooth quantization for attentions.\"\"\"\n return self.enabled_smooth and self.smooth.enabled_attn\n\n @property\n def needs_acts_quantizer_cache(self) -> bool:\n \"\"\"Whether to cache the activations quantizer settings.\"\"\"\n if self.enabled_ipts and self.ipts.needs_calib_data:\n return True\n if self.enabled_opts and self.opts.needs_calib_data:\n return True\n return False\n\n def generate_calib_dirname(self) -> str:\n name = \"\"\n if self.enabled_rotation:\n name += \"-rotate\"\n if self.rotation.random:\n name += \".rnd\"\n if self.enabled_smooth:\n name += \"-smooth\"\n if self.enabled_smooth_proj:\n name += \".proj\"\n if self.enabled_smooth_attn:\n name += \".attn\"\n calib_name = super().generate_calib_dirname()\n if calib_name:\n name += f\"-{calib_name}\"\n return name[1:] if name else name\n\n def generate_cache_dirpath(\n self, *, root: str, shift: bool, default_dtype: torch.dtype = torch.float16\n ) -> DiffusionQuantCacheConfig: # noqa: C901\n \"\"\"Generate the cache paths for the module quantization configuration.\"\"\"\n quant_names = self.generate_dirnames(default_dtype=default_dtype)\n if shift:\n quant_names.append(\"shift\")\n if self.enabled_wgts and self.wgts.enabled_low_rank:\n quant_names.extend(QuantLowRankConfig.generate_dirnames(self.wgts.low_rank, prefix=\"lowrank\"))\n if self.enabled_rotation:\n quant_names.extend(self.rotation.generate_dirnames(prefix=\"rotate\"))\n smooth_dirpath = \"\"\n if self.enabled_smooth:\n quant_names.extend(self.smooth.generate_dirnames(prefix=\"smooth\"))\n smooth_dirpath = os.path.join(\"smooth\", *quant_names)\n branch_dirpath = \"\"\n if self.enabled_wgts and self.wgts.enabled_low_rank:\n quant_names.extend(self.wgts.low_rank.generate_dirnames(prefix=\"lowrank\"))\n branch_dirpath = os.path.join(\"branch\", *quant_names)\n wgts_dirpath = \"\"\n if self.enabled_wgts and self.wgts.needs_calib_data:\n quant_names.extend(self.wgts.calib_range.generate_dirnames(prefix=\"w.range\"))\n wgts_dirpath = os.path.join(\"wgts\", *quant_names)\n if self.enabled_wgts and self.wgts.enabled_gptq:\n quant_names.extend(self.wgts.kernel_gptq.generate_dirnames(prefix=\"w.kernel\"))\n acts_dirpath = \"\"\n if self.needs_acts_quantizer_cache:\n if self.enabled_ipts and self.ipts.needs_calib_data:\n quant_names.extend(self.ipts.calib_range.generate_dirnames(prefix=\"x.range\"))\n if self.enabled_opts and self.opts.needs_calib_data:\n quant_names.extend(self.opts.calib_range.generate_dirnames(prefix=\"y.range\"))\n acts_dirpath = os.path.join(\"acts\", *quant_names)\n cache_dirpath = DiffusionQuantCacheConfig(\n smooth=smooth_dirpath, branch=branch_dirpath, wgts=wgts_dirpath, acts=acts_dirpath\n ).simplify(type(self)._key_map)\n cache_dirpath = cache_dirpath.add_parent_dirs(*self.calib.generate_dirnames())\n cache_dirpath = cache_dirpath.add_parent_dirs(root, \"diffusion\", \"cache\", \"quant\")\n return cache_dirpath\n\n def generate_default_dirname(self) -> str: # noqa: C901\n \"\"\"Generate output directory name for evaluating a large language model.\"\"\"\n key_map = type(self)._key_map\n\n def simplify_skips(skips):\n return set(\n DiffusionQuantCacheConfig.simplify_path(\"skip.[{}]\".format(\"+\".join(skips)), key_map=key_map)[\n 6:-1\n ].split(\"+\")\n )\n\n skip_name, y_skips, w_skips, x_skips = \"\", set(), set(), set()\n if self.enabled_opts and self.opts.skips:\n y_skips = simplify_skips(self.opts.skips)\n if self.enabled_ipts and self.ipts.skips:\n x_skips = simplify_skips(self.ipts.skips)\n if self.enabled_wgts and self.wgts.skips:\n w_skips = simplify_skips(self.wgts.skips)\n skips_map = {}\n if y_skips or x_skips or w_skips:\n skip_name = \"-skip\"\n skip_name_list: list[tuple[str, set]] = []\n if y_skips:\n skip_name_list.append((\"y\", y_skips))\n if x_skips:\n skip_name_list.append((\"x\", x_skips))\n if w_skips:\n skip_name_list.append((\"w\", w_skips))\n # sort the keys by the number of elements in the set\n skip_name_list = sorted(skip_name_list, key=lambda x: (len(x[1]), x[0]), reverse=False)\n skips_map = {k: v for k, v in skip_name_list} # noqa: C416\n skip_name_map: dict[str, set] = {}\n skip_0, skip_0_names = skip_name_list[0]\n skip_name_map[skip_0] = skip_0_names\n if len(skip_name_list) > 1:\n skip_1, skip_1_names = skip_name_list[1]\n if skip_1_names.issuperset(skip_0_names):\n skip_1_names = skip_1_names - skip_0_names\n skip_1_names.add(f\"[{skip_0}]\")\n skip_name_map[skip_1] = skip_1_names\n if len(skip_name_list) > 2:\n skip_2, skip_2_names = skip_name_list[2]\n if skip_2_names.issuperset(skip_name_list[1][1]): # skip_1_names may be modified\n skip_2_names = skip_2_names - skip_name_list[1][1]\n skip_2_names.add(f\"[{skip_1}]\")\n if skip_2_names.issuperset(skip_0_names):\n skip_2_names = skip_2_names - skip_0_names\n skip_2_names.add(f\"[{skip_0}]\")\n skip_name_map[skip_2] = skip_2_names\n if \"y\" in skip_name_map:\n skip_name += f\".y.[{'+'.join(sorted(skip_name_map['y']))}]\"\n if \"x\" in skip_name_map:\n skip_name += f\".x.[{'+'.join(sorted(skip_name_map['x']))}]\"\n if \"w\" in skip_name_map:\n skip_name += f\".w.[{'+'.join(sorted(skip_name_map['w']))}]\"\n del skip_name_list, skip_name_map\n extra_name = \"\"\n if self.enabled_extra_wgts:\n extra_name = \"-extra.[{}]\".format(\"+\".join(sorted(simplify_skips(self.extra_wgts.includes))))\n lowrank_name = \"\"\n if self.enabled_wgts and self.wgts.enabled_low_rank:\n lowrank_name = f\"-low.r{num2str(self.wgts.low_rank.rank)}\"\n if self.wgts.low_rank.num_iters > 1:\n lowrank_name += f\".i{num2str(self.wgts.low_rank.num_iters)}\"\n if self.wgts.low_rank.early_stop:\n lowrank_name += \".e\"\n if self.wgts.low_rank.exclusive:\n lowrank_name += \".s\"\n if self.wgts.low_rank.compensate:\n lowrank_name += \".c\"\n if self.wgts.low_rank.objective != SearchBasedCalibObjective.OutputsError:\n lowrank_name += f\".{self.wgts.low_rank.objective.name}\"\n if self.wgts.low_rank.skips:\n lowrank_skips = simplify_skips(self.wgts.low_rank.skips)\n if \"w\" in skips_map and lowrank_skips.issuperset(skips_map[\"w\"]):\n lowrank_skips = lowrank_skips - skips_map[\"w\"]\n lowrank_skips.add(\"[w]\")\n lowrank_name += \".skip.[{}]\".format(\"+\".join(sorted(lowrank_skips)))\n rotation_name = \"\"\n if self.enabled_rotation:\n rotation_name = \"-rot\"\n if self.rotation.random:\n rotation_name += \".rnd\"\n rotation_name += \".[{}]\".format(\"+\".join(sorted(simplify_skips(self.rotation.transforms))))\n smooth_name = \"\"\n if self.enabled_smooth:\n smooth_name = \"-smth\"\n if self.smooth.enabled_proj:\n smooth_name += \".proj\"\n if self.smooth.proj.granularity != SearchBasedCalibGranularity.Layer:\n smooth_name += f\".{self.smooth.proj.granularity.name}\"\n if self.smooth.proj.strategy != SearchBasedCalibStrategy.Manual:\n smooth_name += f\".{self.smooth.proj.strategy.name}\"\n if self.smooth.proj.alpha <= 0:\n smooth_name += f\".a{num2str(self.smooth.proj.alpha)}\"\n if self.smooth.proj.beta <= 0:\n smooth_name += f\".b{num2str(self.smooth.proj.beta)}\"\n else:\n smooth_name += f\".a{num2str(self.smooth.proj.alpha)}\"\n smooth_name += f\".b{num2str(self.smooth.proj.beta)}\"\n xspan_eq_wspan = True\n for xspan, wspan in self.smooth.proj.spans:\n if xspan != wspan:\n xspan_eq_wspan = False\n break\n if xspan_eq_wspan:\n smooth_name += \".[{}]\".format(\"+\".join(xspan.name for xspan, _ in self.smooth.proj.spans))\n else:\n smooth_name += \".[{}]\".format(\n \"+\".join(f\"x.{xspan.name}.w.{wspan.name}\" for xspan, wspan in self.smooth.proj.spans)\n )\n if self.smooth.proj.allow_low_rank:\n smooth_name += \".lr\"\n if not self.smooth.proj.allow_b_quant or not self.smooth.proj.allow_a_quant:\n smooth_name += \".no.[\"\n if not self.smooth.proj.allow_a_quant:\n smooth_name += \"a+\"\n if not self.smooth.proj.allow_b_quant:\n smooth_name += \"b+\"\n smooth_name = smooth_name[:-1] + \"]\"\n if self.smooth.proj.skips:\n smooth_skips = simplify_skips(self.smooth.proj.skips)\n if \"w\" in skips_map and smooth_skips.issuperset(skips_map[\"w\"]):\n smooth_skips = smooth_skips - skips_map[\"w\"]\n smooth_skips.add(\"[w]\")\n smooth_name += \".skip.[{}]\".format(\"+\".join(sorted(smooth_skips)))\n if self.smooth.enabled_attn:\n smooth_name += \".yx\"\n if self.smooth.attn.granularity != SearchBasedCalibGranularity.Layer:\n smooth_name += f\".{self.smooth.attn.granularity.name}\"\n if self.smooth.attn.strategy != SearchBasedCalibStrategy.Manual:\n smooth_name += f\".{self.smooth.attn.strategy.name}\"\n if self.smooth.attn.alpha <= 0:\n smooth_name += f\".a{num2str(self.smooth.attn.alpha)}\"\n if self.smooth.attn.beta <= 0:\n smooth_name += f\".b{num2str(self.smooth.attn.beta)}\"\n else:\n smooth_name += f\".a{num2str(self.smooth.attn.alpha)}\"\n smooth_name += f\".b{num2str(self.smooth.attn.beta)}\"\n xspan_eq_yspan = True\n for xspan, yspan in self.smooth.attn.spans:\n if xspan != yspan:\n xspan_eq_yspan = False\n break\n if xspan_eq_yspan:\n smooth_name += \".[{}]\".format(\"+\".join(xspan.name for xspan, _ in self.smooth.attn.spans))\n else:\n smooth_name += \".[{}]\".format(\n \"+\".join(f\"x.{xspan.name}.y.{yspan.name}\" for xspan, yspan in self.smooth.attn.spans)\n )\n gptq_name = \"\"\n if self.enabled_wgts and self.wgts.kernel_gptq is not None:\n gptq_name = \"-gptq\"\n if self.wgts.kernel_gptq.skips:\n gptq_skips = simplify_skips(self.wgts.kernel_gptq.skips)\n if \"w\" in skips_map and gptq_skips.issuperset(skips_map[\"w\"]):\n gptq_skips = gptq_skips - skips_map[\"w\"]\n gptq_skips.add(\"[w]\")\n gptq_name += \".skip.[{}]\".format(\"+\".join(sorted(gptq_skips)))\n wrange_name = \"\"\n if (\n self.enabled_wgts\n and self.wgts.enabled_calib_range\n and (self.wgts.calib_range.needs_search or self.wgts.calib_range.ratio != 1)\n ):\n wrange_name = \"-w.range\"\n if self.wgts.calib_range.needs_search:\n if self.wgts.calib_range.granularity != SearchBasedCalibGranularity.Group:\n wrange_name += f\".{self.wgts.calib_range.granularity.name}\"\n if self.wgts.calib_range.objective != SearchBasedCalibObjective.OutputsError:\n wrange_name += f\".{self.wgts.calib_range.objective.name}\"\n if self.wgts.calib_range.degree != 2:\n wrange_name += f\".d{num2str(self.wgts.calib_range.degree)}\"\n wrange_name += f\".[{num2str(self.wgts.calib_range.max_shrink)}\"\n wrange_name += f\".{num2str(self.wgts.calib_range.max_expand)}\"\n wrange_name += f\".g{self.wgts.calib_range.num_grids}]\"\n else:\n wrange_name += f\".r{num2str(self.wgts.calib_range.ratio)}\"\n if self.wgts.calib_range.skips:\n wrange_skips = simplify_skips(self.wgts.calib_range.skips)\n if \"w\" in skips_map and wrange_skips.issuperset(skips_map[\"w\"]):\n wrange_skips = wrange_skips - skips_map[\"w\"]\n wrange_skips.add(\"[w]\")\n wrange_name += \".skip.[{}]\".format(\"+\".join(sorted(wrange_skips)))\n xrange_name = \"\"\n if (\n self.enabled_ipts\n and self.ipts.enabled_calib_range\n and (self.ipts.calib_range.needs_search or self.ipts.calib_range.ratio != 1)\n ):\n xrange_name = \"-x.range\"\n if self.ipts.calib_range.needs_search:\n if self.ipts.calib_range.granularity != SearchBasedCalibGranularity.Group:\n xrange_name += f\".{self.ipts.calib_range.granularity.name}\"\n if self.ipts.calib_range.objective != SearchBasedCalibObjective.OutputsError:\n xrange_name += f\".{self.ipts.calib_range.objective.name}\"\n if self.ipts.calib_range.degree != 2:\n xrange_name += f\".d{num2str(self.ipts.calib_range.degree)}\"\n xrange_name += f\".[{num2str(self.ipts.calib_range.max_shrink)}\"\n xrange_name += f\".{num2str(self.ipts.calib_range.max_expand)}\"\n xrange_name += f\".g{self.ipts.calib_range.num_grids}]\"\n else:\n xrange_name += f\".r{num2str(self.ipts.calib_range.ratio)}\"\n if self.ipts.calib_range.skips:\n xrange_skips = simplify_skips(self.ipts.calib_range.skips)\n if \"x\" in skips_map and xrange_skips.issuperset(skips_map[\"x\"]):\n xrange_skips = xrange_skips - skips_map[\"x\"]\n xrange_skips.add(\"[x]\")\n xrange_name += \".skip.[{}]\".format(\"+\".join(sorted(xrange_skips)))\n yrange_name = \"\"\n if (\n self.enabled_opts\n and self.opts.enabled_calib_range\n and (self.opts.calib_range.needs_search or self.opts.calib_range.ratio != 1)\n ):\n yrange_name = \"-y.range\"\n if self.opts.calib_range.needs_search:\n if self.opts.calib_range.granularity != SearchBasedCalibGranularity.Group:\n yrange_name += f\".{self.opts.calib_range.granularity.name}\"\n if self.opts.calib_range.objective != SearchBasedCalibObjective.OutputsError:\n yrange_name += f\".{self.opts.calib_range.objective.name}\"\n if self.opts.calib_range.degree != 2:\n yrange_name += f\".d{num2str(self.opts.calib_range.degree)}\"\n yrange_name += f\".[{num2str(self.opts.calib_range.max_shrink)}\"\n yrange_name += f\".{num2str(self.opts.calib_range.max_expand)}\"\n yrange_name += f\".g{self.opts.calib_range.num_grids}]\"\n else:\n yrange_name += f\".r{num2str(self.opts.calib_range.ratio)}\"\n if self.opts.calib_range.skips:\n yrange_skips = simplify_skips(self.opts.calib_range.skips)\n if \"y\" in skips_map and yrange_skips.issuperset(skips_map[\"y\"]):\n yrange_skips = yrange_skips - skips_map[\"y\"]\n yrange_skips.add(\"[y]\")\n yrange_name += \".skip.[{}]\".format(\"+\".join(sorted(yrange_skips)))\n name = (\n skip_name\n + extra_name\n + lowrank_name\n + rotation_name\n + smooth_name\n + gptq_name\n + wrange_name\n + xrange_name\n + yrange_name\n )\n name = name[1:] if name else \"default\"\n name += f\"-{self.calib.generate_dirnames()[0]}\"\n return name\n\n @classmethod\n def set_key_map(cls, key_map: dict[str, set[str]]) -> None:\n \"\"\"Set the key map for the language model quantization configuration.\n\n Args:\n key_map (dict[str, set[str]]): The key map.\n \"\"\"\n cls._key_map = key_map\n\n def organize(self) -> dict[str, bool]: # noqa: C901\n \"\"\"Organize the flags for the diffusion model quantization configuration.\n\n Returns:\n dict[str, bool]: The organized flags.\n \"\"\"\n key_map = type(self)._key_map\n wgts_skip_set, ipts_skip_set, opts_skip_set = set(), set(), set()\n if self.wgts is not None:\n wgts_skips = []\n for skip in self.wgts.skips:\n wgts_skips.extend(list(key_map[skip]))\n wgts_skip_set = set(wgts_skips)\n self.wgts.skips = sorted(wgts_skip_set)\n if self.wgts.kernel_gptq is not None:\n wgts_kernel_gptq_skips = []\n for skip in self.wgts.kernel_gptq.skips:\n wgts_kernel_gptq_skips.extend(list(key_map[skip]))\n self.wgts.kernel_gptq.skips = sorted(set(wgts_kernel_gptq_skips) - wgts_skip_set)\n if self.wgts.low_rank is not None:\n wgts_low_rank_skips = []\n for skip in self.wgts.low_rank.skips:\n wgts_low_rank_skips.extend(list(key_map[skip]))\n self.wgts.low_rank.skips = sorted(set(wgts_low_rank_skips) - wgts_skip_set)\n if self.wgts.calib_range is not None:\n wgts_calib_range_skips = []\n for skip in self.wgts.calib_range.skips:\n wgts_calib_range_skips.extend(list(key_map[skip]))\n self.wgts.calib_range.skips = sorted(set(wgts_calib_range_skips) - wgts_skip_set)\n if self.extra_wgts is not None:\n extra_includes = []\n for include in self.extra_wgts.includes:\n extra_includes.extend(list(key_map[include]))\n extra_includes_set = set(extra_includes) - wgts_skip_set\n self.extra_wgts.includes = sorted(extra_includes_set)\n if not self.extra_wgts.is_enabled():\n self.extra_wgts = None\n if self.ipts is not None:\n ipts_skips = []\n for skip in self.ipts.skips:\n ipts_skips.extend(list(key_map[skip]))\n ipts_skip_set = set(ipts_skips)\n self.ipts.skips = sorted(ipts_skip_set)\n if self.ipts.calib_range is not None:\n ipts_calib_range_skips = []\n for skip in self.ipts.calib_range.skips:\n ipts_calib_range_skips.extend(list(key_map[skip]))\n self.ipts.calib_range.skips = sorted(set(ipts_calib_range_skips) - ipts_skip_set)\n if self.opts is not None:\n opts_skips = []\n for skip in self.opts.skips:\n opts_skips.extend(list(key_map[skip]))\n opts_skip_set = set(opts_skips)\n self.opts.skips = sorted(opts_skip_set)\n if self.opts.calib_range is not None:\n opts_calib_range_skips = []\n for skip in self.opts.calib_range.skips:\n opts_calib_range_skips.extend(list(key_map[skip]))\n self.opts.calib_range.skips = sorted(set(opts_calib_range_skips) - opts_skip_set)\n if self.smooth is not None and self.smooth.proj is not None:\n smooth_proj_skips = []\n for skip in self.smooth.proj.skips:\n smooth_proj_skips.extend(list(key_map[skip]))\n self.smooth.proj.skips = sorted(set(smooth_proj_skips) - (wgts_skip_set & ipts_skip_set))\n if self.rotation is not None:\n rotation_transforms = []\n for transform in self.rotation.transforms:\n rotation_transforms.extend(list(key_map[transform]))\n self.rotation.transforms = sorted(set(rotation_transforms))\n"
},
{
"path": "deepcompressor/app/diffusion/quant/quantizer/__init__.py",
"content": "# -*- coding: utf-8 -*-\n\nfrom .config import DiffusionModuleQuantizerConfig\nfrom .quantizer import DiffusionActivationQuantizer, DiffusionWeightQuantizer\n"
},
{
"path": "deepcompressor/app/diffusion/quant/quantizer/config.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Quantizatizer config.\"\"\"\n\nfrom dataclasses import dataclass, field\n\nimport torch\nfrom omniconfig import configclass\n\nfrom deepcompressor.calib.config import SkipBasedDynamicRangeCalibConfig, SkipBasedQuantLowRankCalibConfig\nfrom deepcompressor.data.dtype import QuantDataType\nfrom deepcompressor.quantizer.config import QuantizerConfig\nfrom deepcompressor.quantizer.kernel import QuantGptqConfig\nfrom deepcompressor.utils.config import EnableConfig, IncludeBasedConfig, SkipBasedConfig\n\n__all__ = [\n \"DiffusionQuantizerConfig\",\n \"DiffusionWeightQuantizerConfig\",\n \"DiffusionActivationQuantizerConfig\",\n \"DiffusionModuleQuantizerConfig\",\n]\n\n\n@configclass\n@dataclass\nclass DiffusionGPTQConfig(SkipBasedConfig, QuantGptqConfig):\n \"\"\"Configuration for GPTQ quantization.\n\n Args:\n damp_percentage (`float`, *optional*, defaults to `0.01`):\n The percentage of damping.\n block_size (`int`, *optional*, defaults to `128`):\n The block size of the GPTQ quantization.\n num_inv_tries (`int`, *optional*, defaults to `200`):\n The number of tries for the inverse.\n hessian_block_size (`int`, *optional*, defaults to `-1`):\n The block size when calculing the Hessian.\n skips: list[str] = field(default_factory=list)\n \"\"\"\n\n pass\n\n\n@configclass\n@dataclass\nclass DiffusionQuantizerConfig(QuantizerConfig):\n \"\"\"Diffusion model quantizer configuration.\n\n Args:\n dtype (`QuantDataType` or `None`, *optional*, defaults to `None`):\n The quantization data type.\n zero_point (`ZeroPointDomain` or `None`, *optional*, defaults to `None`):\n The zero-point domain.\n group_shapes (`Sequence[Sequence[int]]`, *optional*, defaults to `((-1, -1, -1),)`):\n The shapes for per-group quantization.\n scale_dtypes (`Sequence[torch.dtype | QuantDataType | None]`, *optional*, defaults to `(None,)`):\n The quantization scale data type for per-group quantization.\n static (`bool`, *optional*, defaults to `False`):\n Whether to use static quantization.\n kernel_gptq (`DiffusionGPTQConfig` or `None`, *optional*, defaults to `None`):\n The gptq quantization configuration.\n low_rank (`SkipBasedQuantLowRankCalibConfig` or `None`, *optional*, defaults to `None`):\n The quantization low-rank branch calibration configuration.\n calib_range (`DynamicRangeCalibConfig` or `None`, *optional*, defaults to `None`):\n The quantizatizer dynamic range calibration configuration.\n \"\"\"\n\n static: bool = False\n kernel_gptq: DiffusionGPTQConfig | None = None\n low_rank: SkipBasedQuantLowRankCalibConfig | None = None\n calib_range: SkipBasedDynamicRangeCalibConfig | None = None\n\n def __post_init__(self) -> None:\n super().__post_init__()\n if self.quant_dtype is None:\n self.static = False\n self.kernel_gptq = None\n self.low_rank = None\n self.calib_range = None\n if self.kernel_gptq is not None and not self.kernel_gptq.is_enabled():\n self.kernel_gptq = None\n if self.static and self.calib_range is None:\n self.calib_range = SkipBasedDynamicRangeCalibConfig()\n if self.low_rank is not None and not self.low_rank.is_enabled():\n self.low_rank = None\n\n @property\n def enabled_gptq(self) -> bool:\n \"\"\"Whether quantization kernel calibration is enabled.\"\"\"\n return self.kernel_gptq is not None and self.kernel_gptq.is_enabled()\n\n @property\n def enabled_low_rank(self) -> bool:\n \"\"\"Whether quantization SVD calibration is enabled.\"\"\"\n return self.low_rank is not None and self.low_rank.is_enabled()\n\n @property\n def enabled_calib_range(self) -> bool:\n \"\"\"Whether quantization dynamic range calibration is enabled.\"\"\"\n return self.calib_range is not None\n\n def generate_calib_dirname(self) -> str:\n \"\"\"Generate the name for quantization calibration.\n\n Returns:\n str: The name.\n \"\"\"\n name = \"\"\n if self.static:\n name += \".static\"\n if self.enabled_gptq:\n name += \".gptq\"\n if self.enabled_low_rank:\n name += \".lowrank\"\n if self.enabled_calib_range and (self.calib_range.needs_search or self.calib_range.ratio != 1):\n name += \".range\"\n return name[1:] if name else \"\"\n\n\n@configclass\n@dataclass\nclass SkipBasedDiffusionQuantizerConfig(SkipBasedConfig, DiffusionQuantizerConfig):\n \"\"\"Diffusion model quantizer configuration.\n\n Args:\n dtype (`QuantDataType` or `None`, *optional*, defaults to `None`):\n The quantization data type.\n zero_point (`ZeroPointDomain` or `None`, *optional*, defaults to `None`):\n The zero-point domain.\n group_shapes (`Sequence[Sequence[int]]`, *optional*, defaults to `((-1, -1, -1),)`):\n The shapes for per-group quantization.\n scale_dtypes (`Sequence[torch.dtype | QuantDataType | None]`, *optional*, defaults to `(None,)`):\n The quantization scale data type for per-group quantization.\n skips (`[str]`, *optional*, defaults to `[]`):\n The keys of the modules to skip.\n static (`bool`, *optional*, defaults to `False`):\n Whether to use static quantization.\n kernel_gptq (`DiffusionGPTQConfig` or `None`, *optional*, defaults to `None`):\n The gptq quantization configuration.\n low_rank (`SkipBasedQuantLowRankCalibConfig` or `None`, *optional*, defaults to `None`):\n The quantization low-rank branch calibration configuration.\n calib_range (`DynamicRangeCalibConfig` or `None`, *optional*, defaults to `None`):\n The quantizatizer dynamic range calibration configuration.\n \"\"\"\n\n def __post_init__(self) -> None:\n super().__post_init__()\n if self.quant_dtype is None:\n self.skips.clear()\n\n\n@configclass\n@dataclass\nclass DiffusionWeightQuantizerConfig(SkipBasedDiffusionQuantizerConfig):\n \"\"\"Diffusion model weight quantizer configuration.\n\n Args:\n dtype (`QuantDataType` or `None`, *optional*, defaults to `None`):\n The quantization data type.\n zero_point (`ZeroPointDomain` or `None`, *optional*, defaults to `None`):\n The zero-point domain.\n group_shapes (`Sequence[Sequence[int]]`, *optional*, defaults to `((-1, -1, -1),)`):\n The shapes for per-group quantization.\n scale_dtypes (`Sequence[torch.dtype | QuantDataType | None]`, *optional*, defaults to `(None,)`):\n The quantization scale data type for per-group quantization.\n skips (`list[str]`, *optional*, defaults to `[]`):\n The keys of the modules to skip.\n low_rank (`SkipBasedQuantLowRankCalibConfig` or `None`, *optional*, defaults to `None`):\n The quantization low-rank branch calibration configuration.\n calib_range (`DynamicRangeCalibConfig` or `None`, *optional*, defaults to `None`):\n The quantizatizer dynamic range calibration configuration.\n \"\"\"\n\n static: bool = field(init=False, default=True)\n\n @property\n def needs_calib_data(self) -> bool:\n return self.enabled_calib_range and self.calib_range.needs_search\n\n\n@configclass\n@dataclass\nclass DiffusionActivationQuantizerConfig(SkipBasedDiffusionQuantizerConfig):\n \"\"\"Diffusion model activation quantizer configuration.\n\n Args:\n dtype (`QuantDataType` or `None`, *optional*, defaults to `None`):\n The quantization data type.\n zero_point (`ZeroPointDomain` or `None`, *optional*, defaults to `None`):\n The zero-point domain.\n group_shapes (`Sequence[Sequence[int]]`, *optional*, defaults to `((-1, -1, -1),)`):\n The shapes for per-group quantization.\n scale_dtypes (`Sequence[torch.dtype | QuantDataType | None]`, *optional*, defaults to `(None,)`):\n The quantization scale data type for per-group quantization.\n skips (`list[str]`, *optional*, defaults to `[]`):\n The keys of the modules to skip.\n static (`bool`, *optional*, defaults to `False`):\n Whether to use static quantization.\n calib_range (`DynamicRangeCalibConfig` or `None`, *optional*, defaults to `None`):\n The quantizatizer dynamic range calibration configuration.\n allow_unsigned (`bool`, *optional*, defaults to `False`):\n Whether to allow unsigned data type for activation quantization.\n \"\"\"\n\n kernel_gptq: None = field(init=False, default=None)\n low_rank: None = field(init=False, default=None)\n allow_unsigned: bool = False\n\n @property\n def needs_calib_data(self) -> bool:\n return self.enabled_calib_range and (self.calib_range.needs_search or self.static)\n\n def generate_dirnames(\n self,\n *,\n prefix: str = \"\",\n shape: torch.Size | tuple[int, ...] = (1024, 1024, 16, 16),\n default_dtype: torch.dtype = torch.float16,\n **kwargs,\n ) -> list[str]:\n \"\"\"Get the directory names of the quantization configuration.\n\n Args:\n prefix (`str`, *optional*, defaults to `\"\"`):\n The prefix for the directory names.\n shape (`torch.Size` or `tuple[int, ...]`, *optional*, defaults to `(1024, 1024, 16, 16)`):\n The shape of the tensor to be quantized.\n default_dtype (`torch.dtype`, *optional*, defaults to `torch.float16`):\n The dtype of the tensor to be quantized.\n\n Returns:\n `list[str]`:\n The directory names of the quantization configuration.\n - The number of effective bits.\n - The name of the quantization data type.\n - The name of the group shapes.\n - The name of the modules to skip.\n \"\"\"\n names = super().generate_dirnames(prefix=prefix, shape=shape, default_dtype=default_dtype)\n if self.allow_unsigned:\n names[1] += \".u\"\n return names\n\n def for_unsigned(self) -> \"DiffusionActivationQuantizerConfig\":\n \"\"\"get the quantizer configuration for unsigned activations.\n\n Returns:\n `DiffusionActivationQuantizerConfig`:\n The quantizer configuration for unsigned activations.\n \"\"\"\n if isinstance(self.dtype, QuantDataType) and self.allow_unsigned:\n return DiffusionActivationQuantizerConfig(\n dtype=self.dtype.to_unsigned(),\n zero_point=self.zero_point,\n group_shapes=self.group_shapes,\n scale_dtypes=self.scale_dtypes,\n skips=self.skips,\n static=self.static,\n calib_range=self.calib_range,\n allow_unsigned=self.allow_unsigned,\n )\n else:\n return self\n\n\n@configclass\n@dataclass\nclass DiffusionExtraWeightQuantizerConfig(IncludeBasedConfig, DiffusionQuantizerConfig):\n \"\"\"Diffusion model extra weight quantizer configuration.\n\n Args:\n dtype (`QuantDataType` or `None`, *optional*, defaults to `None`):\n The quantization data type.\n zero_point (`ZeroPointDomain` or `None`, *optional*, defaults to `None`):\n The zero-point domain.\n group_shapes (`Sequence[Sequence[int]]`, *optional*, defaults to `((-1, -1, -1),)`):\n The shapes for per-group quantization.\n scale_dtypes (`Sequence[torch.dtype | QuantDataType | None]`, *optional*, defaults to `(None,)`):\n The quantization scale data type for per-group quantization.\n includes (`list[str]`, *optional*, defaults to `[]`):\n The keys of the modules to include.\n low_rank (`SkipBasedQuantLowRankCalibConfig` or `None`, *optional*, defaults to `None`):\n The quantization low-rank branch calibration configuration.\n calib_range (`DynamicRangeCalibConfig` or `None`, *optional*, defaults to `None`):\n The quantizatizer dynamic range calibration configuration.\n \"\"\"\n\n static: bool = field(init=False, default=True)\n kernel_gptq: DiffusionGPTQConfig | None = field(init=False, default=None)\n low_rank: SkipBasedQuantLowRankCalibConfig | None = field(init=False, default=None)\n calib_range: SkipBasedDynamicRangeCalibConfig | None = field(init=False, default=None)\n\n @property\n def needs_calib_data(self) -> bool:\n return self.enabled_calib_range and self.calib_range.needs_search\n\n\n@configclass\n@dataclass(kw_only=True)\nclass DiffusionModuleQuantizerConfig(EnableConfig):\n \"\"\"Diffusion model module quantizer configuration.\n\n Args:\n wgts (`DiffusionWeightQuantizerConfig`):\n The weight quantization configuration.\n ipts (`DiffusionActivationQuantizerConfig`):\n The input activation quantization configuration.\n opts (`DiffusionActivationQuantizerConfig`):\n The output activation quantization configuration.\n \"\"\"\n\n wgts: DiffusionWeightQuantizerConfig\n ipts: DiffusionActivationQuantizerConfig\n opts: DiffusionActivationQuantizerConfig\n extra_wgts: DiffusionExtraWeightQuantizerConfig | None = None\n unsigned_ipts: DiffusionActivationQuantizerConfig = field(init=False)\n\n def is_enabled(self):\n return self.enabled_wgts or self.enabled_ipts or self.enabled_opts\n\n @property\n def enabled_wgts(self) -> bool:\n \"\"\"Whether to enable weight quantization.\"\"\"\n return self.wgts is not None and self.wgts.is_enabled()\n\n @property\n def enabled_ipts(self) -> bool:\n \"\"\"Whether to enable activation quantization.\"\"\"\n return self.ipts is not None and self.ipts.is_enabled()\n\n @property\n def enabled_opts(self) -> bool:\n \"\"\"Whether to enable activation quantization.\"\"\"\n return self.opts is not None and self.opts.is_enabled()\n\n @property\n def enabled_extra_wgts(self) -> bool:\n \"\"\"Whether to enable extra weight quantization.\"\"\"\n return self.extra_wgts is not None and self.extra_wgts.is_enabled()\n\n def __post_init__(self) -> None:\n if self.enabled_opts:\n raise NotImplementedError(\"Output activation quantization is not supported yet.\")\n if self.wgts.is_enabled() and self.extra_wgts is not None:\n self.extra_wgts.includes = list(filter(lambda key: key not in self.wgts.skips, self.extra_wgts.includes))\n if self.extra_wgts.is_enabled():\n self.extra_wgts.kernel_gptq = self.wgts.kernel_gptq\n self.extra_wgts.low_rank = self.wgts.low_rank\n self.extra_wgts.calib_range = self.wgts.calib_range\n else:\n self.extra_wgts = None\n else:\n self.extra_wgts = None\n\n def generate_dirnames(\n self,\n *,\n prefix: str = \"\",\n shape: torch.Size | tuple[int, ...] = (1024, 1024, 16, 16),\n default_dtype: torch.dtype = torch.float16,\n **kwargs,\n ) -> list[str]:\n \"\"\"Get the directory names of the quantization configuration.\n\n Args:\n prefix (`str`, *optional*, defaults to `\"\"`):\n The prefix for the directory names.\n shape (`torch.Size` or `tuple[int, ...]`, *optional*, defaults to `(1024, 1024, 16, 16)`):\n The shape of the tensor to be quantized.\n default_dtype (`torch.dtype`, *optional*, defaults to `torch.float16`):\n The dtype of the tensor to be quantized.\n\n Returns:\n `list[str]`:\n The directory names of the quantization configuration.\n - The number of effective bits.\n - The name of the quantization data type.\n - The name of the group shapes.\n - The name of the modules to skip.\n \"\"\"\n wgts_names = self.wgts.generate_dirnames(prefix=\"w\", shape=shape, default_dtype=default_dtype)\n ipts_names = self.ipts.generate_dirnames(prefix=\"x\", shape=shape, default_dtype=default_dtype)\n opts_names = self.opts.generate_dirnames(prefix=\"y\", shape=shape, default_dtype=default_dtype)\n names = [\n f\"{wgts_name}-{ipts_name}-{opts_name}\"\n for wgts_name, ipts_name, opts_name in zip(wgts_names, ipts_names, opts_names, strict=True)\n ]\n if self.extra_wgts is not None:\n extra_wgts_names = self.extra_wgts.generate_dirnames(prefix=\"w\", shape=shape, default_dtype=default_dtype)\n names = [f\"{name}-{extra_wgts_name}\" for name, extra_wgts_name in zip(names, extra_wgts_names, strict=True)]\n if prefix:\n names = [f\"{prefix}.[{name}]\" for name in names]\n return names\n\n def generate_calib_dirname(self) -> str:\n \"\"\"Generate the name for quantization calibration.\n\n Returns:\n `str`:\n The name.\n \"\"\"\n name = \"\"\n if self.enabled_wgts:\n calib_name = self.wgts.generate_calib_dirname()\n if calib_name:\n name += f\"-w.{calib_name}\"\n if self.enabled_ipts:\n calib_name = self.ipts.generate_calib_dirname()\n if calib_name:\n name += f\"-x.{calib_name}\"\n if self.enabled_opts:\n calib_name = self.opts.generate_calib_dirname()\n if calib_name:\n name += f\"-y.{calib_name}\"\n return name[1:] if name else name\n"
},
{
"path": "deepcompressor/app/diffusion/quant/quantizer/quantizer.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Tensor Quantizer module.\"\"\"\n\nimport typing as tp\nfrom dataclasses import dataclass, field\n\nimport torch\nimport torch.nn as nn\n\nfrom deepcompressor.calib.config import SkipBasedQuantLowRankCalibConfig\nfrom deepcompressor.calib.lowrank import LowRankBranch, QuantLowRankCalibrator\nfrom deepcompressor.calib.range import calibrate_dynamic_range\nfrom deepcompressor.data.cache import TensorsCache\nfrom deepcompressor.data.common import TensorType\nfrom deepcompressor.data.range import DynamicRange\nfrom deepcompressor.quantizer.processor import Quantizer\n\nfrom .config import (\n DiffusionActivationQuantizerConfig,\n DiffusionGPTQConfig,\n DiffusionQuantizerConfig,\n DiffusionWeightQuantizerConfig,\n)\n\n__all__ = [\"DiffusionQuantizer\", \"DiffusionWeightQuantizer\", \"DiffusionActivationQuantizer\"]\n\n\n@dataclass\nclass DiffusionQuantizer(Quantizer):\n \"\"\"Denoising model quantizer class.\n\n Args:\n config (`DiffusionQuantizerConfig` or `None`):\n The quantizer configuration.\n key (`str`, *optional*, defaults to `\"\"`):\n The key of the quantizer.\n tensor_type (`TensorType`, *optional*, defaults to `TensorType.Weights`):\n The type of the tensor to quantize.\n channels_dim (`int` or `None`, *optional*, defaults to `None`):\n The dimension of channels.\n scale (`torch.Tensor` or `Sequence[torch.Tensor]` or `None`, *optional*, defaults to `None`):\n The scale tensor.\n zero (`torch.Tensor` or `None`, *optional*, defaults to `None`):\n The zero point tensor.\n dynamic_range (`DynamicRange` or `Sequence[DynamicRange]` or `None`, *optional*, defaults to `None`):\n The dynamic range.\n range_bound (`RangeBound` or `None`, *optional*, defaults to `None`):\n The dynamic range bound.\n quant_range (`QuantRange` or `None`, *optional*, defaults to `None`):\n The quantization range.\n default_dtype (`torch.dtype` or `None`, *optional*, defaults to `None`):\n The default scale dtype\n develop_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The quantization development dtype.\n\n kernel (`DiffusionGPTQConfig` or `None`, *optional*, defaults to `MISSING`):\n The GPTQ kernel configuration.\n If not provided (i.e., `MISSING`), the GPTQ configuration from the `config` will be used.\n low_rank (`QuantLowRankConfig` or `None`, *optional*, defaults to `MISSING`):\n The quantization low-rank branch configuration.\n If not provided (i.e., `MISSING`), the low-rank branch configuration from the `config` will be used.\n input_packager (`BaseInputPackager` or `None`, *optional*, defaults to `None`):\n The input packager, used for unpacking and repacking the input tensor(s).\n output_packager (`BaseOutputPackager` or `None`, *optional*, defaults to `None`):\n The output packager, used for unpacking and repacking the output tensor(s).\n \"\"\"\n\n config: DiffusionQuantizerConfig\n kernel: DiffusionGPTQConfig | None = field(init=False)\n low_rank: SkipBasedQuantLowRankCalibConfig | None = field(init=False)\n tensor_type: TensorType = TensorType.Weights\n\n def __post_init__(self) -> None:\n self.kernel = self.config.kernel_gptq\n self.low_rank = self.config.low_rank\n\n def calibrate_dynamic_range(\n self,\n modules: tp.Sequence[nn.Module],\n activations: TensorsCache,\n weights: tp.Sequence[nn.Parameter] = None,\n eval_inputs: TensorsCache | None = None,\n eval_module: nn.Module | None = None,\n eval_kwargs: dict[str, tp.Any] | None = None,\n orig_weights: tp.Sequence[tuple[nn.Parameter, torch.Tensor]] | None = None,\n orig_activations: TensorsCache | None = None,\n orig_eval_inputs: TensorsCache | None = None,\n ) -> tp.Sequence[DynamicRange] | None:\n \"\"\"Calibrate the dynamic range.\n\n Args:\n modules (`Sequence[nn.Module]`):\n The modules to calibrate.\n activations (`TensorsCache`):\n The inputs cache if the tensor type is not outputs, or the outputs cache if the tensor type is outputs.\n weights (`Sequence[nn.Parameter]` or `None`, *optional*, defaults to `None`):\n The weights to calibrate.\n If not provided, the weights of the modules will be used.\n eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The cache of the inputs for evaluation.\n If not provided, the `activations` cache will be used.\n eval_module (`nn.Module` or `None`, *optional*, defaults to `None`):\n The module to evaluate the quantization error.\n If not provided, the module to calibrate will be used.\n eval_kwargs (`dict[str, tp.Any]` or `None`, *optional*, defaults to `None`):\n The keyword arguments for evaluation.\n orig_weights (`Sequence[tuple[nn.Parameter, torch.Tensor]]` or `None`, *optional*, defaults to `None`):\n The original weights.\n orig_activations (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The original activations.\n orig_eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The original evaluation inputs.\n\n Returns:\n `Sequence[DynamicRange]` or `None`:\n The dynamic ranges of each quantization step.\n \"\"\"\n if (\n not self.is_enabled()\n or self.config.calib_range is None\n or not self.config.calib_range.is_enabled_for(self.key)\n ):\n self.dynamic_range = None\n else:\n self.dynamic_range = calibrate_dynamic_range(\n tensor_type=self.tensor_type,\n config=self.config.calib_range,\n static=self.config.static,\n quantizer=self,\n modules=modules,\n activations=activations,\n weights=weights,\n eval_inputs=eval_inputs,\n eval_module=eval_module,\n eval_kwargs=eval_kwargs,\n orig_weights=orig_weights,\n orig_activations=orig_activations,\n orig_eval_inputs=orig_eval_inputs,\n )\n return self.dynamic_range\n\n\n@dataclass\nclass DiffusionWeightQuantizer(DiffusionQuantizer):\n \"\"\"Diffusion model weight quantizer class.\n\n Args:\n\n Args:\n config (`DiffusionWeightQuantizerConfig` or `None`):\n The quantizer configuration.\n key (`str`, *optional*, defaults to `\"\"`):\n The key of the quantizer.\n scale (`torch.Tensor` or `Sequence[torch.Tensor]` or `None`, *optional*, defaults to `None`):\n The scale tensor.\n zero (`torch.Tensor` or `None`, *optional*, defaults to `None`):\n The zero point tensor.\n dynamic_range (`DynamicRange` or `Sequence[DynamicRange]` or `None`, *optional*, defaults to `None`):\n The dynamic range.\n range_bound (`RangeBound` or `None`, *optional*, defaults to `None`):\n The dynamic range bound.\n quant_range (`QuantRange` or `None`, *optional*, defaults to `None`):\n The quantization range.\n default_dtype (`torch.dtype` or `None`, *optional*, defaults to `None`):\n The default scale dtype\n develop_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The quantization development dtype.\n\n kernel (`DiffusionGPTQConfig` or `None`, *optional*, defaults to `MISSING`):\n The GPTQ kernel configuration.\n If not provided (i.e., `MISSING`), the GPTQ configuration from the `config` will be used.\n low_rank (`QuantLowRankConfig` or `None`, *optional*, defaults to `MISSING`):\n The quantization low-rank branch configuration.\n If not provided (i.e., `MISSING`), the low-rank branch configuration from the `config` will be used.\n input_packager (`BaseInputPackager` or `None`, *optional*, defaults to `None`):\n The input packager, used for unpacking and repacking the input tensor(s).\n output_packager (`BaseOutputPackager` or `None`, *optional*, defaults to `None`):\n The output packager, used for unpacking and repacking the output tensor(s).\n \"\"\"\n\n config: DiffusionWeightQuantizerConfig\n channels_dim: None = field(init=False, default=None)\n tensor_type: TensorType = field(init=False, default=TensorType.Weights)\n\n def calibrate_dynamic_range(\n self,\n module: nn.Module,\n inputs: TensorsCache,\n weight: nn.Parameter | None = None,\n eval_inputs: TensorsCache | None = None,\n eval_module: nn.Module | None = None,\n eval_kwargs: dict[str, tp.Any] | None = None,\n orig_inputs: TensorsCache | None = None,\n orig_eval_inputs: TensorsCache | None = None,\n ) -> tp.Sequence[DynamicRange] | None:\n \"\"\"Calibrate the dynamic range.\n\n Args:\n module (`nn.Module`):\n The module to calibrate.\n inputs (`TensorsCache`):\n The inputs cache.\n weight (`nn.Parameter` or `None`, *optional*, defaults to `None`):\n The weight parameter to calibrate.\n If not provided, the weight of the `module` will be used.\n eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The cache of the inputs for evaluation.\n If not provided, the `activations` cache will be used.\n eval_module (`nn.Module` or `None`, *optional*, defaults to `None`):\n The module to evaluate the quantization error.\n If not provided, the module to calibrate will be used.\n eval_kwargs (`dict[str, tp.Any]` or `None`, *optional*, defaults to `None`):\n The keyword arguments for evaluation.\n orig_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The original inputs.\n orig_eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The original evaluation inputs.\n\n Returns:\n `Sequence[DynamicRange]` or `None`:\n The dynamic ranges of each quantization step.\n \"\"\"\n return super().calibrate_dynamic_range(\n modules=[module],\n weights=[weight] if weight is not None else [module.weight],\n activations=inputs,\n eval_inputs=eval_inputs,\n eval_module=eval_module,\n eval_kwargs=eval_kwargs,\n orig_activations=orig_inputs,\n orig_eval_inputs=orig_eval_inputs,\n )\n\n def calibrate_low_rank(\n self,\n input_quantizer: \"DiffusionActivationQuantizer\",\n modules: tp.Sequence[nn.Module],\n inputs: TensorsCache,\n weights: tp.Sequence[nn.Parameter] = None,\n eval_inputs: TensorsCache | None = None,\n eval_module: nn.Module | None = None,\n eval_kwargs: dict[str, tp.Any] | None = None,\n orig_inputs: TensorsCache | None = None,\n orig_eval_inputs: TensorsCache | None = None,\n ) -> LowRankBranch:\n \"\"\"Calibrate the quantization low-rank branch.\"\"\"\n if weights is None:\n weights = [module.weight for module in modules]\n return QuantLowRankCalibrator(\n config=self.low_rank,\n w_quantizer=self,\n x_quantizer=input_quantizer,\n develop_dtype=self.develop_dtype,\n ).calibrate(\n x_wgts=weights,\n x_acts=inputs,\n x_mods=modules,\n eval_inputs=eval_inputs,\n eval_module=eval_module,\n eval_kwargs=eval_kwargs,\n orig_x_acts=orig_inputs,\n orig_eval_inputs=orig_eval_inputs,\n )\n\n\n@dataclass\nclass DiffusionActivationQuantizer(DiffusionQuantizer):\n \"\"\"Diffusion model activation quantizer class.\n\n Args:\n config (`DiffusionActivationQuantizerConfig` or `None`):\n The quantizer configuration.\n key (`str`, *optional*, defaults to `\"\"`):\n The key of the quantizer.\n tensor_type (`TensorType`, *optional*, defaults to `TensorType.Inputs`):\n The type of the tensor to quantize.\n channels_dim (`int` or `None`, *optional*, defaults to `None`):\n The dimension of channels.\n scale (`torch.Tensor` or `Sequence[torch.Tensor]` or `None`, *optional*, defaults to `None`):\n The scale tensor.\n zero (`torch.Tensor` or `None`, *optional*, defaults to `None`):\n The zero point tensor.\n dynamic_range (`DynamicRange` or `Sequence[DynamicRange]` or `None`, *optional*, defaults to `None`):\n The dynamic range.\n range_bound (`RangeBound` or `None`, *optional*, defaults to `None`):\n The dynamic range bound.\n quant_range (`QuantRange` or `None`, *optional*, defaults to `None`):\n The quantization range.\n default_dtype (`torch.dtype` or `None`, *optional*, defaults to `None`):\n The default scale dtype\n develop_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The quantization development dtype.\n\n input_packager (`BaseInputPackager` or `None`, *optional*, defaults to `None`):\n The input packager, used for unpacking and repacking the input tensor(s).\n output_packager (`BaseOutputPackager` or `None`, *optional*, defaults to `None`):\n The output packager, used for unpacking and repacking the output tensor(s).\n develop_dtype (torch.dtype, optional): The develop dtype. Defaults to ``torch.float32``.\n \"\"\"\n\n config: DiffusionActivationQuantizerConfig\n tensor_type: TensorType = TensorType.Inputs\n\n def __post_init__(self) -> None:\n super().__post_init__()\n assert self.tensor_type != TensorType.Weights, \"The tensor type cannot be weights.\"\n assert isinstance(self.channels_dim, int), \"The channels dimension must be provided.\"\n"
},
{
"path": "deepcompressor/app/diffusion/quant/rotate.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Large Language Model Rotation module.\"\"\"\n\nimport gc\n\nimport torch\n\nfrom deepcompressor.calib.rotate import (\n get_rotation_matrix,\n hadamard_in_channels,\n rotate_in_channels,\n rotate_out_channels,\n)\nfrom deepcompressor.utils import tools\n\nfrom ..nn.struct import DiffusionModelStruct\nfrom .config import DiffusionQuantConfig\n\n__all__ = [\"rotate_diffusion\"]\n\n\n@torch.inference_mode()\ndef rotate_diffusion( # noqa: C901\n model: DiffusionModelStruct, /, config: DiffusionQuantConfig\n):\n \"\"\"Rotate the weights of the diffusion model.\n\n Args:\n model (`PreTrainedModel` or `LlmStruct`):\n Model to be rotated.\n config (`QuantRotationConfig`):\n Rotation configuration.\n \"\"\"\n if not isinstance(model, DiffusionModelStruct):\n model = DiffusionModelStruct.construct(model)\n assert isinstance(model, DiffusionModelStruct)\n devices: dict[str, torch.device] = {}\n dtypes: dict[str, torch.dtype] = {}\n linears: dict[str, torch.nn.Linear] = {}\n size: float = 0\n for n, m in model.module.named_modules():\n if isinstance(m, torch.nn.Linear):\n devices[n] = m.weight.device\n dtypes[n] = m.weight.dtype\n linears[n] = m\n size += m.weight.numel() / 1e9\n for linear in linears.values():\n linear.to(dtype=torch.float32, device=\"cpu\" if size > 30 else None)\n\n logger = tools.logging.getLogger(f\"{__name__}.Rotate\")\n head_rotation = None\n for transformer_block in model.iter_transformer_block_structs():\n logger.debug(f\"- Rotating {transformer_block.name}\")\n tools.logging.Formatter.indent_inc()\n for attn in transformer_block.iter_attention_structs():\n if attn.qkv_proj_key in config.rotation.transforms:\n if attn.qkv_proj_key not in config.wgts.skips or attn.qkv_proj_key not in config.ipts.skips:\n logger.debug(f\"- Hadamard transform on {attn.name}.qkv_proj (in)\")\n hadamard_in_channels(\n attn.qkv_proj, dtype=dtypes[attn.q_proj_name], device=devices[attn.q_proj_name]\n )\n if not attn.is_self_attn() and attn.add_qkv_proj_key in config.rotation.transforms:\n if attn.add_qkv_proj_key not in config.wgts.skips or attn.add_qkv_proj_key not in config.ipts.skips:\n logger.debug(f\"- Hadamard transform on {attn.name}.add_qkv_proj (in)\")\n hadamard_in_channels(\n attn.add_qkv_proj, dtype=dtypes[attn.add_k_proj_name], device=devices[attn.add_k_proj_name]\n )\n if attn.out_proj_key in config.rotation.transforms or attn.add_out_proj_key in config.rotation.transforms:\n if (\n attn.out_proj_key not in config.wgts.skips\n or attn.out_proj_key not in config.ipts.skips\n or attn.add_out_proj_key not in config.wgts.skips\n or attn.add_out_proj_key not in config.ipts.skips\n ):\n if head_rotation is None:\n head_rotation = get_rotation_matrix(\n attn.config.num_head_channels, random=config.rotation.random\n )\n if attn.v_proj is not None:\n logger.debug(f\"- Rotating {attn.v_proj_name} (out)\")\n rotate_out_channels(attn.v_proj.weight, rotation=head_rotation, bias=attn.v_proj.bias)\n if attn.add_v_proj is not None:\n logger.debug(f\"- Rotating {attn.add_v_proj_name} (out)\")\n rotate_out_channels(attn.add_v_proj.weight, rotation=head_rotation, bias=attn.add_v_proj.bias)\n if attn.o_proj is not None:\n logger.debug(f\"- Rotating {attn.o_proj_name} (in)\")\n rotate_in_channels(attn.o_proj.weight, rotation=head_rotation)\n if attn.add_o_proj is not None:\n logger.debug(f\"- Rotating {attn.add_o_proj_name} (in)\")\n rotate_in_channels(attn.add_o_proj.weight, rotation=head_rotation)\n gc.collect()\n torch.cuda.empty_cache()\n ffn, add_ffn = transformer_block.ffn_struct, transformer_block.add_ffn_struct\n if ffn.up_proj_key in config.rotation.transforms:\n if ffn.up_proj_key not in config.wgts.skips or ffn.up_proj_key not in config.ipts.skips:\n logger.debug(f\"- Hadamard transform on {ffn.up_proj_name} (in)\")\n hadamard_in_channels(ffn.up_projs, dtype=dtypes[ffn.up_proj_name], device=devices[ffn.up_proj_name])\n if add_ffn is not None and add_ffn.up_proj_key in config.rotation.transforms:\n if add_ffn.up_proj_key not in config.wgts.skips or add_ffn.up_proj_key not in config.ipts.skips:\n logger.debug(f\"- Hadamard transform on {add_ffn.up_proj_name} (in)\")\n hadamard_in_channels(\n add_ffn.up_projs, dtype=dtypes[add_ffn.up_proj_name], device=devices[add_ffn.up_proj_name]\n )\n if ffn.down_proj_key in config.rotation.transforms:\n if ffn.down_proj_key not in config.wgts.skips or ffn.down_proj_key not in config.ipts.skips:\n logger.debug(f\"- Hadamard transform on {ffn.down_proj_name} (in)\")\n hadamard_in_channels(\n ffn.down_projs, dtype=dtypes[ffn.down_proj_name], device=devices[ffn.down_proj_name]\n )\n if add_ffn is not None and add_ffn.down_proj_key in config.rotation.transforms:\n if add_ffn.down_proj_key not in config.wgts.skips or add_ffn.down_proj_key not in config.ipts.skips:\n logger.debug(f\"- Hadamard transform on {add_ffn.down_proj_name} (in)\")\n hadamard_in_channels(\n add_ffn.down_projs, dtype=dtypes[add_ffn.down_proj_name], device=devices[add_ffn.down_proj_name]\n )\n gc.collect()\n torch.cuda.empty_cache()\n tools.logging.Formatter.indent_dec()\n\n for n, m in linears.items():\n m.to(device=devices[n], dtype=dtypes[n])\n"
},
{
"path": "deepcompressor/app/diffusion/quant/smooth.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Diffusion smooth quantization module.\"\"\"\n\nimport typing as tp\n\nimport torch\nimport torch.nn as nn\nfrom tqdm import tqdm\n\nfrom deepcompressor.calib.smooth import ActivationSmoother, smooth_linear_modules\nfrom deepcompressor.data.cache import IOTensorsCache\nfrom deepcompressor.quantizer import Quantizer\nfrom deepcompressor.utils import tools\nfrom deepcompressor.utils.hooks import KeyedInputPackager\n\nfrom ..nn.struct import (\n DiffusionAttentionStruct,\n DiffusionBlockStruct,\n DiffusionFeedForwardStruct,\n DiffusionModelStruct,\n DiffusionTransformerBlockStruct,\n)\nfrom .config import DiffusionQuantConfig\nfrom .utils import get_needs_inputs_fn, wrap_joint_attn\n\n__all__ = [\"smooth_diffusion\"]\n\n\n@torch.inference_mode()\ndef smooth_diffusion_attention(\n attn: DiffusionAttentionStruct,\n config: DiffusionQuantConfig,\n smooth_cache: dict[str, torch.Tensor],\n block_cache: dict[str, IOTensorsCache] | None = None,\n block_kwargs: dict[str, tp.Any] | None = None,\n) -> dict[str, torch.Tensor]:\n logger = tools.logging.getLogger(f\"{__name__}.SmoothQuant\")\n # attention qk\n if config.smooth.enabled_attn:\n logger.debug(\"- %s.k\", attn.name)\n raise NotImplementedError(\"Not implemented yet\")\n return smooth_cache\n\n\n@torch.inference_mode()\ndef smooth_diffusion_qkv_proj(\n attn: DiffusionAttentionStruct,\n config: DiffusionQuantConfig,\n smooth_cache: dict[str, torch.Tensor],\n block_cache: dict[str, IOTensorsCache] | None = None,\n block_kwargs: dict[str, tp.Any] | None = None,\n) -> dict[str, torch.Tensor]:\n logger = tools.logging.getLogger(f\"{__name__}.SmoothQuant\")\n # region qkv projection\n module_key = attn.qkv_proj_key\n needs_quant = config.enabled_wgts and config.wgts.is_enabled_for(module_key)\n needs_quant = needs_quant or (config.enabled_ipts and config.ipts.is_enabled_for(module_key))\n if needs_quant and config.smooth.enabled_proj and config.smooth.proj.is_enabled_for(module_key):\n logger.debug(\"- %s.qkv_proj\", attn.name)\n prevs = None\n if config.smooth.proj.fuse_when_possible and attn.parent.norm_type.startswith(\"layer_norm\"):\n if not hasattr(attn.parent.module, \"pos_embed\") or attn.parent.module.pos_embed is None:\n prevs = attn.parent.pre_attn_norms[attn.idx]\n assert isinstance(prevs, nn.LayerNorm)\n cache_key = attn.q_proj_name\n config_wgts = config.wgts\n if config.enabled_extra_wgts and config.extra_wgts.is_enabled_for(module_key):\n config_wgts = config.extra_wgts\n smooth_cache[cache_key] = smooth_linear_modules(\n prevs,\n attn.qkv_proj,\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config_wgts, key=module_key, low_rank=config.wgts.low_rank),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=module_key),\n inputs=block_cache[attn.q_proj_name].inputs if block_cache else None,\n eval_inputs=block_cache[attn.name].inputs if block_cache else None,\n eval_module=attn,\n eval_kwargs=attn.filter_kwargs(block_kwargs),\n develop_dtype=config.develop_dtype,\n )\n if prevs is None:\n # we need to register forward pre hook to smooth inputs\n if attn.module.group_norm is None and attn.module.spatial_norm is None:\n ActivationSmoother(\n smooth_cache[cache_key],\n channels_dim=-1,\n input_packager=KeyedInputPackager(attn.module, [0]),\n ).as_hook().register(attn.module)\n else:\n ActivationSmoother(smooth_cache[cache_key], channels_dim=-1).as_hook().register(attn.qkv_proj)\n for m in attn.qkv_proj:\n m.in_smooth_cache_key = cache_key\n # endregion\n if attn.is_self_attn():\n return smooth_cache\n # region additional qkv projection\n module_key = attn.add_qkv_proj_key\n needs_quant = config.enabled_wgts and config.wgts.is_enabled_for(module_key)\n needs_quant = needs_quant or (config.enabled_ipts and config.ipts.is_enabled_for(module_key))\n needs_quant = needs_quant and attn.add_k_proj is not None\n if needs_quant and config.smooth.enabled_proj and config.smooth.proj.is_enabled_for(module_key):\n logger.debug(\"- %s add_qkv_proj\", attn.name)\n prevs = None\n pre_attn_add_norm = attn.parent.pre_attn_add_norms[attn.idx]\n if isinstance(pre_attn_add_norm, nn.LayerNorm) and config.smooth.proj.fuse_when_possible:\n prevs = pre_attn_add_norm\n cache_key = attn.add_k_proj_name\n config_wgts = config.wgts\n if config.enabled_extra_wgts and config.extra_wgts.is_enabled_for(module_key):\n config_wgts = config.extra_wgts\n smooth_cache[cache_key] = smooth_linear_modules(\n prevs,\n attn.add_qkv_proj,\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config_wgts, key=module_key, low_rank=config.wgts.low_rank),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=module_key),\n inputs=block_cache[attn.add_k_proj_name].inputs if block_cache else None,\n eval_inputs=block_cache[attn.name].inputs if block_cache else None,\n eval_module=wrap_joint_attn(attn, indexes=1) if attn.is_joint_attn() else attn,\n eval_kwargs=attn.filter_kwargs(block_kwargs),\n develop_dtype=config.develop_dtype,\n )\n if prevs is None:\n # we need to register forward pre hook to smooth inputs\n ActivationSmoother(smooth_cache[cache_key], channels_dim=-1).as_hook().register(attn.add_qkv_proj)\n for m in attn.add_qkv_proj:\n m.in_smooth_cache_key = cache_key\n # endregion\n return smooth_cache\n\n\n@torch.inference_mode()\ndef smooth_diffusion_out_proj( # noqa: C901\n attn: DiffusionAttentionStruct,\n config: DiffusionQuantConfig,\n smooth_cache: dict[str, torch.Tensor],\n block_cache: dict[str, IOTensorsCache] | None = None,\n block_kwargs: dict[str, tp.Any] | None = None,\n) -> dict[str, torch.Tensor]:\n logger = tools.logging.getLogger(f\"{__name__}.SmoothQuant\")\n module_keys = []\n for module_key in (attn.out_proj_key, attn.add_out_proj_key) if attn.is_joint_attn() else (attn.out_proj_key,):\n needs_quant = config.enabled_wgts and config.wgts.is_enabled_for(module_key)\n needs_quant = needs_quant or (config.enabled_ipts and config.ipts.is_enabled_for(module_key))\n if needs_quant and config.smooth.enabled_proj and config.smooth.proj.is_enabled_for(module_key):\n module_keys.append(module_key)\n if not module_keys:\n return smooth_cache\n exclusive = False\n if config.enabled_wgts and config.wgts.enabled_low_rank:\n exclusive = config.wgts.low_rank.exclusive\n config.wgts.low_rank.exclusive = True\n fuse_smooth = not attn.config.linear_attn and config.smooth.proj.fuse_when_possible\n prevs = [attn.v_proj, attn.add_v_proj] if fuse_smooth else None\n if len(module_keys) == 1 and module_keys[0] == attn.out_proj_key:\n logger.debug(\"- %s.out_proj\", attn.name)\n module_key = attn.out_proj_key\n cache_key = attn.o_proj_name\n config_wgts = config.wgts\n if config.enabled_extra_wgts and config.extra_wgts.is_enabled_for(module_key):\n config_wgts = config.extra_wgts\n smooth_cache[cache_key] = smooth_linear_modules(\n prevs,\n attn.o_proj,\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config_wgts, key=module_key, low_rank=config.wgts.low_rank),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=module_key),\n inputs=block_cache[attn.o_proj_name].inputs if block_cache else None,\n eval_inputs=block_cache[attn.o_proj_name].inputs if block_cache else None,\n eval_module=attn.o_proj,\n extra_modules=[attn.add_o_proj] if attn.is_joint_attn() else None,\n develop_dtype=config.develop_dtype,\n )\n elif len(module_keys) == 1 and module_keys[0] == attn.add_out_proj_key:\n assert attn.is_joint_attn()\n logger.debug(\"- %s.add_out_proj\", attn.name)\n module_key = attn.add_out_proj_key\n cache_key = attn.add_o_proj_name\n config_wgts = config.wgts\n if config.enabled_extra_wgts and config.extra_wgts.is_enabled_for(module_key):\n config_wgts = config.extra_wgts\n smooth_cache[cache_key] = smooth_linear_modules(\n prevs,\n attn.add_o_proj,\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config_wgts, key=module_key, low_rank=config.wgts.low_rank),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=module_key),\n inputs=block_cache[attn.add_o_proj_name].inputs if block_cache else None,\n eval_inputs=block_cache[attn.add_o_proj_name].inputs if block_cache else None,\n eval_module=attn.add_o_proj,\n extra_modules=[attn.o_proj],\n develop_dtype=config.develop_dtype,\n )\n else:\n assert attn.is_joint_attn()\n logger.debug(\"- %s.out_proj + %s.add_out_proj\", attn.name, attn.name)\n module_key = attn.out_proj_key\n cache_key = attn.o_proj_name\n config_wgts = config.wgts\n if config.enabled_extra_wgts and config.extra_wgts.is_enabled_for(module_key):\n config_wgts = config.extra_wgts\n smooth_cache[cache_key] = smooth_linear_modules(\n prevs,\n [attn.o_proj, attn.add_o_proj],\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config_wgts, key=module_key, low_rank=config.wgts.low_rank),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=module_key),\n inputs=block_cache[attn.o_proj_name].inputs if block_cache else None,\n eval_inputs=block_cache[attn.name].inputs if block_cache else None,\n eval_module=wrap_joint_attn(attn, indexes=(0, 1)),\n eval_kwargs=attn.filter_kwargs(block_kwargs),\n develop_dtype=config.develop_dtype,\n )\n if config.enabled_wgts and config.wgts.enabled_low_rank:\n config.wgts.low_rank.exclusive = exclusive\n if fuse_smooth:\n for prev in prevs:\n if prev is not None:\n prev.out_smooth_cache_key = cache_key\n else:\n for o_proj in [attn.o_proj, attn.add_o_proj]:\n if o_proj is not None:\n ActivationSmoother(smooth_cache[cache_key], channels_dim=-1).as_hook().register(o_proj)\n attn.o_proj.in_smooth_cache_key = cache_key\n if attn.add_o_proj is not None:\n attn.add_o_proj.in_smooth_cache_key = cache_key\n return smooth_cache\n\n\n@torch.inference_mode()\ndef smooth_diffusion_up_proj(\n pre_ffn_norm: nn.Module,\n ffn: DiffusionFeedForwardStruct,\n config: DiffusionQuantConfig,\n smooth_cache: dict[str, torch.Tensor],\n block_cache: dict[str, IOTensorsCache] | None = None,\n) -> dict[str, torch.Tensor]:\n assert len(ffn.up_projs) == 1\n logger = tools.logging.getLogger(f\"{__name__}.SmoothQuant\")\n # ffn up projection\n module_key = ffn.up_proj_key\n needs_quant = config.enabled_wgts and config.wgts.is_enabled_for(module_key)\n needs_quant = needs_quant or (config.enabled_ipts and config.ipts.is_enabled_for(module_key))\n if needs_quant and config.smooth.enabled_proj and config.smooth.proj.is_enabled_for(module_key):\n logger.debug(\"- %s.up_proj\", ffn.name)\n prevs = None\n if config.smooth.proj.fuse_when_possible and isinstance(pre_ffn_norm, nn.LayerNorm):\n if ffn.parent.norm_type in [\"ada_norm\", \"layer_norm\"]:\n prevs = pre_ffn_norm\n cache_key = ffn.up_proj_name\n channels_dim = -1 if isinstance(ffn.down_proj, nn.Linear) else 1\n config_wgts = config.wgts\n if config.enabled_extra_wgts and config.extra_wgts.is_enabled_for(module_key):\n config_wgts = config.extra_wgts\n smooth_cache[cache_key] = smooth_linear_modules(\n prevs,\n ffn.up_projs,\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config_wgts, key=module_key, low_rank=config.wgts.low_rank),\n input_quantizer=Quantizer(config.ipts, channels_dim=channels_dim, key=module_key),\n inputs=block_cache[ffn.up_proj_name].inputs if block_cache else None,\n eval_inputs=block_cache[ffn.up_proj_name].inputs if block_cache else None,\n eval_module=ffn.up_proj,\n develop_dtype=config.develop_dtype,\n )\n if prevs is None:\n ActivationSmoother(smooth_cache[cache_key], channels_dim=channels_dim).as_hook().register(ffn.up_proj)\n for proj in ffn.up_projs:\n proj.in_smooth_cache_key = cache_key\n return smooth_cache\n\n\n@torch.inference_mode()\ndef smooth_diffusion_down_proj(\n ffn: DiffusionFeedForwardStruct,\n config: DiffusionQuantConfig,\n smooth_cache: dict[str, torch.Tensor],\n block_cache: dict[str, IOTensorsCache] | None = None,\n) -> dict[str, torch.Tensor]:\n logger = tools.logging.getLogger(f\"{__name__}.SmoothQuant\")\n # ffn down projection\n module_key = ffn.down_proj_key.upper()\n needs_quant = config.enabled_wgts and config.wgts.is_enabled_for(module_key)\n needs_quant = needs_quant or (config.enabled_ipts and config.ipts.is_enabled_for(module_key))\n if needs_quant and config.smooth.enabled_proj and config.smooth.proj.is_enabled_for(module_key):\n logger.debug(\"- %s.down_proj\", ffn.name)\n cache_key = ffn.down_proj_name\n config_ipts = config.unsigned_ipts if getattr(ffn.down_proj, \"unsigned\", False) else config.ipts\n config_wgts = config.wgts\n if config.enabled_extra_wgts and config.extra_wgts.is_enabled_for(module_key):\n config_wgts = config.extra_wgts\n channels_dim = -1 if isinstance(ffn.down_proj, nn.Linear) else 1\n smooth_cache[cache_key] = smooth_linear_modules(\n None,\n ffn.down_proj,\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config_wgts, key=module_key, low_rank=config.wgts.low_rank),\n input_quantizer=Quantizer(config_ipts, channels_dim=channels_dim, key=module_key),\n inputs=block_cache[ffn.down_proj_name].inputs if block_cache else None,\n eval_inputs=block_cache[ffn.down_proj_name].inputs if block_cache else None,\n eval_module=ffn.down_proj,\n develop_dtype=config.develop_dtype,\n )\n ffn.down_proj.in_smooth_cache_key = cache_key\n ActivationSmoother(smooth_cache[cache_key], channels_dim=channels_dim).as_hook().register(ffn.down_proj)\n return smooth_cache\n\n\n@torch.inference_mode()\ndef smooth_diffusion_parallel_qkv_up_proj(\n block: DiffusionTransformerBlockStruct,\n config: DiffusionQuantConfig,\n smooth_cache: dict[str, torch.Tensor],\n block_cache: dict[str, IOTensorsCache] | None = None,\n block_kwargs: dict[str, tp.Any] | None = None,\n) -> dict[str, torch.Tensor]:\n assert block.parallel\n assert len(block.ffn_struct.up_projs) == 1\n logger = tools.logging.getLogger(f\"{__name__}.SmoothQuant\")\n # region qkv proj + up proj\n attn, ffn = block.attn_structs[0], block.ffn_struct\n module_key = attn.qkv_proj_key\n needs_quant = config.enabled_wgts and config.wgts.is_enabled_for(module_key)\n needs_quant = needs_quant or (config.enabled_ipts and config.ipts.is_enabled_for(module_key))\n if needs_quant and config.smooth.enabled_proj and config.smooth.proj.is_enabled_for(module_key):\n logger.debug(\"- %s.qkv_proj + %s.up_proj\", attn.name, ffn.name)\n cache_key = attn.q_proj_name\n modules = attn.qkv_proj + ffn.up_projs\n config_wgts = config.wgts\n if config.enabled_extra_wgts and config.extra_wgts.is_enabled_for(module_key):\n config_wgts = config.extra_wgts\n smooth_cache[cache_key] = smooth_linear_modules(\n None,\n modules,\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config_wgts, key=module_key, low_rank=config.wgts.low_rank),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=module_key),\n inputs=block_cache[attn.q_proj_name].inputs if block_cache else None,\n eval_inputs=block_cache[block.name].inputs if block_cache else None,\n eval_module=block,\n eval_kwargs=block_kwargs,\n splits=[len(attn.qkv_proj)],\n develop_dtype=config.develop_dtype,\n )\n ActivationSmoother(smooth_cache[cache_key], channels_dim=-1).as_hook().register(modules)\n for m in modules:\n m.in_smooth_cache_key = cache_key\n # endregion\n # region add qkv proj + add up proj\n if attn.is_self_attn():\n if block.add_ffn_struct is not None:\n smooth_cache = smooth_diffusion_up_proj(\n pre_ffn_norm=block.pre_add_ffn_norm,\n ffn=block.add_ffn_struct,\n config=config,\n smooth_cache=smooth_cache,\n block_cache=block_cache,\n )\n return smooth_cache\n module_key = attn.add_qkv_proj_key\n needs_quant = config.enabled_wgts and config.wgts.is_enabled_for(module_key)\n needs_quant = needs_quant or (config.enabled_ipts and config.ipts.is_enabled_for(module_key))\n if needs_quant and config.smooth.enabled_proj and config.smooth.proj.is_enabled_for(module_key):\n add_ffn = block.add_ffn_struct\n cache_key = attn.add_k_proj_name\n modules = attn.add_qkv_proj\n if add_ffn is None:\n logger.debug(\"- %s.add_qkv_proj\", attn.name)\n else:\n logger.debug(\"- %s.add_qkv_proj + %s.up_proj\", attn.name, add_ffn.name)\n modules = modules + add_ffn.up_projs\n config_wgts = config.wgts\n if config.enabled_extra_wgts and config.extra_wgts.is_enabled_for(module_key):\n config_wgts = config.extra_wgts\n smooth_cache[cache_key] = smooth_linear_modules(\n None,\n modules,\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config_wgts, key=module_key, low_rank=config.wgts.low_rank),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=module_key),\n inputs=block_cache[attn.add_k_proj_name].inputs if block_cache else None,\n eval_inputs=block_cache[block.name].inputs if block_cache else None,\n eval_module=block,\n eval_kwargs=block_kwargs,\n develop_dtype=config.develop_dtype,\n )\n ActivationSmoother(smooth_cache[cache_key], channels_dim=-1).as_hook().register(modules)\n for m in modules:\n m.in_smooth_cache_key = cache_key\n # endregion\n return smooth_cache\n\n\n@torch.inference_mode()\ndef smooth_diffusion_sequential_transformer_block(\n block: DiffusionTransformerBlockStruct,\n config: DiffusionQuantConfig,\n smooth_cache: dict[str, torch.Tensor],\n block_cache: dict[str, IOTensorsCache] | None = None,\n block_kwargs: dict[str, tp.Any] | None = None,\n) -> dict[str, torch.Tensor]:\n assert not block.parallel\n for attn in block.attn_structs:\n smooth_cache = smooth_diffusion_attention(\n attn=attn, config=config, smooth_cache=smooth_cache, block_cache=block_cache, block_kwargs=block_kwargs\n )\n smooth_cache = smooth_diffusion_qkv_proj(\n attn=attn, config=config, smooth_cache=smooth_cache, block_cache=block_cache, block_kwargs=block_kwargs\n )\n smooth_cache = smooth_diffusion_out_proj(\n attn=attn, config=config, smooth_cache=smooth_cache, block_cache=block_cache, block_kwargs=block_kwargs\n )\n if block.ffn_struct is not None:\n smooth_cache = smooth_diffusion_up_proj(\n pre_ffn_norm=block.pre_ffn_norm,\n ffn=block.ffn_struct,\n config=config,\n smooth_cache=smooth_cache,\n block_cache=block_cache,\n )\n smooth_cache = smooth_diffusion_down_proj(\n ffn=block.ffn_struct, config=config, smooth_cache=smooth_cache, block_cache=block_cache\n )\n if block.add_ffn_struct is not None:\n smooth_cache = smooth_diffusion_up_proj(\n pre_ffn_norm=block.pre_add_ffn_norm,\n ffn=block.add_ffn_struct,\n config=config,\n smooth_cache=smooth_cache,\n block_cache=block_cache,\n )\n smooth_cache = smooth_diffusion_down_proj(\n ffn=block.add_ffn_struct, config=config, smooth_cache=smooth_cache, block_cache=block_cache\n )\n return smooth_cache\n\n\n@torch.inference_mode()\ndef smooth_diffusion_parallel_transformer_block(\n block: DiffusionTransformerBlockStruct,\n config: DiffusionQuantConfig,\n smooth_cache: dict[str, torch.Tensor],\n block_cache: dict[str, IOTensorsCache] | None = None,\n block_kwargs: dict[str, tp.Any] | None = None,\n) -> dict[str, torch.Tensor]:\n assert block.parallel\n assert block.ffn_struct is not None\n for attn in block.attn_structs:\n smooth_cache = smooth_diffusion_attention(\n attn=attn, config=config, smooth_cache=smooth_cache, block_cache=block_cache, block_kwargs=block_kwargs\n )\n if attn.idx == 0:\n smooth_cache = smooth_diffusion_parallel_qkv_up_proj(\n block=block,\n config=config,\n smooth_cache=smooth_cache,\n block_cache=block_cache,\n block_kwargs=block_kwargs,\n )\n else:\n smooth_cache = smooth_diffusion_qkv_proj(\n attn=attn, config=config, smooth_cache=smooth_cache, block_cache=block_cache, block_kwargs=block_kwargs\n )\n smooth_cache = smooth_diffusion_out_proj(\n attn=attn, config=config, smooth_cache=smooth_cache, block_cache=block_cache, block_kwargs=block_kwargs\n )\n smooth_cache = smooth_diffusion_down_proj(\n ffn=block.ffn_struct, config=config, smooth_cache=smooth_cache, block_cache=block_cache\n )\n if block.add_ffn_struct is not None:\n smooth_cache = smooth_diffusion_down_proj(\n ffn=block.add_ffn_struct, config=config, smooth_cache=smooth_cache, block_cache=block_cache\n )\n return smooth_cache\n\n\n@torch.inference_mode()\ndef smooth_diffusion_module(\n module_key: str,\n module_name: str,\n module: nn.Linear | nn.Conv2d,\n config: DiffusionQuantConfig,\n smooth_cache: dict[str, torch.Tensor],\n layer_cache: dict[str, IOTensorsCache] | None = None,\n) -> dict[str, torch.Tensor]:\n assert isinstance(module, (nn.Linear, nn.Conv2d))\n logger = tools.logging.getLogger(f\"{__name__}.SmoothQuant\")\n needs_quant = config.enabled_wgts and config.wgts.is_enabled_for(module_key)\n needs_quant = needs_quant or (config.enabled_ipts and config.ipts.is_enabled_for(module_key))\n if needs_quant and config.smooth.enabled_proj and config.smooth.proj.is_enabled_for(module_key):\n logger.debug(\"- Smoothing Module %s\", module_name)\n tools.logging.Formatter.indent_inc()\n logger.debug(\"- %s\", module_name)\n cache_key = module_name\n channels_dim = -1 if isinstance(module, nn.Linear) else 1\n config_wgts = config.wgts\n if config.enabled_extra_wgts and config.extra_wgts.is_enabled_for(module_key):\n config_wgts = config.extra_wgts\n smooth_cache[cache_key] = smooth_linear_modules(\n None,\n module,\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config_wgts, key=module_key),\n input_quantizer=Quantizer(config.ipts, channels_dim=channels_dim, key=module_key),\n inputs=layer_cache[module_name].inputs if layer_cache else None,\n eval_inputs=layer_cache[module_name].inputs if layer_cache else None,\n eval_module=module,\n develop_dtype=config.develop_dtype,\n )\n ActivationSmoother(smooth_cache[cache_key], channels_dim=channels_dim).as_hook().register(module)\n module.in_smooth_cache_key = cache_key\n tools.logging.Formatter.indent_dec()\n else:\n logger.debug(\"- Skipping Module %s\", module_name)\n return smooth_cache\n\n\n@torch.inference_mode()\ndef smooth_diffusion_layer(\n layer: DiffusionBlockStruct,\n config: DiffusionQuantConfig,\n smooth_cache: dict[str, torch.Tensor],\n layer_cache: dict[str, IOTensorsCache] | None = None,\n layer_kwargs: dict[str, tp.Any] | None = None,\n) -> None:\n \"\"\"Smooth a single diffusion model block.\n\n Args:\n layer (`DiffusionBlockStruct`):\n The diffusion block.\n config (`DiffusionQuantConfig`):\n The quantization configuration.\n smooth_cache (`dict[str, torch.Tensor]`):\n The smoothing scales cache.\n layer_cache (`dict[str, IOTensorsCache]`, *optional*):\n The layer cache.\n layer_kwargs (`dict[str, tp.Any]`, *optional*):\n The layer keyword arguments.\n \"\"\"\n logger = tools.logging.getLogger(f\"{__name__}.SmoothQuant\")\n logger.debug(\"- Smoothing Diffusion Block %s\", layer.name)\n tools.logging.Formatter.indent_inc()\n layer_cache = layer_cache or {}\n layer_kwargs = layer_kwargs or {}\n # We skip resnets since we currently cannot scale the Swish function\n visited: set[str] = set()\n for module_key, module_name, module, parent, _ in layer.named_key_modules():\n if isinstance(parent, (DiffusionAttentionStruct, DiffusionFeedForwardStruct)):\n block = parent.parent\n assert isinstance(block, DiffusionTransformerBlockStruct)\n if block.name not in visited:\n logger.debug(\"- Smoothing Transformer Block %s\", block.name)\n visited.add(block.name)\n tools.logging.Formatter.indent_inc()\n if block.parallel:\n smooth_cache = smooth_diffusion_parallel_transformer_block(\n block=block,\n config=config,\n smooth_cache=smooth_cache,\n block_cache=layer_cache,\n block_kwargs=layer_kwargs,\n )\n else:\n smooth_cache = smooth_diffusion_sequential_transformer_block(\n block=block,\n config=config,\n smooth_cache=smooth_cache,\n block_cache=layer_cache,\n block_kwargs=layer_kwargs,\n )\n tools.logging.Formatter.indent_dec()\n elif isinstance(module, (nn.Linear, nn.Conv2d)):\n smooth_cache = smooth_diffusion_module(\n module_key=module_key,\n module_name=module_name,\n module=module,\n config=config,\n smooth_cache=smooth_cache,\n layer_cache=layer_cache,\n )\n else:\n needs_quant = config.enabled_wgts and config.wgts.is_enabled_for(module_key)\n needs_quant = needs_quant or (config.enabled_ipts and config.ipts.is_enabled_for(module_key))\n if needs_quant and config.smooth.enabled_proj and config.smooth.proj.is_enabled_for(module_key):\n raise NotImplementedError(f\"Module {module_name} is not supported for smoothing\")\n logger.debug(\"- Skipping Module %s\", module_name)\n tools.logging.Formatter.indent_dec()\n\n\n@torch.inference_mode()\ndef smooth_diffusion(\n model: nn.Module | DiffusionModelStruct,\n config: DiffusionQuantConfig,\n smooth_cache: dict[str, torch.Tensor] | None = None,\n) -> dict[str, torch.Tensor]:\n \"\"\"Smooth the diffusion model.\n\n Args:\n model (`nn.Module` or `DiffusionModelStruct`):\n The diffusion model.\n config (`DiffusionQuantConfig`):\n The quantization configuration.\n smooth_cache (`dict[str, torch.Tensor]`, *optional*):\n The smoothing scales cache.\n\n Returns:\n `dict[str, torch.Tensor]`:\n The smoothing scales cache.\n \"\"\"\n if not isinstance(model, DiffusionModelStruct):\n model = DiffusionModelStruct.construct(model)\n assert isinstance(model, DiffusionModelStruct)\n smooth_cache = smooth_cache or {}\n if config.smooth.enabled_proj:\n if smooth_cache:\n assert smooth_cache.get(\"proj.fuse_when_possible\", True) == config.smooth.proj.fuse_when_possible\n if config.smooth.enabled_attn:\n if smooth_cache:\n assert smooth_cache.get(\"attn.fuse_when_possible\", True) == config.smooth.attn.fuse_when_possible\n if not smooth_cache:\n with tools.logging.redirect_tqdm():\n for _, (layer, layer_cache, layer_kwargs) in tqdm(\n config.calib.build_loader().iter_layer_activations(\n model,\n needs_inputs_fn=get_needs_inputs_fn(model, config),\n skip_pre_modules=True,\n skip_post_modules=True,\n ),\n desc=\"smoothing\",\n leave=False,\n total=model.num_blocks,\n dynamic_ncols=True,\n ):\n smooth_diffusion_layer(\n layer=layer,\n config=config,\n smooth_cache=smooth_cache,\n layer_cache=layer_cache,\n layer_kwargs=layer_kwargs,\n )\n else:\n for layer in model.block_structs:\n smooth_diffusion_layer(layer=layer, config=config, smooth_cache=smooth_cache)\n if config.smooth.enabled_proj:\n smooth_cache.setdefault(\"proj.fuse_when_possible\", config.smooth.proj.fuse_when_possible)\n if config.smooth.enabled_attn:\n smooth_cache.setdefault(\"attn.fuse_when_possible\", config.smooth.attn.fuse_when_possible)\n return smooth_cache\n"
},
{
"path": "deepcompressor/app/diffusion/quant/utils.py",
"content": "import typing as tp\n\nimport torch\nimport torch.nn as nn\n\nfrom ..nn.struct import DiffusionAttentionStruct, DiffusionFeedForwardStruct, DiffusionModelStruct\nfrom .config import DiffusionQuantConfig\n\n__all__ = [\"get_needs_inputs_fn\", \"get_needs_outputs_fn\", \"wrap_joint_attn\"]\n\n\ndef wrap_joint_attn(attn: nn.Module, /, *, indexes: int | tuple[int, ...] = 1) -> tp.Callable:\n if isinstance(indexes, int):\n\n def eval(*args, **kwargs) -> torch.Tensor:\n return attn(*args, **kwargs)[indexes]\n\n else:\n\n def eval(*args, **kwargs) -> tuple[torch.Tensor, ...]:\n tensors = attn(*args, **kwargs)\n result = torch.concat([tensors[i] for i in indexes], dim=-2)\n return result\n\n return eval\n\n\ndef get_needs_inputs_fn(\n model: DiffusionModelStruct, config: DiffusionQuantConfig\n) -> tp.Callable[[str, nn.Module], bool]:\n \"\"\"Get function that checks whether the module needs to cache inputs.\n\n Args:\n model (`DiffusionModelStruct`):\n The diffused model.\n config (`DiffusionQuantConfig`):\n The quantization configuration.\n\n Returns:\n `Callable[[str, nn.Module], bool]`:\n The function that checks whether the module needs to cache inputs.\n \"\"\"\n\n needs_inputs_names = set()\n for module_key, module_name, _, parent, field_name in model.named_key_modules():\n if (config.enabled_wgts and config.wgts.is_enabled_for(module_key)) or (\n config.enabled_ipts and config.ipts.is_enabled_for(module_key)\n ):\n if isinstance(parent, DiffusionAttentionStruct):\n if field_name.endswith(\"o_proj\"):\n needs_inputs_names.add(module_name)\n elif field_name in (\"q_proj\", \"k_proj\", \"v_proj\"):\n needs_inputs_names.add(parent.q_proj_name)\n if parent.parent.parallel and parent.idx == 0:\n needs_inputs_names.add(parent.parent.name)\n else:\n needs_inputs_names.add(parent.name)\n elif field_name in (\"add_q_proj\", \"add_k_proj\", \"add_v_proj\"):\n needs_inputs_names.add(parent.add_k_proj_name)\n if parent.parent.parallel and parent.idx == 0:\n needs_inputs_names.add(parent.parent.name)\n else:\n needs_inputs_names.add(parent.name)\n else:\n raise RuntimeError(f\"Unknown field name: {field_name}\")\n elif isinstance(parent, DiffusionFeedForwardStruct):\n if field_name == \"up_proj\":\n needs_inputs_names.update(parent.up_proj_names[: parent.config.num_experts])\n elif field_name == \"down_proj\":\n needs_inputs_names.update(parent.down_proj_names[: parent.config.num_experts])\n else:\n raise RuntimeError(f\"Unknown field name: {field_name}\")\n else:\n needs_inputs_names.add(module_name)\n\n def needs_inputs(name: str, module: nn.Module) -> bool:\n return name in needs_inputs_names\n\n return needs_inputs\n\n\ndef get_needs_outputs_fn(\n model: DiffusionModelStruct, config: DiffusionQuantConfig\n) -> tp.Callable[[str, nn.Module], bool]:\n \"\"\"Get function that checks whether the module needs to cache outputs.\n\n Args:\n model (`DiffusionModelStruct`):\n The diffused model.\n config (`DiffusionQuantConfig`):\n The quantization configuration.\n\n Returns:\n `Callable[[str, nn.Module], bool]`:\n The function that checks whether the module needs to cache outputs.\n \"\"\"\n\n # TODO: Implement the function that checks whether the module needs to cache outputs.\n\n def needs_outputs(name: str, module: nn.Module) -> bool:\n return False\n\n return needs_outputs\n"
},
{
"path": "deepcompressor/app/diffusion/quant/weight.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Diffusion model weight quantization calibration module.\"\"\"\n\nimport gc\nimport typing as tp\n\nimport torch\nimport torch.nn as nn\nfrom tqdm import tqdm\n\nfrom deepcompressor.data.cache import IOTensorsCache\nfrom deepcompressor.data.zero import ZeroPointDomain\nfrom deepcompressor.nn.patch.lowrank import LowRankBranch\nfrom deepcompressor.utils import tools\n\nfrom ..nn.struct import DiffusionAttentionStruct, DiffusionBlockStruct, DiffusionModelStruct, DiffusionModuleStruct\nfrom .config import DiffusionQuantConfig\nfrom .quantizer import DiffusionActivationQuantizer, DiffusionWeightQuantizer\nfrom .utils import get_needs_inputs_fn, wrap_joint_attn\n\n__all__ = [\"quantize_diffusion_weights\", \"load_diffusion_weights_state_dict\"]\n\n\n@torch.inference_mode()\ndef calibrate_diffusion_block_low_rank_branch( # noqa: C901\n layer: DiffusionModuleStruct | DiffusionBlockStruct,\n config: DiffusionQuantConfig,\n branch_state_dict: dict[str, dict[str, torch.Tensor]],\n layer_cache: dict[str, IOTensorsCache] = None,\n layer_kwargs: dict[str, tp.Any] = None,\n) -> None:\n \"\"\"Calibrate low-rank branches for a block of a diffusion model.\n\n Args:\n layer (`DiffusionModuleStruct` or `DiffusionBlockStruct`):\n The block to calibrate.\n config (`DiffusionQuantConfig`):\n The quantization configuration.\n branch_state_dict (`dict[str, dict[str, torch.Tensor]]`):\n The state dict of the low-rank branches.\n layer_cache (`dict[str, IOTensorsCache]`, *optional*, defaults to `None`):\n The cache of the layer.\n layer_kwargs (`dict[str, tp.Any]`, *optional*, defaults to `None`):\n The keyword arguments for the layer.\n \"\"\"\n assert config.wgts.low_rank is not None\n logger = tools.logging.getLogger(f\"{__name__}.WeightQuantSVD\")\n logger.debug(\"- Calibrating low-rank branches of block %s\", layer.name)\n layer_cache = layer_cache or {}\n layer_kwargs = layer_kwargs or {}\n for module_key, module_name, module, parent, field_name in layer.named_key_modules():\n modules, module_names = [module], [module_name]\n if not config.wgts.low_rank.exclusive:\n if field_name.endswith((\"q_proj\", \"k_proj\", \"v_proj\")):\n assert isinstance(parent, DiffusionAttentionStruct)\n if parent.is_self_attn():\n if field_name == \"q_proj\":\n modules, module_names = parent.qkv_proj, parent.qkv_proj_names\n else:\n continue\n elif parent.is_cross_attn():\n if field_name == \"add_k_proj\":\n modules.append(parent.add_v_proj)\n module_names.append(parent.add_v_proj_name)\n elif field_name != \"q_proj\":\n continue\n else:\n assert parent.is_joint_attn()\n if field_name == \"q_proj\":\n modules, module_names = parent.qkv_proj, parent.qkv_proj_names\n elif field_name == \"add_k_proj\":\n modules, module_names = parent.add_qkv_proj, parent.add_qkv_proj_names\n else:\n continue\n if field_name.endswith((\"q_proj\", \"k_proj\")):\n assert isinstance(parent, DiffusionAttentionStruct)\n if parent.parent.parallel and parent.idx == 0:\n eval_module = parent.parent.module\n eval_name = parent.parent.name\n eval_kwargs = layer_kwargs\n else:\n eval_module = parent.module\n eval_name = parent.name\n eval_kwargs = parent.filter_kwargs(layer_kwargs)\n if parent.is_joint_attn() and \"add_\" in field_name:\n eval_module = wrap_joint_attn(eval_module, indexes=1)\n else:\n eval_module, eval_name, eval_kwargs = module, module_name, None\n if isinstance(modules[0], nn.Linear):\n assert all(isinstance(m, nn.Linear) for m in modules)\n channels_dim = -1\n else:\n assert all(isinstance(m, nn.Conv2d) for m in modules)\n channels_dim = 1\n config_wgts = config.wgts\n if config.enabled_extra_wgts and config.extra_wgts.is_enabled_for(module_key):\n config_wgts = config.extra_wgts\n quantizer = DiffusionWeightQuantizer(config_wgts, develop_dtype=config.develop_dtype, key=module_key)\n if quantizer.is_enabled() and quantizer.is_enabled_low_rank():\n if isinstance(module, nn.Conv2d):\n assert module.weight.shape[2:].numel()\n else:\n assert isinstance(module, nn.Linear)\n if module_name not in branch_state_dict:\n logger.debug(\"- Calibrating low-rank branch for %s\", \", \".join(module_names))\n tools.logging.Formatter.indent_inc()\n branch_state_dict[module_name] = quantizer.calibrate_low_rank(\n input_quantizer=DiffusionActivationQuantizer(\n config.ipts, key=module_key, channels_dim=channels_dim\n ),\n modules=modules,\n inputs=layer_cache[module_name].inputs if layer_cache else None,\n eval_inputs=layer_cache[eval_name].inputs if layer_cache else None,\n eval_module=eval_module,\n eval_kwargs=eval_kwargs,\n ).state_dict()\n tools.logging.Formatter.indent_dec()\n gc.collect()\n torch.cuda.empty_cache()\n shared_branch = LowRankBranch(\n in_features=module.weight.shape[1],\n out_features=sum(m.weight.shape[0] for m in modules),\n rank=config.wgts.low_rank.rank,\n )\n shared_branch.to(device=module.weight.device, dtype=module.weight.dtype)\n shared_branch.load_state_dict(branch_state_dict[module_name])\n logger.debug(\" + Adding low-rank branches to %s\", \", \".join(module_names))\n if len(modules) > 1:\n oc_idx = 0\n for module in modules:\n branch = LowRankBranch(\n in_features=module.weight.shape[1],\n out_features=module.weight.shape[0],\n rank=config.wgts.low_rank.rank,\n )\n branch.a = shared_branch.a\n branch.b.to(dtype=module.weight.dtype, device=module.weight.device)\n branch.b.weight.copy_(shared_branch.b.weight[oc_idx : oc_idx + module.weight.shape[0]])\n oc_idx += module.weight.shape[0]\n module.weight.data.sub_(branch.get_effective_weight().view(module.weight.data.shape))\n branch.as_hook().register(module)\n else:\n module.weight.data.sub_(shared_branch.get_effective_weight().view(module.weight.data.shape))\n shared_branch.as_hook().register(module)\n del shared_branch\n gc.collect()\n torch.cuda.empty_cache()\n\n\n@torch.inference_mode()\ndef update_diffusion_block_weight_quantizer_state_dict(\n layer: DiffusionModuleStruct | DiffusionBlockStruct,\n config: DiffusionQuantConfig,\n quantizer_state_dict: dict[str, dict[str, torch.Tensor | float | None]],\n layer_cache: dict[str, IOTensorsCache],\n layer_kwargs: dict[str, tp.Any],\n):\n \"\"\"Update the state dict of the weight quantizers for a block of a diffusion model.\n\n Args:\n layer (`DiffusionModuleStruct` or `DiffusionBlockStruct`):\n The block to update.\n config (`DiffusionQuantConfig`):\n The quantization configuration.\n quantizer_state_dict (`dict[str, dict[str, torch.Tensor | float | None]]`):\n The state dict of the weight quantizers.\n layer_cache (`dict[str, IOTensorsCache]`):\n The cache of the layer.\n layer_kwargs (`dict[str, tp.Any]`):\n The keyword arguments for the layer.\n \"\"\"\n logger = tools.logging.getLogger(f\"{__name__}.WeightQuant\")\n logger.debug(\"- Calibrating weights: block %s\", layer.name)\n tools.logging.Formatter.indent_inc()\n for module_key, module_name, module, parent, field_name in layer.named_key_modules():\n if field_name.endswith((\"q_proj\", \"k_proj\")):\n assert isinstance(parent, DiffusionAttentionStruct)\n if parent.parent.parallel and parent.idx == 0:\n eval_module = parent.parent.module\n eval_name = parent.parent.name\n eval_kwargs = layer_kwargs\n else:\n eval_module = parent.module\n eval_name = parent.name\n eval_kwargs = parent.filter_kwargs(layer_kwargs)\n if parent.is_joint_attn() and \"add_\" in field_name:\n eval_module = wrap_joint_attn(eval_module, indexes=1)\n else:\n eval_module, eval_name, eval_kwargs = module, module_name, None\n config_wgts = config.wgts\n if config.enabled_extra_wgts and config.extra_wgts.is_enabled_for(module_key):\n config_wgts = config.extra_wgts\n quantizer = DiffusionWeightQuantizer(config_wgts, develop_dtype=config.develop_dtype, key=module_key)\n if quantizer.is_enabled():\n if module_name not in quantizer_state_dict:\n logger.debug(\"- Calibrating %s.weight quantizer\", module_name)\n quantizer.calibrate_dynamic_range(\n module=module,\n inputs=layer_cache[module_name].inputs if layer_cache else None,\n eval_inputs=layer_cache[eval_name].inputs if layer_cache else None,\n eval_module=eval_module,\n eval_kwargs=eval_kwargs,\n )\n quantizer_state_dict[module_name] = quantizer.state_dict()\n gc.collect()\n torch.cuda.empty_cache()\n else:\n logger.debug(\"- Loading %s.weight quantizer\", module_name)\n else:\n logger.debug(\"- Skipping %s.weight\", module_name)\n if module_name in quantizer_state_dict:\n quantizer_state_dict.pop(module_name)\n tools.logging.Formatter.indent_dec()\n\n\n@torch.inference_mode()\ndef quantize_diffusion_block_weights(\n layer: DiffusionModuleStruct | DiffusionBlockStruct,\n config: DiffusionQuantConfig,\n quantizer_state_dict: dict[str, dict[str, torch.Tensor | float | None]],\n layer_cache: dict[str, IOTensorsCache] = None,\n return_with_scale_state_dict: bool = False,\n) -> dict[str, torch.Tensor | float | None]:\n \"\"\"Quantize the weights of a block of a diffusion model.\n\n Args:\n layer (`DiffusionModuleStruct` or `DiffusionBlockStruct`):\n The block to quantize.\n config (`DiffusionQuantConfig`):\n The quantization configuration.\n quantizer_state_dict (`dict[str, dict[str, torch.Tensor | float | None]]`):\n The state dict of the weight quantizers.\n layer_cache (`dict[str, IOTensorsCache]`, *optional*, defaults to `None`):\n The cache of the layer.\n return_with_scale_state_dict (`bool`, *optional*, defaults to `False`):\n Whether to return the scale state dict.\n\n Returns:\n `dict[str, torch.Tensor | float | None]`:\n The scale state dict.\n \"\"\"\n logger = tools.logging.getLogger(f\"{__name__}.WeightQuant\")\n logger.debug(\"- Quantizing weights: block %s\", layer.name)\n layer_cache = layer_cache or {}\n\n scale_state_dict: dict[str, torch.Tensor | float | None] = {}\n\n tools.logging.Formatter.indent_inc()\n for module_key, module_name, module, _, _ in layer.named_key_modules():\n if module_name in quantizer_state_dict:\n param_name = f\"{module_name}.weight\"\n logger.debug(\"- Quantizing %s\", param_name)\n config_wgts = config.wgts\n if config.enabled_extra_wgts and config.extra_wgts.is_enabled_for(module_key):\n config_wgts = config.extra_wgts\n logger.debug(\" + quant_dtype: %s\", str(config_wgts.dtype))\n logger.debug(\" + group_shape: %s\", str(config_wgts.group_shapes))\n logger.debug(\" + scale_dtype: %s\", str(config_wgts.scale_dtypes))\n quantizer = DiffusionWeightQuantizer(config_wgts, develop_dtype=config.develop_dtype, key=module_key)\n quantizer.load_state_dict(quantizer_state_dict[module_name], device=module.weight.device)\n result = quantizer.quantize(\n module.weight.data,\n inputs=layer_cache[module_name].inputs.front() if layer_cache else None,\n return_with_dequant=True,\n return_with_quant=return_with_scale_state_dict,\n )\n if (\n config.wgts.enabled_low_rank\n and config.wgts.low_rank.is_enabled_for(module_key)\n and config.wgts.low_rank.compensate\n and config.wgts.low_rank.num_iters <= 1\n ):\n logger.debug(\"- Adding compensate low-rank branch to %s (side)\", module_name)\n LowRankBranch(\n in_features=module.weight.shape[1],\n out_features=module.weight.shape[0],\n rank=config.wgts.low_rank.rank,\n weight=module.weight.data - result.data,\n ).as_hook().register(module)\n module.weight.data = result.data\n if return_with_scale_state_dict:\n scale_state_dict.update(result.scale.state_dict(f\"{param_name}.scale\"))\n zero_name = \"scaled_zero\" if config.wgts.zero_point is ZeroPointDomain.PostScale else \"zero\"\n if isinstance(result.zero, torch.Tensor):\n scale_state_dict[f\"{param_name}.{zero_name}\"] = result.zero.to(\"cpu\")\n else:\n scale_state_dict[f\"{param_name}.{zero_name}\"] = result.zero\n del result\n gc.collect()\n torch.cuda.empty_cache()\n tools.logging.Formatter.indent_dec()\n return scale_state_dict\n\n\n@torch.inference_mode()\ndef quantize_diffusion_weights(\n model: nn.Module | DiffusionModelStruct,\n config: DiffusionQuantConfig,\n quantizer_state_dict: dict[str, dict[str, torch.Tensor | float | None]] | None = None,\n branch_state_dict: dict[str, dict[str, torch.Tensor]] | None = None,\n return_with_scale_state_dict: bool = False,\n) -> tuple[\n dict[str, dict[str, torch.Tensor | float | None]],\n dict[str, dict[str, torch.Tensor]],\n dict[str, torch.Tensor | float | None],\n]:\n \"\"\"Quantize the weights of a diffusion model.\n\n Args:\n model (`nn.Module` or `DiffusionModelStruct`):\n The diffusion model to quantize.\n config (`DiffusionQuantConfig`):\n The quantization configuration.\n quantizer_state_dict (`dict[str, dict[str, torch.Tensor | float | None]]`, *optional*, defaults to `None`):\n The state dict of the weight quantizers.\n branch_state_dict (`dict[str, dict[str, torch.Tensor]]`, *optional*, defaults to `None`):\n The state dict of the low-rank branches.\n return_with_scale_state_dict (`bool`, *optional*, defaults to `False`):\n Whether to return the scale state dict.\n\n Returns:\n `tuple[\n dict[str, dict[str, torch.Tensor | float | None]],\n dict[str, dict[str, torch.Tensor]],\n dict[str, torch.Tensor | float | None]\n ]`:\n The state dict of the weight quantizers, the state dict of the low-rank branches, and the scale state dict.\n \"\"\"\n logger = tools.logging.getLogger(f\"{__name__}.WeightQuant\")\n if not isinstance(model, DiffusionModelStruct):\n model = DiffusionModelStruct.construct(model)\n assert isinstance(model, DiffusionModelStruct)\n quantizer_state_dict = quantizer_state_dict or {}\n branch_state_dict = branch_state_dict or {}\n\n if config.wgts.enabled_low_rank and (not config.wgts.low_rank.compensate or config.wgts.low_rank.num_iters > 1):\n logger.info(\"* Adding low-rank branches to weights\")\n tools.logging.Formatter.indent_inc()\n with tools.logging.redirect_tqdm():\n if branch_state_dict:\n for _, layer in tqdm(\n model.get_named_layers(skip_pre_modules=True, skip_post_modules=True).items(),\n desc=\"adding low-rank branches\",\n leave=False,\n dynamic_ncols=True,\n ):\n calibrate_diffusion_block_low_rank_branch(\n layer=layer, config=config, branch_state_dict=branch_state_dict\n )\n else:\n for _, (layer, layer_cache, layer_kwargs) in tqdm(\n config.calib.build_loader().iter_layer_activations(\n model,\n needs_inputs_fn=get_needs_inputs_fn(model, config),\n skip_pre_modules=True,\n skip_post_modules=True,\n ),\n desc=\"calibrating low-rank branches\",\n leave=False,\n total=model.num_blocks,\n dynamic_ncols=True,\n ):\n calibrate_diffusion_block_low_rank_branch(\n layer=layer,\n config=config,\n branch_state_dict=branch_state_dict,\n layer_cache=layer_cache,\n layer_kwargs=layer_kwargs,\n )\n tools.logging.Formatter.indent_dec()\n\n skip_pre_modules = all(key in config.wgts.skips for key in model.get_prev_module_keys())\n skip_post_modules = all(key in config.wgts.skips for key in model.get_post_module_keys())\n with tools.logging.redirect_tqdm():\n if not quantizer_state_dict:\n if config.wgts.needs_calib_data:\n iterable = config.calib.build_loader().iter_layer_activations(\n model,\n needs_inputs_fn=get_needs_inputs_fn(model, config),\n skip_pre_modules=skip_pre_modules,\n skip_post_modules=skip_post_modules,\n )\n else:\n iterable = map( # noqa: C417\n lambda kv: (kv[0], (kv[1], {}, {})),\n model.get_named_layers(\n skip_pre_modules=skip_pre_modules, skip_post_modules=skip_post_modules\n ).items(),\n )\n for _, (layer, layer_cache, layer_kwargs) in tqdm(\n iterable,\n desc=\"calibrating weight quantizers\",\n leave=False,\n total=model.num_blocks + int(not skip_post_modules) + int(not skip_pre_modules) * 3,\n dynamic_ncols=True,\n ):\n update_diffusion_block_weight_quantizer_state_dict(\n layer=layer,\n config=config,\n quantizer_state_dict=quantizer_state_dict,\n layer_cache=layer_cache,\n layer_kwargs=layer_kwargs,\n )\n scale_state_dict: dict[str, torch.Tensor | float | None] = {}\n if config.wgts.enabled_gptq:\n iterable = config.calib.build_loader().iter_layer_activations(\n model,\n needs_inputs_fn=get_needs_inputs_fn(model, config),\n skip_pre_modules=skip_pre_modules,\n skip_post_modules=skip_post_modules,\n )\n else:\n iterable = map( # noqa: C417\n lambda kv: (kv[0], (kv[1], {}, {})),\n model.get_named_layers(skip_pre_modules=skip_pre_modules, skip_post_modules=skip_post_modules).items(),\n )\n for _, (layer, layer_cache, _) in tqdm(\n iterable,\n desc=\"quantizing weights\",\n leave=False,\n total=model.num_blocks + int(not skip_post_modules) + int(not skip_pre_modules) * 3,\n dynamic_ncols=True,\n ):\n layer_scale_state_dict = quantize_diffusion_block_weights(\n layer=layer,\n config=config,\n layer_cache=layer_cache,\n quantizer_state_dict=quantizer_state_dict,\n return_with_scale_state_dict=return_with_scale_state_dict,\n )\n scale_state_dict.update(layer_scale_state_dict)\n return quantizer_state_dict, branch_state_dict, scale_state_dict\n\n\n@torch.inference_mode()\ndef load_diffusion_weights_state_dict(\n model: nn.Module | DiffusionModelStruct,\n config: DiffusionQuantConfig,\n state_dict: dict[str, torch.Tensor],\n branch_state_dict: dict[str, dict[str, torch.Tensor]] | None = None,\n) -> None:\n \"\"\"Load the state dict of the weights of a diffusion model.\n\n Args:\n model (`nn.Module` or `DiffusionModelStruct`):\n The diffusion model to load the weights.\n config (`DiffusionQuantConfig`):\n The quantization configuration.\n state_dict (`dict[str, torch.Tensor]`):\n The state dict of the weights.\n branch_state_dict (`dict[str, dict[str, torch.Tensor]]`):\n The state dict of the low-rank branches.\n \"\"\"\n if not isinstance(model, DiffusionModelStruct):\n model = DiffusionModelStruct.construct(model)\n assert isinstance(model, DiffusionModelStruct)\n if config.enabled_wgts and config.wgts.enabled_low_rank:\n assert branch_state_dict is not None\n for _, layer in tqdm(\n model.get_named_layers(skip_pre_modules=True, skip_post_modules=True).items(),\n desc=\"adding low-rank branches\",\n leave=False,\n dynamic_ncols=True,\n ):\n calibrate_diffusion_block_low_rank_branch(layer=layer, config=config, branch_state_dict=branch_state_dict)\n model.module.load_state_dict(state_dict)\n gc.collect()\n torch.cuda.empty_cache()\n"
},
{
"path": "deepcompressor/app/diffusion/utils.py",
"content": "import os\nimport random\n\nimport numpy as np\nimport torch\nfrom PIL import Image\n\nfrom deepcompressor.utils.common import hash_str_to_int\n\n__all__ = [\"get_control\"]\n\n\ndef update_mask(mask: np.ndarray, x: int, y: int, radius: int | float):\n mask = mask.copy()\n H, W = mask.shape\n for i in range(H):\n for j in range(W):\n if (j - x) ** 2 + (i - y) ** 2 <= radius**2:\n mask[i, j] = True\n return mask\n\n\ndef generate_mask(\n masked_ratio_range: tuple[int, int], size: int | tuple[int, int], seed: int | None = None, eps=1e-2\n) -> np.ndarray:\n if seed is not None:\n random.seed(seed)\n masked_ratio = random.randint(masked_ratio_range[0], masked_ratio_range[1]) / 100\n if isinstance(size, int):\n size = (size, size)\n assert len(size) == 2\n height, width = size\n mask = np.zeros((height, width), dtype=bool)\n while True:\n radius = random.randint(16, min(height, width) // 2)\n x = random.randint(0, width - 1)\n y = random.randint(0, height - 1)\n new_mask = update_mask(mask, x, y, radius)\n if new_mask.sum() / (height * width) <= masked_ratio + eps:\n mask = new_mask\n if mask.sum() / (height * width) >= masked_ratio - eps:\n break\n return mask\n\n\ndef center_crop_and_resize(image: Image.Image, target_size: int | tuple[int, int]) -> Image.Image:\n if isinstance(target_size, int):\n target_size = (target_size, target_size)\n else:\n assert len(target_size) == 2\n target_width, target_height = target_size\n\n width, height = image.size\n if width / height > target_width / target_height:\n new_width = height * target_width / target_height\n left = round((width - new_width) / 2)\n right = round(left + new_width)\n image = image.crop((left, 0, right, height))\n elif width / height < width / height:\n new_height = width * target_height / target_width\n top = round((height - new_height) / 2)\n bottom = round(top + new_height)\n image = image.crop((0, top, width, bottom))\n width, height = image.size\n if width != target_width or height != target_height:\n image = image.resize((target_width, target_height), Image.Resampling.BICUBIC)\n return image\n\n\ndef get_control( # noqa: C901\n task: str,\n images: Image.Image | list[Image.Image],\n names: str | list[str] | None = None,\n data_root: str | None = None,\n device: str | torch.device = \"cuda\",\n **kwargs,\n) -> Image.Image | list[Image.Image] | tuple[Image.Image, Image.Image] | tuple[list[Image.Image], list[Image.Image]]:\n size = kwargs.get(\"size\", 1024)\n if isinstance(size, int):\n size = (size, size)\n assert len(size) == 2\n image_batch = [images] if isinstance(images, Image.Image) else images\n if isinstance(names, str):\n names = [names]\n\n if task == \"canny-to-image\":\n processor = kwargs.get(\"processor\", None)\n\n control_images = []\n for i, image in enumerate(image_batch):\n if data_root is not None and names is not None:\n data_path = os.path.join(data_root, \"canny_images\", f\"{names[i]}.png\")\n if os.path.exists(data_path):\n control_images.append(Image.open(data_path))\n continue\n if processor is None:\n from controlnet_aux import CannyDetector\n\n processor = CannyDetector()\n image = center_crop_and_resize(image, size)\n control_image = processor(\n image, low_threshold=50, high_threshold=200, detect_resolution=max(size), image_resolution=max(size)\n )\n control_images.append(control_image)\n if isinstance(images, Image.Image):\n return control_images[0]\n return control_images\n elif task == \"depth-to-image\":\n processor = kwargs.get(\"processor\", None)\n control_images = []\n for i, image in enumerate(image_batch):\n if data_root is not None and names is not None:\n data_path = os.path.join(data_root, \"depth_images\", f\"{names[i]}.png\")\n if os.path.exists(data_path):\n control_images.append(Image.open(data_path))\n continue\n if processor is None:\n from image_gen_aux import DepthPreprocessor\n\n processor = DepthPreprocessor.from_pretrained(\"LiheYoung/depth-anything-large-hf\").to(device)\n image = center_crop_and_resize(image, size)\n control_image = processor(image.convert(\"RGB\"))[0].convert(\"RGB\")\n control_images.append(control_image)\n if isinstance(images, Image.Image):\n return control_images[0]\n return control_images\n elif task == \"inpainting\":\n control_images, mask_images = [], []\n\n for i, image in enumerate(image_batch):\n name = None if names is None else names[i]\n if data_root is not None and name is not None:\n cropped_image_path = os.path.join(data_root, \"cropped_images\", f\"{name}.png\")\n mask_path = os.path.join(data_root, \"mask_images\", f\"{name}.png\")\n if os.path.exists(cropped_image_path) and os.path.exists(mask_path):\n control_images.append(Image.open(cropped_image_path).convert(\"RGB\"))\n mask_images.append(Image.open(mask_path))\n continue\n\n image = center_crop_and_resize(image, size)\n control_images.append(image.convert(\"RGB\"))\n if names is not None:\n seed = hash_str_to_int(names[i])\n else:\n seed = None\n\n mask = generate_mask((5, 60), size, seed=seed)\n mask_image = Image.fromarray(mask.astype(np.uint8) * 255)\n mask_images.append(mask_image)\n if isinstance(images, Image.Image):\n return control_images[0], mask_images[0]\n return control_images, mask_images\n else:\n raise ValueError(f\"Unsupported task: {task}\")\n"
},
{
"path": "deepcompressor/app/llm/__init__.py",
"content": "# -*- coding: utf-8 -*-\n"
},
{
"path": "deepcompressor/app/llm/cache/__init__.py",
"content": "# -*- coding: utf-8 -*-\n"
},
{
"path": "deepcompressor/app/llm/cache/config.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"LLM quantization cache configuration.\"\"\"\n\nfrom dataclasses import dataclass, field\n\nfrom omniconfig import configclass\n\nfrom deepcompressor.utils.config.path import BasePathConfig\n\n__all__ = [\"LlmQuantCacheConfig\", \"LlmCacheConfig\"]\n\n\n@configclass\n@dataclass\nclass LlmQuantCacheConfig(BasePathConfig):\n \"\"\"Large language model quantization cache path.\n\n Args:\n rotation (`str`, *optional*, default=`\"\"`):\n The rotation matrix cache path.\n reorder (`str`, *optional*, default=`\"\"`):\n The reorder channel indexes cache path.\n smooth (`str`, *optional*, default=`\"\"`):\n The smoothing scales cache path.\n wgts (`str`, *optional*, default=`\"\"`):\n The weight quantizers state dict cache path.\n acts (`str`, *optional*, default=`\"\"`):\n The activation quantizers state dict cache path.\n \"\"\"\n\n rotation: str = \"\"\n reorder: str = \"\"\n smooth: str = \"\"\n wgts: str = \"\"\n acts: str = \"\"\n\n\n@configclass\n@dataclass\nclass LlmCacheConfig:\n \"\"\"LLM quantization cache configuration.\n\n Attributes:\n root (`str`, *optional*, default=`\"\"`):\n The root directory path for the cache.\n dirpath (`LlmQuantCacheConfig`, *optional*, default=`LlmQuantCacheConfig()`):\n The directory paths for the cache.\n path (`LlmQuantCacheConfig`, *optional*, default=`LlmQuantCacheConfig()`):\n The file paths for the cache.\n \"\"\"\n\n root: str = field(default=\"\")\n dirpath: LlmQuantCacheConfig = field(init=False, default_factory=LlmQuantCacheConfig)\n path: LlmQuantCacheConfig = field(default_factory=LlmQuantCacheConfig)\n"
},
{
"path": "deepcompressor/app/llm/config.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Configurations for evaluating a large language model.\"\"\"\n\nimport os\nimport random\nfrom dataclasses import dataclass, field\n\nimport numpy as np\nimport omniconfig\nimport torch\nfrom omniconfig import ConfigParser, configclass\n\nfrom deepcompressor.data.utils import ScaleUtils\nfrom deepcompressor.utils.config.output import OutputConfig\n\nfrom .cache.config import LlmCacheConfig, LlmQuantCacheConfig\nfrom .eval.config import LlmEvalConfig\nfrom .model.config import LlmModelConfig\nfrom .quant.config import LlmQuantConfig\n\n__all__ = [\n \"LlmPtqRunConfig\",\n \"LlmCacheConfig\",\n \"LlmQuantCacheConfig\",\n \"LlmEvalConfig\",\n \"LlmModelConfig\",\n \"LlmQuantConfig\",\n]\n\n\n@configclass\n@dataclass\nclass LlmPtqRunConfig:\n \"\"\"Top-level config of post-training quantization for a large language model.\n\n Args:\n cache (`LlmCacheConfig`):\n Large language model quantization cache path configuration.\n output (`OutputConfig`):\n Output directory configuration.\n model (`LlmModelConfig`):\n Large language model configuration.\n eval (`LlmEvalConfig`):\n Large language model evaluation configuration.\n quant (`LlmQuantConfig`):\n Large language model quantization configuration.\n seed (`int`, *optional*, defaults to `12345`):\n Random seed.\n skip_eval (`bool`, *optional*, defaults to `False`):\n Whether to skip evaluation.\n load_model (`str`, *optional*, defaults to `\"\"`):\n Directory path to load the model checkpoint.\n save_model (`str`, *optional*, defaults to `\"\"`):\n Directory path to save the model checkpoint.\n copy_on_save (`bool`, *optional*, defaults to `False`):\n Whether to copy the quantization cache on save.\n \"\"\"\n\n cache: LlmCacheConfig\n output: OutputConfig\n model: LlmModelConfig\n eval: LlmEvalConfig\n quant: LlmQuantConfig = field(metadata={omniconfig.ARGPARSE_KWARGS: {\"prefix\": \"\"}})\n seed: int = 12345\n skip_eval: bool = False\n load_from: str = \"\"\n save_model: str = \"\"\n copy_on_save: bool = False\n\n def __post_init__(self): # noqa: C901\n # region set scale default dtype\n if self.quant.enabled_wgts:\n self.quant.wgts.scale_dtypes = tuple(\n ScaleUtils.infer_scale_dtypes(self.quant.wgts.scale_dtypes, default_dtype=self.model.dtype)\n )\n if self.quant.enabled_ipts:\n self.quant.ipts.scale_dtypes = tuple(\n ScaleUtils.infer_scale_dtypes(self.quant.ipts.scale_dtypes, default_dtype=self.model.dtype)\n )\n if self.quant.enabled_opts:\n self.quant.opts.scale_dtypes = tuple(\n ScaleUtils.infer_scale_dtypes(self.quant.opts.scale_dtypes, default_dtype=self.model.dtype)\n )\n # endregion\n # region set num_gpus and batch_size for auto parallelism of large models\n self.eval.num_gpus = min(torch.cuda.device_count(), self.eval.num_gpus)\n if self.model.size < 50:\n self.eval.batch_size = min(8, self.eval.batch_size)\n elif self.model.size < 100:\n self.eval.batch_size = min(4, self.eval.batch_size)\n else:\n self.eval.batch_size = min(1, self.eval.batch_size)\n # endregion\n if self.quant.is_enabled():\n if self.cache.path.is_all_empty():\n self.cache.dirpath = self.quant.generate_cache_dirpath(\n root=self.cache.root, seed=self.seed, default_dtype=self.model.dtype\n )\n self.cache.path = self.cache.dirpath.clone().add_children(f\"{self.model.name}.pt\")\n else:\n self.cache.dirpath = self.cache.path.clone().to_dirpath()\n if self.output.dirname == \"default\":\n self.output.dirname = self.quant.generate_default_dirname()\n self.output.dirpath = os.path.join(\n self.output.root,\n \"llm\",\n self.model.family,\n self.model.name,\n *self.quant.generate_dirnames(default_dtype=self.model.dtype)[:-1],\n self.quant.generate_calib_dirname(),\n self.output.dirname,\n )\n random.seed(self.seed)\n torch.manual_seed(self.seed)\n torch.cuda.manual_seed_all(self.seed)\n np.random.seed(self.seed)\n\n @classmethod\n def get_parser(cls) -> ConfigParser:\n \"\"\"Get a parser for evaluating a large language model.\n\n Returns:\n `ConfigParser`: A parser for evaluating a large language model.\n \"\"\"\n parser = ConfigParser(\"Evaluate a large language model\")\n parser.add_config(cls)\n return parser\n"
},
{
"path": "deepcompressor/app/llm/eval/__init__.py",
"content": "# -*- coding: utf-8 -*-\n"
},
{
"path": "deepcompressor/app/llm/eval/base.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Language model evaluator base.\"\"\"\n\nfrom abc import ABC, abstractmethod\n\nfrom transformers import PreTrainedModel, PreTrainedTokenizer\n\n__all__ = [\"LlmEvaluatorBase\"]\n\n\nclass LlmEvaluatorBase(ABC):\n def __init__(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer):\n self.model, self.tokenizer = model, tokenizer\n\n @abstractmethod\n def filter_tasks(self, tasks: list[str]) -> list[str]:\n \"\"\"Filter the tasks to only include supported tasks.\"\"\"\n ...\n\n @abstractmethod\n def evaluate(self, tasks: list[str], **kwargs) -> dict[str, dict[str, dict[str, float]]]:\n \"\"\"Evaluate the model on the given tasks.\"\"\"\n ...\n"
},
{
"path": "deepcompressor/app/llm/eval/config.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Language model evaluation config.\"\"\"\n\nimport random\nimport typing as tp\nfrom dataclasses import dataclass, field\n\nimport numpy as np\nimport omniconfig\nimport torch\nfrom omniconfig import configclass\nfrom transformers import PreTrainedModel, PreTrainedTokenizer\n\nfrom deepcompressor.utils import tools\n\nfrom .custom import LlmCustomEvaluator\nfrom .lm_eval import LmevalEvaluator\nfrom .longbench import LongbenchEvaluator\n\n__all__ = [\"LlmEvalConfig\"]\n\n\n@configclass\n@dataclass\nclass LlmEvalConfig:\n \"\"\"Large language model evaluation configuration.\n\n Attributes:\n num_gpus (`int`, *optional*, defaults to `1`):\n The number of GPUs to use.\n batch_size (`int`, *optional*, defaults to `1`):\n The batch size used for inference.\n tasks (`list[str]`, *optional*, defaults to `[\"zero-shot\"]`):\n Task names, e.g. wikitext, hellaswag, piqa, winogrande.\n max_seq_length (`int`, *optional*, defaults to `-4096`):\n Maximum sequence length.\n If negative, sequence lengths smaller than or equal to the absolute value are used.\n evaluators (`list[str]`, *optional*, defaults to `[\"gptq\"]`):\n Evaluators names.\n num_shot (`int`, *optional*, defaults to `None`):\n The number of shots for few-shot evaluation.\n fewshot_as_multiturn (`bool`, *optional*, defaults to `False`):\n Whether to treat few-shot evaluation as multi-turn.\n apply_chat_template (`bool`, *optional*, defaults to `False`):\n Whether to apply chat template for evaluation.\n \"\"\"\n\n num_gpus: int = field(default=1, metadata={omniconfig.ARGPARSE_ARGS: (\"--num-gpus\", \"-n\")})\n batch_size: int = 1\n tasks: list[str] = field(\n default_factory=lambda: [\"zero-shot\"],\n metadata={omniconfig.ARGPARSE_KWARGS: {\"nargs\": \"+\", \"type\": str}},\n )\n max_seq_length: int = -4096\n evaluators: list[str] = field(\n default_factory=lambda: [\"gptq\"], metadata={omniconfig.ARGPARSE_KWARGS: {\"nargs\": \"+\", \"type\": str}}\n )\n num_shot: int | None = None\n fewshot_as_multiturn: bool = False\n apply_chat_template: bool = False\n\n def __post_init__(self):\n if \"zero-shot\" in self.tasks:\n self.tasks.remove(\"zero-shot\")\n self.tasks.extend((\"wikitext\", \"hellaswag\", \"piqa\", \"winogrande\", \"arc_easy\", \"arc_challenge\"))\n self.tasks = sorted({tast.lower() for tast in self.tasks})\n self.evaluators = sorted({evaluator.lower() for evaluator in self.evaluators})\n for evaluator in self.evaluators:\n assert evaluator in (\"lm_eval\", \"gptq\", \"longbench\"), f\"Invalid evaluator: {evaluator}\"\n if len(self.evaluators) == 1 and self.evaluators[0] == \"gpq\":\n self.tasks = [task for task in self.tasks if task.startswith((\"wikitext\", \"pile\", \"gsm8k\"))]\n assert len(self.tasks) > 0, \"No valid tasks for GPTQ evaluation\"\n\n def evaluate(\n self,\n model: PreTrainedModel,\n /,\n tokenizer: PreTrainedTokenizer,\n model_name: str,\n eos_token_ids: tp.Sequence[int] = (),\n output_dirpath: str = \"\",\n ) -> dict[str, dict[int, dict[str, dict[tp.Any, dict[str, tp.Any]]]]]:\n \"\"\"Evaluate the model.\n\n Args:\n model (`PreTrainedModel`):\n The model.\n tokenizer (`PreTrainedTokenizer`):\n The tokenizer.\n model_name (`str`):\n The name of the model.\n eos_token_ids (`Sequence[int]`, *optional*, defaults to `()`):\n The EOS token IDs.\n\n Returns:\n `dict[str, dict[int, dict[str, dict[tp.Any, dict[str, tp.Any]]]]]`:\n The evaluation results.\n - The first key is the evaluator name.\n - The second key is the maximum sequence length.\n - The third key is the content name, e.g., \"results\", \"versions\", \"config\".\n - The fourth key is the task name for \"results\".\n \"\"\"\n logger = tools.logging.getLogger(f\"{__name__}.LlmEval\")\n tools.logging.Formatter.indent_inc()\n tools.logging.Formatter.indent_dec()\n lm_max_seq_length = get_max_seq_length(model, tokenizer)\n max_seq_lengths = {2048, 4096, lm_max_seq_length}\n if self.max_seq_length < 0:\n if self.max_seq_length == -1:\n max_seq_length = lm_max_seq_length\n else:\n max_seq_length = min(lm_max_seq_length, -self.max_seq_length)\n max_seq_lengths = [length for length in sorted(max_seq_lengths) if length <= max_seq_length]\n elif self.max_seq_length == 0:\n max_seq_lengths = [lm_max_seq_length]\n else:\n max_seq_lengths = [self.max_seq_length]\n results = {}\n for evaluator_name in self.evaluators:\n logger.info(f\"- Evaluator: {evaluator_name}\")\n tasks = list(self.tasks)\n if evaluator_name == \"gptq\":\n evaluator = LlmCustomEvaluator(model=model, tokenizer=tokenizer)\n elif evaluator_name == \"lm_eval\":\n evaluator = LmevalEvaluator(model=model, tokenizer=tokenizer, batch_size=self.batch_size)\n elif evaluator_name == \"longbench\":\n evaluator = LongbenchEvaluator(\n model=model,\n tokenizer=tokenizer,\n model_name=model_name,\n eos_token_ids=eos_token_ids,\n output_dirpath=output_dirpath,\n )\n else:\n raise ValueError(f\"Invalid evaluator: {evaluator_name}\")\n logger.info(f\"- Tasks: {tasks}\")\n logger.info(f\"- Batch_size: {self.batch_size}\")\n rsts = {}\n tools.logging.Formatter.indent_inc()\n for max_seq_length in max_seq_lengths:\n logger.info(f\"+ Max_seq_length: {max_seq_length}\")\n tools.logging.Formatter.indent_inc()\n tools.logging.Formatter.indent_inc()\n # set seed\n torch.manual_seed(42)\n torch.cuda.manual_seed(42)\n torch.cuda.manual_seed_all(42)\n np.random.seed(42)\n random.seed(42)\n # evaluate\n rst = evaluator.evaluate(\n tasks=tasks,\n max_length=max_seq_length,\n num_shot=self.num_shot,\n fewshot_as_multiturn=self.fewshot_as_multiturn,\n apply_chat_template=self.apply_chat_template,\n )\n rst[\"model\"] = model_name\n tools.logging.Formatter.indent_dec()\n logger.info(\"- Results:\")\n tools.logging.Formatter.indent_inc()\n tools.logging.info(self.make_table(rst), logger=logger)\n tools.logging.Formatter.indent_dec()\n rsts[max_seq_length] = rst\n tools.logging.Formatter.indent_dec()\n tools.logging.Formatter.indent_dec()\n results[evaluator_name] = rsts\n return results\n\n @staticmethod\n def make_table(rst: dict[str, dict[tp.Any, dict[str, tp.Any]]]) -> str:\n \"\"\"Generate table of results.\n\n Args:\n results (`dict[str, dict[tp.Any, dict[str, tp.Any]]]`):\n The evaluation results.\n\n Returns:\n `str`:\n The string representation of the results in a table.\n \"\"\"\n from pytablewriter import MarkdownTableWriter\n\n md_writer = MarkdownTableWriter()\n md_writer.headers = [\"Task\", \"Version\", \"Metric\", \"Value\", \"\", \"Stderr\"]\n values = []\n for k, dic in rst[\"results\"].items():\n version = rst[\"versions\"][k]\n for m, v in dic.items():\n if \"_stderr\" in m:\n continue\n mse = \"_stderr,\".join(m.split(\",\"))\n appended = False\n if mse in dic:\n se = dic[mse]\n if isinstance(se, (int, float)):\n values.append([k, version, m, \"%.4f\" % v, \"ยฑ\", \"%.4f\" % se])\n appended = True\n if not appended and isinstance(v, (int, float)):\n values.append([k, version, m, \"%.4f\" % v, \"\", \"\"])\n k = \"\"\n version = \"\"\n md_writer.value_matrix = values\n return md_writer.dumps()\n\n\ndef get_max_seq_length(model: PreTrainedModel, tokenizer: PreTrainedTokenizer, default_seq_length: int = 2048) -> int:\n seqlen_config_attrs = (\"n_positions\", \"max_position_embeddings\", \"n_ctx\")\n for attr in seqlen_config_attrs:\n if hasattr(model.config, attr):\n return getattr(model.config, attr)\n if hasattr(tokenizer, \"model_max_length\"):\n if tokenizer.model_max_length == 1000000000000000019884624838656:\n return default_seq_length\n return tokenizer.model_max_length\n return default_seq_length\n"
},
{
"path": "deepcompressor/app/llm/eval/custom.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Language model customized evaluator.\"\"\"\n\nimport math\n\nimport torch\nimport torch.nn as nn\nfrom datasets import load_dataset\nfrom tqdm import tqdm\nfrom transformers import PreTrainedModel, PreTrainedTokenizer\n\nfrom .base import LlmEvaluatorBase\n\n__all__ = [\"LlmCustomEvaluator\"]\n\n\nclass LlmCustomEvaluator(LlmEvaluatorBase):\n def filter_tasks(self, tasks: list[str]) -> list[str]:\n \"\"\"Filter the tasks to only include supported tasks.\"\"\"\n return [task for task in tasks if task.startswith((\"wikitext\", \"pile\"))]\n\n def evaluate(\n self, tasks: list[str], max_length: int | None = None, **kwargs\n ) -> dict[str, dict[str, dict[str, float]]]:\n \"\"\"Evaluate the model on the given tasks.\n\n Args:\n tasks (`list[str]`): List of tasks to evaluate on.\n max_length (`int`, optional, defaults to `None`): Maximum length for the model.\n\n Returns:\n dict[str, dict[str, dict[str, float]]]: Evaluation results `{\"results\": {\"task\": {\"metric\": score}}}`.\n \"\"\"\n result = {\"results\": {}, \"versions\": {}}\n for task in tasks:\n result[\"results\"][task] = {\n \"word_perplexity\": _eval_ppl_with_gptq_evaluator(\n self.model, self.tokenizer, task=task, seq_length=max_length\n )\n }\n result[\"versions\"][task] = 1\n return result\n\n\ndef _eval_ppl_with_gptq_evaluator(\n model: PreTrainedModel,\n /,\n tokenizer: PreTrainedTokenizer,\n task: str,\n seq_length: int = 2048,\n max_num_samples: int = -1,\n) -> float:\n \"\"\"Evaluate the perplexity of a model on a task using GPTQ style evaluation.\n\n Args:\n model (`PreTrainedModel`):\n The model.\n tokenizer (`PreTrainedTokenizer`):\n The tokenizer.\n task (`str`):\n The task name.\n seq_length (`int`, *optional*, defaults to `2048`):\n The sequence length.\n max_num_samples (`int`, *optional*, defaults to `-1`):\n The maximum number of samples to evaluate.\n\n Returns:\n float: The perplexity.\n \"\"\"\n assert seq_length > 0, \"seq_length must be positive\"\n if task.startswith(\"wikitext\"):\n test_dataset = load_dataset(\"wikitext\", \"wikitext-2-raw-v1\", split=\"test\")\n test_dataset = tokenizer(\"\\n\\n\".join(test_dataset[\"text\"]), return_tensors=\"pt\")\n elif task.startswith(\"pile\"):\n test_dataset = load_dataset(\"pile\", task, split=\"test\")\n test_dataset = tokenizer(\"\\n\\n\".join(test_dataset[\"text\"]), return_tensors=\"pt\")\n else:\n raise ValueError(f\"Invalid task: {task}\")\n\n test_dataset = test_dataset.input_ids.to(model.device)\n num_samples = test_dataset.numel() // seq_length\n if max_num_samples > 0:\n num_samples = min(num_samples, max_num_samples)\n model = model.eval()\n\n nlls = []\n for i in tqdm(range(num_samples), desc=f\"evaluating on {task} with seq_length {seq_length}\", dynamic_ncols=True):\n batch = test_dataset[:, (i * seq_length) : ((i + 1) * seq_length)]\n with torch.inference_mode():\n shift_logits = model(batch.to(model.device)).logits[:, :-1, :].contiguous().float()\n shift_labels = batch[:, 1:]\n loss = nn.CrossEntropyLoss()(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))\n neg_log_likelihood = loss.float() * seq_length\n nlls.append(neg_log_likelihood)\n return math.exp(sum(nlls) / (num_samples * seq_length))\n"
},
{
"path": "deepcompressor/app/llm/eval/lm_eval.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Language model evaluator using lm_eval.\"\"\"\n\nimport lm_eval\nimport lm_eval.models\nfrom transformers import PreTrainedModel, PreTrainedTokenizer\n\nfrom .base import LlmEvaluatorBase\n\n__all__ = [\"LmevalEvaluator\"]\n\n\nclass LmevalEvaluator(LlmEvaluatorBase):\n def __init__(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, batch_size: int = 1):\n super().__init__(model=model, tokenizer=tokenizer)\n self.lm = lm_eval.models.huggingface.HFLM(pretrained=model, tokenizer=tokenizer, batch_size=batch_size)\n\n def filter_tasks(self, tasks: list[str]) -> list[str]:\n \"\"\"Filter the tasks to only include supported tasks.\"\"\"\n return tasks\n\n def evaluate(\n self,\n tasks: list[str],\n max_length: int | None = None,\n num_shot: int | None = None,\n fewshot_as_multiturn: bool = False,\n apply_chat_template: bool = False,\n **kwargs,\n ) -> dict[str, dict[str, dict[str, float]]]:\n \"\"\"Evaluate the model on the given tasks.\n\n Args:\n tasks (`list[str]`): List of tasks to evaluate on.\n max_length (`int`, optional, defaults to `None`): Maximum length for the model.\n\n Returns:\n dict[str, dict[str, dict[str, float]]]: Evaluation results `{\"results\": {\"task\": {\"metric\": score}}}`.\n \"\"\"\n self.lm._max_length = max_length\n result = lm_eval.evaluator.simple_evaluate(\n model=self.lm,\n tasks=tasks,\n verbosity=\"ERROR\",\n num_fewshot=num_shot,\n fewshot_as_multiturn=fewshot_as_multiturn,\n apply_chat_template=apply_chat_template,\n **kwargs,\n )\n self.lm._max_length = None\n result.pop(\"samples\", None)\n result.pop(\"config\", None)\n return result\n"
},
{
"path": "deepcompressor/app/llm/eval/longbench/__init__.py",
"content": "from .eval import LongbenchEvaluator, LongbenchScorer\n"
},
{
"path": "deepcompressor/app/llm/eval/longbench/eval.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Language model evaluator for LongBench.\"\"\"\n\nimport json\nimport os\nimport typing as tp\n\nimport numpy as np\nimport torch\nimport torch.utils.data\nfrom datasets import load_dataset\nfrom tqdm import tqdm\nfrom transformers import PreTrainedModel, PreTrainedTokenizer\n\nfrom deepcompressor.utils import tools\n\nfrom ..base import LlmEvaluatorBase\nfrom .metrics import (\n classification_score,\n code_sim_score,\n count_score,\n qa_f1_score,\n qa_f1_zh_score,\n retrieval_score,\n retrieval_zh_score,\n rouge_score,\n rouge_zh_score,\n)\n\n__all__ = [\"LongbenchEvaluator\"]\n\n\nclass LongbenchEvaluator(LlmEvaluatorBase):\n task2maxlen: dict[str, int] = {\n \"narrativeqa\": 128,\n \"qasper\": 128,\n \"multifieldqa_en\": 64,\n \"multifieldqa_zh\": 64,\n \"hotpotqa\": 32,\n \"2wikimqa\": 32,\n \"musique\": 32,\n \"dureader\": 128,\n \"gov_report\": 512,\n \"qmsum\": 512,\n \"multi_news\": 512,\n \"vcsum\": 512,\n \"trec\": 64,\n \"triviaqa\": 32,\n \"samsum\": 128,\n \"lsht\": 64,\n \"passage_count\": 32,\n \"passage_retrieval_en\": 32,\n \"passage_retrieval_zh\": 32,\n \"lcc\": 64,\n \"repobench-p\": 64,\n }\n\n task2prompt: dict[str, str] = None\n\n def __init__(\n self,\n model: PreTrainedModel,\n tokenizer: PreTrainedTokenizer,\n model_name: str,\n eos_token_ids: tp.Sequence[int],\n output_dirpath: str = \"\",\n task2maxlen: dict[str, int] = None,\n task2prompt: dict[str, str] = None,\n ):\n super().__init__(model=model, tokenizer=tokenizer)\n self.model_name = model_name\n self.eos_token_ids = eos_token_ids\n if task2maxlen is not None:\n self.task2maxlen = task2maxlen\n if task2prompt is not None:\n self.task2prompt = task2prompt\n self.output_dirpath = output_dirpath\n self.logger = tools.logging.getLogger(__name__)\n\n def filter_tasks(self, tasks: list[str]) -> list[str]:\n \"\"\"Filter the tasks to only include supported tasks.\"\"\"\n if \"longbench-e\" in tasks:\n return [\"longbench-e\"]\n if \"longbench\" in tasks:\n return sorted(self.task2maxlen.keys(), key=lambda x: self.task2maxlen[x])\n return sorted([task for task in tasks if task in self.task2maxlen], key=lambda x: self.task2maxlen[x])\n\n def evaluate(self, tasks: list[str], max_length: int, **kwargs) -> dict[str, dict[str, dict[str, float]]]:\n \"\"\"Evaluate the model on the given tasks.\"\"\"\n ...\n tools.logging.Formatter.indent_inc()\n longbench_e = False\n if \"longbench-e\" in tasks:\n assert len(tasks) == 1, \"LongBench-E should be the only task\"\n longbench_e = True\n tasks = [\n \"hotpotqa\",\n \"2wikimqa\",\n \"triviaqa\",\n \"passage_count\",\n \"multifieldqa_en\",\n \"trec\",\n \"lcc\",\n \"repobench-p\",\n \"qasper\",\n \"samsum\",\n \"gov_report\",\n \"multi_news\",\n \"passage_retrieval_en\",\n ]\n result = {\"results\": {}, \"versions\": {}}\n for task in tasks:\n self.logger.info(f\"- Evaluating on {task}\")\n tools.logging.Formatter.indent_inc()\n preds = self.predict(task=task, max_length=max_length)\n if not preds:\n self.logger.warning(f\"No results for {task}\")\n tools.logging.Formatter.indent_dec()\n continue\n if self.output_dirpath:\n self.logger.info(f\"+ Saving results for {task} to {self.output_dirpath}\")\n os.makedirs(os.path.join(self.output_dirpath, \"longbench\"), exist_ok=True)\n with open(\n os.path.join(self.output_dirpath, \"longbench\", f\"{task}.json\"),\n \"w\",\n encoding=\"utf-8\",\n ) as f:\n for pred in preds:\n json.dump(pred, f, ensure_ascii=False)\n f.write(\"\\n\")\n predictions, answers, lengths = [], [], []\n for pred in preds:\n predictions.append(pred[\"prediction\"])\n answers.append(pred[\"answers\"])\n lengths.append(pred[\"length\"])\n all_classes = preds[0][\"all_classes\"]\n if longbench_e:\n scores = LongbenchScorer.scorer_e(\n task=task,\n predictions=predictions,\n answers=answers,\n lengths=lengths,\n all_classes=all_classes,\n )\n else:\n scores = {\n \"score\": LongbenchScorer.score(\n task=task,\n predictions=predictions,\n answers=answers,\n all_classes=all_classes,\n )\n }\n tools.logging.debug(f\"+ Scores: {scores}\", self.logger)\n result[\"results\"][task] = scores\n result[\"versions\"][task] = 1\n tools.logging.Formatter.indent_dec()\n tools.logging.Formatter.indent_dec()\n return result\n\n def predict(\n self,\n task: str,\n max_length: int,\n max_gen_length: int | None = None,\n prompt_format: str = \"\",\n ) -> list[dict[str, tp.Any]]:\n if max_gen_length is None:\n max_gen_length = self.task2maxlen[task]\n if prompt_format == \"\":\n prompt_format = self.task2prompt[task]\n dataset = load_dataset(\"THUDM/LongBench\", task, split=\"test\")\n preds = []\n pbar = tqdm(dataset)\n tools.logging.Formatter.indent_inc()\n for idx, data in enumerate(pbar):\n prompt = prompt_format.format(**data)\n # truncate to fit max_length\n # (we suggest truncate in the middle, since the left and right side may contain crucial instructions)\n tokenized_prompt = self.tokenizer(prompt, truncation=False, return_tensors=\"pt\").input_ids[0]\n if len(tokenized_prompt) > max_length:\n half = int(max_length / 2)\n prompt = self.tokenizer.decode(\n tokenized_prompt[:half], skip_special_tokens=True\n ) + self.tokenizer.decode(tokenized_prompt[-half:], skip_special_tokens=True)\n if task not in (\"trec\", \"triviaqa\", \"samsum\", \"lsht\", \"lcc\", \"repobench-p\"):\n # chat models are better off without build prompts on these tasks\n prompt = self.build_chat(prompt)\n input = self.tokenizer(prompt, truncation=False, return_tensors=\"pt\").to(\"cuda\")\n pbar.set_description(f\"Generating for {idx}, len={input.input_ids.shape[-1]}\")\n with torch.no_grad():\n output = self.model(input_ids=input.input_ids, past_key_values=None, use_cache=True)\n past_key_values = output.past_key_values\n pred_token_idx = output.logits[:, -1, :].argmax(dim=-1).unsqueeze(1)\n generated_content = [pred_token_idx.item()]\n for _ in range(max_gen_length - 1):\n outputs = self.model(input_ids=pred_token_idx, past_key_values=past_key_values, use_cache=True)\n past_key_values = outputs.past_key_values\n pred_token_idx = outputs.logits[:, -1, :].argmax(dim=-1).unsqueeze(1)\n generated_content += [pred_token_idx.item()]\n if pred_token_idx.item() in self.eos_token_ids:\n break\n pred = self.tokenizer.decode(generated_content, skip_special_tokens=True)\n pred = self.post_process(pred)\n # tools.logging.debug(f\"- Prediction: {pred}\", self.logger)\n preds.append(\n {\n \"prediction\": pred,\n \"answers\": data[\"answers\"],\n \"all_classes\": data[\"all_classes\"],\n \"length\": data[\"length\"],\n }\n )\n # break\n tools.logging.Formatter.indent_dec()\n return preds\n\n def build_chat(self, prompt):\n \"\"\"Build chat prompt for chat models.\"\"\"\n if \"llama-2\" in self.model_name:\n prompt = f\"[INST]{prompt}[/INST]\"\n return prompt\n\n def post_process(self, response: str) -> str:\n if \"xgen\" in self.model_name:\n response = response.strip().replace(\"Assistant:\", \"\")\n elif \"internlm\" in self.model_name:\n response = response.split(\"\")[0]\n elif \"llama-3\" in self.model_name:\n response = response.split(\".assistant\")[0].split(\"\\n\\nQuestion\")[0].split(\"\")[0].strip()\n elif \"llama-2-7b\" in self.model_name and \"instruct\" in self.model_name and \"32k\" in self.model_name:\n response = (\n response.split(\"(Document\")[0]\n .split(\"\\n\\nQuestion\")[0]\n .split(\"\\n\\nAnswer\")[0]\n .split(\"(Passage\")[0]\n .strip()\n )\n return response\n\n\nclass LongbenchScorer:\n task2metric = {\n \"narrativeqa\": qa_f1_score,\n \"qasper\": qa_f1_score,\n \"multifieldqa_en\": qa_f1_score,\n \"multifieldqa_zh\": qa_f1_zh_score,\n \"hotpotqa\": qa_f1_score,\n \"2wikimqa\": qa_f1_score,\n \"musique\": qa_f1_score,\n \"dureader\": rouge_zh_score,\n \"gov_report\": rouge_score,\n \"qmsum\": rouge_score,\n \"multi_news\": rouge_score,\n \"vcsum\": rouge_zh_score,\n \"trec\": classification_score,\n \"triviaqa\": qa_f1_score,\n \"samsum\": rouge_score,\n \"lsht\": classification_score,\n \"passage_retrieval_en\": retrieval_score,\n \"passage_count\": count_score,\n \"passage_retrieval_zh\": retrieval_zh_score,\n \"lcc\": code_sim_score,\n \"repobench-p\": code_sim_score,\n }\n\n @staticmethod\n def score(\n task: str,\n predictions: tp.Sequence[str],\n answers: tp.Sequence[tp.Sequence[str]],\n all_classes: tp.Sequence[str],\n task2metric: tp.Mapping[str, tp.Callable[[str, str, tp.Any], float]] = None,\n ) -> float:\n if task2metric is None:\n task2metric = LongbenchScorer.task2metric\n total_score = 0.0\n for prediction, ground_truths in zip(predictions, answers, strict=True):\n score = 0.0\n prediction = (\n prediction.split(\".assistant\")[0]\n .split(\"\\n\\nQuestion\")[0]\n .split(\"\")[0]\n .split(\"(Document\")[0]\n .split(\"\\n\\nQuestion\")[0]\n .split(\"\\n\\nAnswer\")[0]\n .split(\"(Passage\")[0]\n .strip()\n )\n if task in [\"trec\", \"triviaqa\", \"samsum\", \"lsht\"]:\n prediction = prediction.lstrip(\"\\n\").split(\"\\n\")[0]\n if task in [\"multifieldqa_zh\", \"dureader\"]:\n prediction = prediction.split(\"้ฎ้ข๏ผ\")[0].strip()\n if task in [\"lsht\"]:\n prediction = prediction.split(\"ๆฐ้ปๅ
ๅฎน๏ผ\")[0].strip()\n if task in [\"passage_retrieval_zh\"]:\n prediction = prediction.split(\"่ฏท้ฎ\")[0].split(\"ๆ็คบ\")[0].strip()\n for ground_truth in ground_truths:\n score = max(\n score,\n task2metric[task](prediction, ground_truth, all_classes=all_classes),\n )\n total_score += score\n return round(100 * total_score / len(predictions), 2)\n\n @staticmethod\n def scorer_e(\n task: str,\n predictions: tp.Sequence[str],\n answers: tp.Sequence[tp.Sequence[str]],\n lengths: tp.Sequence[int],\n all_classes: tp.Sequence[str],\n task2metric: tp.Mapping[str, tp.Callable[[str, str, tp.Any], float]] = None,\n ) -> dict[str, float]:\n if task2metric is None:\n task2metric = LongbenchScorer.task2metric\n scores = {\"0-4k\": [], \"4-8k\": [], \"8k+\": []}\n for prediction, ground_truths, length in zip(predictions, answers, lengths, strict=True):\n score = 0.0\n if task in [\"trec\", \"triviaqa\", \"samsum\", \"lsht\"]:\n prediction = prediction.lstrip(\"\\n\").split(\"\\n\")[0]\n for ground_truth in ground_truths:\n score = max(\n score,\n task2metric[task](prediction, ground_truth, all_classes=all_classes),\n )\n if length < 4000:\n scores[\"0-4k\"].append(score)\n elif length < 8000:\n scores[\"4-8k\"].append(score)\n else:\n scores[\"8k+\"].append(score)\n for key in scores.keys():\n scores[key] = round(100 * np.mean(scores[key]), 2)\n return scores\n\n\n# Initialize the evaluator task2prompt by loading the json file\nwith open(os.path.join(os.path.dirname(__file__), \"task2prompt.json\")) as f:\n LongbenchEvaluator.task2prompt = json.load(f)\n"
},
{
"path": "deepcompressor/app/llm/eval/longbench/metrics.py",
"content": "\"\"\"LongBench metrics.\"\"\"\n\nimport re\nimport string\nfrom collections import Counter\n\nimport jieba\nfrom fuzzywuzzy import fuzz\nfrom rouge import Rouge\n\n__all__ = [\n \"classification_score\",\n \"code_sim_score\",\n \"count_score\",\n \"qa_f1_score\",\n \"qa_f1_zh_score\",\n \"retrieval_score\",\n \"retrieval_zh_score\",\n \"rouge_score\",\n \"rouge_zh_score\",\n]\n\n\ndef normalize_answer(s: str) -> str:\n \"\"\"Lower text and remove punctuation, articles and extra whitespace.\"\"\"\n\n def remove_articles(text: str) -> str:\n return re.sub(r\"\\b(a|an|the)\\b\", \" \", text)\n\n def white_space_fix(text: str) -> str:\n return \" \".join(text.split())\n\n def remove_punc(text: str) -> str:\n exclude = set(string.punctuation)\n return \"\".join(ch for ch in text if ch not in exclude)\n\n return white_space_fix(remove_articles(remove_punc(s.lower())))\n\n\ndef normalize_zh_answer(s: str) -> str:\n \"\"\"Lower text and remove punctuation, extra whitespace.\"\"\"\n\n def white_space_fix(text):\n return \"\".join(text.split())\n\n def remove_punc(text):\n exclude = set(\n string.punctuation + \"๏ผ๏ผ๏ฝกใ๏ผ๏ผ๏ผ๏ผ
๏ผ๏ผ๏ผ๏ผ๏ผ๏ผ๏ผ๏ผ๏ผ๏ผ๏ผ๏ผ๏ผ๏ผ๏ผ ๏ผป๏ผผ๏ผฝ๏ผพ๏ผฟ๏ฝ๏ฝ๏ฝ๏ฝ๏ฝ\"\n \"๏ฝ๏ฝ ๏ฝข๏ฝฃ๏ฝคใใใใใใใใใใใใใใใใใใใใใใฐใพใฟโโโโโโโโโโฆโง๏น.\"\n )\n return \"\".join(ch for ch in text if ch not in exclude)\n\n return white_space_fix(remove_punc(s.lower()))\n\n\ndef count_score(prediction: str, ground_truth: str, **kwargs) -> float:\n numbers = re.findall(r\"\\d+\", prediction)\n right_num = 0\n for number in numbers:\n if str(number) == str(ground_truth):\n right_num += 1\n final_score = 0.0 if len(numbers) == 0 else right_num / len(numbers)\n return float(final_score)\n\n\ndef retrieval_score(prediction: str, ground_truth: str, **kwargs) -> float:\n pattern = r\"Paragraph (\\d+)\"\n matches = re.findall(pattern, ground_truth)\n ground_truth_id = matches[0]\n numbers = re.findall(r\"\\d+\", prediction)\n right_num = 0\n for number in numbers:\n if str(number) == str(ground_truth_id):\n right_num += 1\n return 0.0 if len(numbers) == 0 else right_num / len(numbers)\n\n\ndef retrieval_zh_score(prediction: str, ground_truth: str, **kwargs) -> float:\n pattern = r\"ๆฎต่ฝ(\\d+)\"\n matches = re.findall(pattern, ground_truth)\n ground_truth_id = matches[0]\n numbers = re.findall(r\"\\d+\", prediction)\n right_num = 0\n for number in numbers:\n if str(number) == str(ground_truth_id):\n right_num += 1\n return 0.0 if len(numbers) == 0 else right_num / len(numbers)\n\n\ndef code_sim_score(prediction: str, ground_truth: str, **kwargs) -> float:\n all_lines = prediction.lstrip(\"\\n\").split(\"\\n\")\n prediction = \"\"\n for line in all_lines:\n if (\"`\" not in line) and (\"#\" not in line) and (\"//\" not in line):\n prediction = line\n break\n return fuzz.ratio(prediction, ground_truth) / 100\n\n\ndef classification_score(prediction: str, ground_truth: str, **kwargs) -> float:\n em_match_list = [\n class_name\n for class_name in kwargs[\"all_classes\"]\n if class_name in prediction and not (class_name in ground_truth and class_name != ground_truth)\n ]\n return 1.0 / len(em_match_list) if ground_truth in em_match_list else 0.0\n\n\ndef rouge_score(prediction: str, ground_truth: str, **kwargs) -> float:\n try:\n scores = Rouge().get_scores([prediction], [ground_truth], avg=True)\n except Exception:\n return 0.0\n return scores[\"rouge-l\"][\"f\"]\n\n\ndef rouge_zh_score(prediction: str, ground_truth: str, **kwargs) -> float:\n prediction = \" \".join(list(jieba.cut(prediction, cut_all=False)))\n ground_truth = \" \".join(list(jieba.cut(ground_truth, cut_all=False)))\n return rouge_score(prediction, ground_truth)\n\n\ndef f1_score(prediction: str, ground_truth: str, **kwargs) -> float:\n common = Counter(prediction) & Counter(ground_truth)\n num_same = sum(common.values())\n if num_same == 0:\n return 0\n precision = 1.0 * num_same / len(prediction)\n recall = 1.0 * num_same / len(ground_truth)\n return (2 * precision * recall) / (precision + recall)\n\n\ndef qa_f1_score(prediction: str, ground_truth: str, **kwargs) -> float:\n normalized_prediction = normalize_answer(prediction)\n normalized_ground_truth = normalize_answer(ground_truth)\n prediction_tokens = normalized_prediction.split()\n ground_truth_tokens = normalized_ground_truth.split()\n return f1_score(prediction_tokens, ground_truth_tokens)\n\n\ndef qa_f1_zh_score(prediction: str, ground_truth: str, **kwargs) -> float:\n prediction_tokens = list(jieba.cut(prediction, cut_all=False))\n ground_truth_tokens = list(jieba.cut(ground_truth, cut_all=False))\n prediction_tokens = [normalize_zh_answer(token) for token in prediction_tokens]\n ground_truth_tokens = [normalize_zh_answer(token) for token in ground_truth_tokens]\n prediction_tokens = [token for token in prediction_tokens if len(token) > 0]\n ground_truth_tokens = [token for token in ground_truth_tokens if len(token) > 0]\n return f1_score(prediction_tokens, ground_truth_tokens)\n"
},
{
"path": "deepcompressor/app/llm/eval/longbench/task2prompt.json",
"content": "{\n \"narrativeqa\": \"You are given a story, which can be either a novel or a movie script, and a question. Answer the question asconcisely as you can, using a single phrase if possible. Do not provide any explanation.\\n\\nStory: {context}\\n\\nNow, answer the question based on the story asconcisely as you can, using a single phrase if possible. Do not provide any explanation.\\n\\nQuestion: {input}\\n\\nAnswer:\",\n \"qasper\": \"You are given a scientific article and a question. Answer the question as concisely as you can, using a single phrase or sentence if possible. If the question cannot be answered based on the information in the article, write \\\"unanswerable\\\". If the question is a yes/no question, answer \\\"yes\\\", \\\"no\\\", or \\\"unanswerable\\\". Do not provide any explanation.\\n\\nArticle: {context}\\n\\n Answer the question based on the above article as concisely as you can, using a single phrase or sentence if possible. If the question cannot be answered based on the information in the article, write \\\"unanswerable\\\". If the question is a yes/no question, answer \\\"yes\\\", \\\"no\\\", or \\\"unanswerable\\\". Do not provide any explanation.\\n\\nQuestion: {input}\\n\\nAnswer:\",\n \"multifieldqa_en\": \"Read the following text and answer briefly.\\n\\n{context}\\n\\nNow, answer the following question based on the above text, only give me the answer and do not output any other words.\\n\\nQuestion: {input}\\nAnswer:\",\n \"multifieldqa_zh\": \"้
่ฏปไปฅไธๆๅญๅนถ็จไธญๆ็ฎ็ญๅ็ญ๏ผ\\n\\n{context}\\n\\n็ฐๅจ่ฏทๅบไบไธ้ข็ๆ็ซ ๅ็ญไธ้ข็้ฎ้ข๏ผๅชๅ่ฏๆ็ญๆก๏ผไธ่ฆ่พๅบไปปไฝๅ
ถไปๅญ่ฏใ\\n\\n้ฎ้ข๏ผ{input}\\nๅ็ญ๏ผ\",\n \"hotpotqa\": \"Answer the question based on the given passages. Only give me the answer and do not output any other words.\\n\\nThe following are given passages.\\n{context}\\n\\nAnswer the question based on the given passages. Only give me the answer and do not output any other words.\\n\\nQuestion: {input}\\nAnswer:\",\n \"2wikimqa\": \"Answer the question based on the given passages. Only give me the answer and do not output any other words.\\n\\nThe following are given passages.\\n{context}\\n\\nAnswer the question based on the given passages. Only give me the answer and do not output any other words.\\n\\nQuestion: {input}\\nAnswer:\",\n \"musique\": \"Answer the question based on the given passages. Only give me the answer and do not output any other words.\\n\\nThe following are given passages.\\n{context}\\n\\nAnswer the question based on the given passages. Only give me the answer and do not output any other words.\\n\\nQuestion: {input}\\nAnswer:\",\n \"dureader\": \"่ฏทๅบไบ็ปๅฎ็ๆ็ซ ๅ็ญไธ่ฟฐ้ฎ้ขใ\\n\\nๆ็ซ ๏ผ{context}\\n\\n่ฏทๅบไบไธ่ฟฐๆ็ซ ๅ็ญไธ้ข็้ฎ้ขใ\\n\\n้ฎ้ข๏ผ{input}\\nๅ็ญ๏ผ\",\n \"gov_report\": \"You are given a report by a government agency. Write a one-page summary of the report.\\n\\nReport:\\n{context}\\n\\nNow, write a one-page summary of the report.\\n\\nSummary:\",\n \"qmsum\": \"You are given a meeting transcript and a query containing a question or instruction. Answer the query in one or more sentences.\\n\\nTranscript:\\n{context}\\n\\nNow, answer the query based on the above meeting transcript in one or more sentences.\\n\\nQuery: {input}\\nAnswer:\",\n \"multi_news\": \"You are given several news passages. Write a one-page summary of all news. \\n\\nNews:\\n{context}\\n\\nNow, write a one-page summary of all the news.\\n\\nSummary:\",\n \"vcsum\": \"ไธ้ขๆไธๆฎตไผ่ฎฎ่ฎฐๅฝ๏ผ่ฏทไฝ ้
่ฏปๅ๏ผๅไธๆฎตๆป็ป๏ผๆป็ปไผ่ฎฎ็ๅ
ๅฎนใ\\nไผ่ฎฎ่ฎฐๅฝ๏ผ\\n{context}\\n\\nไผ่ฎฎๆป็ป๏ผ\",\n \"trec\": \"Please determine the type of the question below. Here are some examples of questions.\\n\\n{context}\\n{input}\",\n \"triviaqa\": \"Answer the question based on the given passage. Only give me the answer and do not output any other words. The following are some examples.\\n\\n{context}\\n\\n{input}\",\n \"samsum\": \"Summarize the dialogue into a few short sentences. The following are some examples.\\n\\n{context}\\n\\n{input}\",\n \"lsht\": \"่ฏทๅคๆญ็ปๅฎๆฐ้ป็็ฑปๅซ๏ผไธ้ขๆฏไธไบไพๅญใ\\n\\n{context}\\n{input}\",\n \"passage_count\": \"There are some paragraphs below sourced from Wikipedia. Some of them may be duplicates. Please carefully read these paragraphs and determine how many unique paragraphs there are after removing duplicates. In other words, how many non-repeating paragraphs are there in total?\\n\\n{context}\\n\\nPlease enter the final count of unique paragraphs after removing duplicates. The output format should only contain the number, such as 1, 2, 3, and so on.\\n\\nThe final answer is: \",\n \"passage_retrieval_en\": \"Here are 30 paragraphs from Wikipedia, along with an abstract. Please determine which paragraph the abstract is from.\\n\\n{context}\\n\\nThe following is an abstract.\\n\\n{input}\\n\\nPlease enter the number of the paragraph that the abstract is from. The answer format must be like \\\"Paragraph 1\\\", \\\"Paragraph 2\\\", etc.\\n\\nThe answer is: \",\n \"passage_retrieval_zh\": \"ไปฅไธๆฏ่ฅๅนฒๆฎต่ฝๆๅญ๏ผไปฅๅๅ
ถไธญไธไธชๆฎต่ฝ็ๆ่ฆใ่ฏท็กฎๅฎ็ปๅฎ็ๆ่ฆๅบ่ชๅชไธๆฎตใ\\n\\n{context}\\n\\nไธ้ขๆฏไธไธชๆ่ฆ\\n\\n{input}\\n\\n่ฏท่พๅ
ฅๆ่ฆๆๅฑๆฎต่ฝ็็ผๅทใ็ญๆกๆ ผๅผๅฟ
้กปๆฏ\\\"ๆฎต่ฝ1\\\"๏ผ\\\"ๆฎต่ฝ2\\\"็ญๆ ผๅผ\\n\\n็ญๆกๆฏ๏ผ\",\n \"lcc\": \"Please complete the code given below. \\n{context}Next line of code:\\n\",\n \"repobench-p\": \"Please complete the code given below. \\n{context}{input}Next line of code:\\n\"\n}"
},
{
"path": "deepcompressor/app/llm/model/__init__.py",
"content": "# -*- coding: utf-8 -*-\n"
},
{
"path": "deepcompressor/app/llm/model/config.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Net configurations.\"\"\"\n\nimport typing as tp\nfrom dataclasses import dataclass, field\n\nimport torch\nfrom omniconfig import configclass\nfrom transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PreTrainedTokenizer\n\nfrom deepcompressor.data.utils.dtype import eval_dtype\nfrom deepcompressor.utils.config.model import BaseModelConfig\n\nfrom ..nn.patch import patch_attention, patch_gemma_rms_norm\n\n__all__ = [\"LlmModelConfig\"]\n\n\n@configclass\n@dataclass\nclass LlmModelConfig(BaseModelConfig):\n \"\"\"Arguments for creating a large language model.\n\n Args:\n name (`str`):\n Name of the model.\n path (`str`, *optional*, defaults to `\"\"`):\n Path of the model.\n root (`str`, *optional*, defaults to `\"\"`):\n Root directory path for models.\n local_path (`str`, *optional*, defaults to `\"\"`):\n Local path of the model.\n local_root (`str`, *optional*, defaults to `\"\"`):\n Local root directory path for models.\n dtype (`torch.dtype`, *optional*, defaults to `None`):\n Data type of the model. If not specified, the original data type of the model will be used.\n fast_tokenizer (`bool`, *optional*, defaults to `True`):\n Whether to use fast tokenizer.\n\n Attributes:\n size (`float`):\n Size of the model.\n variant (`str`):\n Variant of the model.\n \"\"\"\n\n _model_factories: tp.ClassVar[dict[str, tp.Callable[[str], tuple[PreTrainedModel, PreTrainedTokenizer]]]] = {}\n\n size: float = field(init=False)\n variant: str = field(init=False)\n dtype: torch.dtype = field(default_factory=lambda s=None: eval_dtype(s, with_quant_dtype=False))\n use_flash_attn: bool = False\n fast_tokenizer: bool = True\n orig_dtype: torch.dtype = field(init=False)\n\n def __post_init__(self):\n parts = self.name.split(\"-\")\n # we first infer the size, it should be a string matching \"$\\d+[mb]$\"\n family, size, variant = \"\", \"\", \"\"\n for i, part in enumerate(parts):\n part = part.lower()\n if part[-1] == \"m\" or part[-1] == \"b\":\n _part = part[:-1].replace(\"x\", \"\", 1)\n if _part.isdigit():\n size = part\n family = \"-\".join(parts[:i])\n if len(parts) > i + 1:\n variant = \"-\".join(parts[i + 1 :])\n break\n assert size, f\"Cannot infer size from {self.name}\"\n assert family, f\"Cannot infer family from {self.name}\"\n if not self.family:\n self.family = family\n self.variant = variant\n if size[-1] == \"m\":\n size = float(size[:-1]) / 1000\n else:\n assert size[-1] == \"b\"\n size = size[:-1]\n if \"x\" in size:\n num_experts, expert_gb = size.split(\"x\")\n num_experts = int(num_experts)\n expert_size = float(expert_gb)\n size = num_experts * expert_size\n else:\n size = float(size)\n self.size = size\n super().__post_init__()\n self.name = self.name.lower()\n self.family = self.family.lower()\n self.variant = self.variant.lower()\n config = AutoConfig.from_pretrained(self.path)\n self.orig_dtype = config.torch_dtype\n if self.orig_dtype == torch.float32:\n self.dtype = self.dtype or torch.float16\n elif self.orig_dtype == torch.float16:\n self.dtype = self.dtype or torch.float16\n elif self.orig_dtype == torch.bfloat16:\n self.dtype = self.dtype or torch.bfloat16\n else:\n raise ValueError(f\"Unsupported data type: {self.orig_dtype}\")\n\n def build(self) -> tuple[PreTrainedModel, PreTrainedTokenizer]:\n \"\"\"Build model and tokenizer.\n\n Args:\n dtype (`torch.dtype`, *optional*, defaults to `None`):\n Data type of the model.\n\n Returns:\n `tuple[PreTrainedModel, PreTrainedTokenizer]`:\n Model and tokenizer.\n \"\"\"\n torch_dtype = self.dtype\n if self.name in self._model_factories:\n return self._model_factories[self.name](\n self.path, torch_dtype=torch_dtype, use_fast=self.fast_tokenizer, use_flash_attn=self.use_flash_attn\n )\n kwargs = {\"torch_dtype\": torch_dtype}\n if torch.cuda.is_available() and torch.cuda.device_count() > 0:\n kwargs[\"device_map\"] = \"balanced\"\n return self._default_build(self.path, **kwargs)\n\n @staticmethod\n def _default_build(path: str, **kwargs) -> tuple[PreTrainedModel, PreTrainedTokenizer]:\n \"\"\"Build model and tokenizer.\n\n Args:\n dtype (`torch.dtype`, *optional*, defaults to `None`):\n Data type of the model.\n\n Returns:\n `tuple[PreTrainedModel, PreTrainedTokenizer]`:\n Model and tokenizer.\n \"\"\"\n config = AutoConfig.from_pretrained(path)\n tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True, use_fast=kwargs.pop(\"use_fast\", True))\n if \"use_flash_attn\" in kwargs:\n use_flash_attn = kwargs.pop(\"use_flash_attn\")\n if use_flash_attn:\n kwargs[\"attn_implementation\"] = \"flash_attention_2\"\n model = AutoModelForCausalLM.from_pretrained(path, config=config, **kwargs)\n patch_attention(model)\n patch_gemma_rms_norm(model)\n model.eval()\n return model, tokenizer\n\n @classmethod\n def register_model_factory(\n cls,\n names: str | tuple[str, ...],\n /,\n factory: tp.Callable[[str, torch.dtype], tuple[PreTrainedModel, PreTrainedTokenizer]],\n *,\n overwrite: bool = False,\n ) -> None:\n \"\"\"Register a model factory.\n\n Args:\n names (`str` or `tuple[str, ...]`):\n Names of the model.\n factory (`Callable[[str, torch.dtype], tuple[PreTrainedModel, PreTrainedTokenizer]]`):\n Factory function.\n overwrite (`bool`, *optional*, defaults to `False`):\n Whether to overwrite the existing factory for the model.\n \"\"\"\n if isinstance(names, str):\n names = (names,)\n for name in names:\n if not overwrite and name in cls._model_factories:\n raise ValueError(f\"Factory for {name} already exists\")\n cls._model_factories[name] = factory\n"
},
{
"path": "deepcompressor/app/llm/nn/__init__.py",
"content": "# -*- coding: utf-8 -*-\n\nfrom .struct import LlmModelStruct, LlmTransformerBlockStruct, LlmTransformerStruct\n"
},
{
"path": "deepcompressor/app/llm/nn/patch.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Llama model patcher.\"\"\"\n\nimport functools\n\nimport torch\nimport torch.nn as nn\nfrom transformers.models.gemma.modeling_gemma import GemmaRMSNorm\nfrom transformers.models.gemma2.modeling_gemma2 import Gemma2RMSNorm\n\nfrom deepcompressor.utils import tools\nfrom deepcompressor.utils.patch import copy_func\n\n__all__ = [\"patch_attention\", \"patch_gemma_rms_norm\", \"RotaryEmbedding\"]\n\n\ndef rotate_half(x):\n \"\"\"Rotates half the hidden dims of the input.\"\"\"\n x1 = x[..., : x.shape[-1] // 2]\n x2 = x[..., x.shape[-1] // 2 :]\n return torch.cat((-x2, x1), dim=-1)\n\n\ndef update_rotary_cos_sin(\n cos: torch.Tensor, sin: torch.Tensor, position_ids: torch.LongTensor | None, unsqueeze_dim: int = 1\n) -> tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Update the cos and sin tensors with new position_ids.\n\n Args:\n cos (``torch.Tensor``):\n Cosine tensor.\n sin (``torch.Tensor``):\n Sine tensor.\n position_ids (``torch.LongTensor | None``):\n Position ids.\n unsqueeze_dim (``int``, *optional*, defaults to ``1``):\n The dimension along which to unsqueeze cos and sin.\n\n Returns:\n ``tuple[torch.Tensor]``:\n Updated cos and sin tensors.\n \"\"\"\n assert unsqueeze_dim in (1, 2), f\"unsqueeze_dim must be 1 or 2, got {unsqueeze_dim}\"\n if position_ids is None:\n if cos.ndim == 2:\n cos = cos.unsqueeze(0)\n if sin.ndim == 2:\n sin = sin.unsqueeze(0)\n cos = cos.unsqueeze(unsqueeze_dim)\n sin = sin.unsqueeze(unsqueeze_dim)\n else:\n cos = cos[position_ids].unsqueeze(unsqueeze_dim) # [bs, 1, seq_len, dim] if unsqueeze_dim == 1\n sin = sin[position_ids].unsqueeze(unsqueeze_dim) # [bs, seq_len, 1, dim] if unsqueeze_dim == 2\n assert cos.ndim == 4, f\"cos must have 4 dimensions, got {cos.ndim}\"\n assert sin.ndim == 4, f\"sin must have 4 dimensions, got {sin.ndim}\"\n return cos, sin\n\n\nclass RotaryEmbedding(nn.Module):\n \"\"\"Rotary embedding for attention.\"\"\"\n\n def __init__(self) -> None:\n \"\"\"Initialize the class.\"\"\"\n super().__init__()\n\n def forward(\n self, states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, unsqueeze_dim: int = 1\n ) -> torch.Tensor:\n \"\"\"Apply rotary embedding to states.\n\n Args:\n states (torch.Tensor): States.\n cos (torch.Tensor): Cosine tensor.\n sin (torch.Tensor): Sine tensor.\n unsqueeze_dim (int, optional): The dimension along which to unsqueeze cos and sin.\n Defaults to ``1``.\n\n Returns:\n torch.Tensor: States with rotary embedding.\n \"\"\"\n states = (states * cos) + (rotate_half(states) * sin)\n if unsqueeze_dim == 1:\n batch_size, num_heads, seq_len, head_dim = states.shape\n states = states.transpose(1, 2)\n else:\n batch_size, seq_len, num_heads, head_dim = states.shape\n return states.view(batch_size, seq_len, num_heads * head_dim)\n\n\ndef apply_rotary_pos_emb(\n self,\n q: torch.Tensor,\n k: torch.Tensor,\n cos: torch.Tensor,\n sin: torch.Tensor,\n position_ids: torch.LongTensor = None,\n unsqueeze_dim: int = 1,\n) -> tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Apply Rotary Position Embedding to the query and key tensors.\n\n Args:\n q (`torch.Tensor`):\n The query tensor.\n k (`torch.Tensor`):\n The key tensor.\n cos (`torch.Tensor`):\n The cosine part of the rotary embedding.\n sin (`torch.Tensor`):\n The sine part of the rotary embedding.\n position_ids (`torch.Tensor`):\n The position indices of the tokens corresponding to the query and key tensors. For example, this can be\n used to pass offsetted position ids when working with a KV-cache.\n unsqueeze_dim (`int`, *optional*, defaults to 1):\n The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and\n sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note\n that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and\n k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes\n cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have\n the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.\n Returns:\n `tuple[torch.Tensor, torch.Tensor]`:\n comprising of the query and key tensors rotated using the Rotary Position Embedding.\n \"\"\"\n assert unsqueeze_dim == 1 or unsqueeze_dim == 2, f\"unsqueeze_dim must be 1 or 2, got {unsqueeze_dim}\"\n if unsqueeze_dim == 1:\n batch_size, _, seq_len, head_dim = q.shape\n else:\n batch_size, seq_len, _, head_dim = q.shape\n cos, sin = update_rotary_cos_sin(cos, sin, position_ids=position_ids, unsqueeze_dim=unsqueeze_dim)\n q = self.q_rotary_emb(q, cos=cos, sin=sin, unsqueeze_dim=unsqueeze_dim)\n k = self.k_rotary_emb(k, cos=cos, sin=sin, unsqueeze_dim=unsqueeze_dim)\n q = q.view(batch_size, seq_len, -1, head_dim)\n k = k.view(batch_size, seq_len, -1, head_dim)\n if unsqueeze_dim == 1:\n q = q.transpose(1, 2)\n k = k.transpose(1, 2)\n return q, k\n\n\ndef patch_attention(model: nn.Module) -> nn.Module:\n \"\"\"Patch attention.\"\"\"\n\n logger = tools.logging.getLogger(f\"{__name__}.ModelPatcher\")\n for module_name, module in model.named_modules():\n classname = type(module).__name__\n if classname.lower().endswith(\"attention\"):\n forward_name = \"\"\n if isinstance(module.forward, functools.partial):\n if hasattr(module, \"_deepcompressor_orig_forward\"):\n logger.info(f\"- Attention in {module_name} has already been patched\")\n else:\n # this module has been wrapped in ``accelerate`` package\n assert hasattr(module, \"_old_forward\")\n assert module._old_forward is module.forward.__wrapped__\n if \"apply_rotary_pos_emb\" in module._old_forward.__func__.__globals__:\n forward_name = \"_old_forward\"\n else:\n if \"apply_rotary_pos_emb\" in module.forward.__func__.__globals__:\n forward_name = \"forward\"\n if forward_name:\n logger.info(f\"- Patching {classname}.{forward_name} in {module_name}\")\n module.q_rotary_emb = RotaryEmbedding()\n module.k_rotary_emb = RotaryEmbedding()\n module.apply_rotary_pos_emb = functools.partial(apply_rotary_pos_emb, module)\n module._deepcompressor_orig_forward = getattr(module, forward_name)\n\n orig_forward = module._deepcompressor_orig_forward.__func__\n new_globals = dict(orig_forward.__globals__)\n new_globals[\"apply_rotary_pos_emb\"] = module.apply_rotary_pos_emb\n new_forward = copy_func(orig_forward, new_globals)\n setattr(module, forward_name, new_forward.__get__(module))\n\n return model\n\n\ndef gemma_rms_norm_forward(self: GemmaRMSNorm | Gemma2RMSNorm, x: torch.Tensor) -> torch.Tensor:\n \"\"\"Forward function for Gemma RMSNorm.\"\"\"\n assert hasattr(self, \"_deepcompressor_orig_forward\"), \"Gemma RMSNorm must be patched before calling forward\"\n output = self._norm(x.float())\n # Llama does x.to(float16) * w whilst Gemma2 is (x * w).to(float16)\n # See https://github.com/huggingface/transformers/pull/29402\n output = output * self.weight.float()\n return output.type_as(x)\n\n\ndef patch_gemma_rms_norm(model: nn.Module) -> nn.Module:\n \"\"\"Patch Gemma RMSNorm.\"\"\"\n\n logger = tools.logging.getLogger(f\"{__name__}.ModelPatcher\")\n for module_name, module in model.named_modules():\n if isinstance(module, (GemmaRMSNorm, Gemma2RMSNorm)):\n classname = type(module).__name__\n forward_name = \"\"\n if hasattr(module, \"_deepcompressor_orig_forward\"):\n logger.info(f\"- {module_name} has already been patched\")\n else:\n if isinstance(module.forward, functools.partial):\n assert hasattr(module, \"_old_forward\")\n assert module._old_forward is module.forward.__wrapped__\n forward_name = \"_old_forward\"\n else:\n forward_name = \"forward\"\n if forward_name:\n logger.info(f\"- Patching {classname}.{forward_name} in {module_name}\")\n module.weight.data.add_(1.0)\n module._deepcompressor_orig_forward = getattr(module, forward_name)\n setattr(module, forward_name, functools.partial(gemma_rms_norm_forward, module))\n return model\n"
},
{
"path": "deepcompressor/app/llm/nn/struct.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Utility functions for Large Language Models.\"\"\"\n\n# region imports\nimport typing as tp\nfrom dataclasses import dataclass, field\n\nimport torch.nn as nn\nfrom transformers import PreTrainedModel\nfrom transformers.models.gemma2.modeling_gemma2 import (\n Gemma2Attention,\n Gemma2Config,\n Gemma2DecoderLayer,\n Gemma2ForCausalLM,\n Gemma2ForSequenceClassification,\n Gemma2MLP,\n Gemma2Model,\n)\nfrom transformers.models.llama.modeling_llama import (\n LlamaAttention,\n LlamaConfig,\n LlamaDecoderLayer,\n LlamaForCausalLM,\n LlamaForSequenceClassification,\n LlamaMLP,\n LlamaModel,\n)\nfrom transformers.models.mistral.modeling_mistral import (\n MistralAttention,\n MistralConfig,\n MistralDecoderLayer,\n MistralForCausalLM,\n MistralForSequenceClassification,\n MistralMLP,\n MistralModel,\n)\nfrom transformers.models.mixtral.modeling_mixtral import (\n MixtralAttention,\n MixtralConfig,\n MixtralDecoderLayer,\n MixtralForCausalLM,\n MixtralForSequenceClassification,\n MixtralModel,\n MixtralSparseMoeBlock,\n)\nfrom transformers.models.qwen2.modeling_qwen2 import (\n Qwen2Attention,\n Qwen2Config,\n Qwen2DecoderLayer,\n Qwen2ForCausalLM,\n Qwen2ForSequenceClassification,\n Qwen2MLP,\n Qwen2Model,\n)\nfrom transformers.models.t5.modeling_t5 import (\n T5Attention,\n T5Block,\n T5Config,\n T5DenseActDense,\n T5DenseGatedActDense,\n T5EncoderModel,\n T5LayerFF,\n T5LayerSelfAttention,\n T5Stack,\n)\n\nfrom deepcompressor.nn.struct.attn import (\n AttentionConfigStruct,\n BaseTransformerStruct,\n FeedForwardConfigStruct,\n FeedForwardStruct,\n SelfAttentionStruct,\n TransformerBlockStruct,\n)\nfrom deepcompressor.nn.struct.base import BaseModuleStruct\nfrom deepcompressor.utils.common import join_name\n\nfrom .patch import RotaryEmbedding\n\n# endregion\n\n\n__all__ = [\n \"LlmConfigStruct\",\n \"LlmModelStruct\",\n \"LlmTransformerStruct\",\n \"LlmTransformerBlockStruct\",\n \"LlmSelfAttentionStruct\",\n \"LlmFeedForwardStruct\",\n]\n\n# region type aliases\nATTENTION_CLS = tp.Union[\n LlamaAttention, MistralAttention, MixtralAttention, Qwen2Attention, T5Attention, Gemma2Attention\n]\nFEEDFORWARD_CLS = tp.Union[\n LlamaMLP, MistralMLP, MixtralSparseMoeBlock, Qwen2MLP, T5DenseActDense, T5DenseGatedActDense, Gemma2MLP\n]\nTRANSFORMER_BLOCK_CLS = tp.Union[\n LlamaDecoderLayer, MistralDecoderLayer, MixtralDecoderLayer, Qwen2DecoderLayer, T5Block, Gemma2DecoderLayer\n]\nTRANSFORMER_CLS = tp.Union[LlamaModel, MistralModel, MixtralModel, Qwen2Model, T5Stack, Gemma2Model]\nCASUALLM_CLS = tp.Union[LlamaForCausalLM, MistralForCausalLM, MixtralForCausalLM, Qwen2ForCausalLM, Gemma2ForCausalLM]\nSEQCLSLM_CLS = tp.Union[\n LlamaForSequenceClassification,\n MistralForSequenceClassification,\n MixtralForSequenceClassification,\n Qwen2ForSequenceClassification,\n Gemma2ForSequenceClassification,\n]\n# endregion\n\n\n@dataclass(kw_only=True)\nclass LlmTransformerBlockConfigStruct(FeedForwardConfigStruct, AttentionConfigStruct):\n \"\"\"Large Language Model Transformer Block Configuration.\n\n Args:\n hidden_size (`int`):\n The size of the input/output activations, i.e., the number of input channels.\n inner_size (`int`):\n The size of the inner activations, i.e., the number of **query** channels in the attention block.\n intermediate_size (`int`):\n The number of intermediate channels in the feedforward network.\n intermediate_act_type (`str`):\n The activation function for the intermediate activations in the feedforward network.\n num_query_heads (`int`):\n The number of query heads.\n num_key_value_heads (`int`):\n The number of key-value heads.\n num_experts (`int`):\n The number of experts (for the feedforward network).\n with_qk_norm (`bool`, *optional*, defaults to `False`):\n Whether to apply normalization to queries and keys.\n with_rope (`bool`):\n Whether to use Rotary Positional Encoding (RoPE).\n\n Attributes:\n head_size (`int`):\n The size of the head, equal to `num_query_channels // num_query_heads`.\n num_key_value_groups (`int`):\n The number of key-value groups, equal to `num_query_heads // num_key_value_heads`.\n intermediate_lowerbound (`float` or `None`):\n The lowerbound of the intermediate activations in feedforward network.\n \"\"\"\n\n pass\n\n\n@dataclass(kw_only=True)\nclass LlmTransformerConfigStruct(LlmTransformerBlockConfigStruct):\n \"\"\"Large Language Model Transformer Configuration.\n\n Args:\n hidden_size (`int`):\n The size of the input/output activations, i.e., the number of input channels.\n inner_size (`int`):\n The size of the inner activations, i.e., the number of **query** channels in the attention block.\n intermediate_size (`int`):\n The number of intermediate channels in the feedforward network.\n intermediate_act_type (`str`):\n The activation function for the intermediate activations in the feedforward network.\n num_query_heads (`int`):\n The number of query heads.\n num_key_value_heads (`int`):\n The number of key-value heads.\n num_experts (`int`):\n The number of experts (for the feedforward network).\n with_qk_norm (`bool`, *optional*, defaults to `False`):\n Whether to apply normalization to queries and keys.\n with_rope (`bool`):\n Whether to use Rotary Positional Encoding (RoPE).\n vocab_size (`int`):\n The size of the vocabulary.\n num_hidden_layers (`int`):\n The number of hidden layers.\n\n Attributes:\n head_size (`int`):\n The size of the head, equal to `num_query_channels // num_query_heads`.\n num_key_value_groups (`int`):\n The number of key-value groups, equal to `num_query_heads // num_key_value_heads`.\n intermediate_lowerbound (`float` or `None`):\n The lowerbound of the intermediate activations in feedforward network.\n \"\"\"\n\n vocab_size: int\n num_hidden_layers: int\n\n\n@dataclass(kw_only=True)\nclass LlmConfigStruct(LlmTransformerConfigStruct):\n \"\"\"Large Language Model Configuration.\n\n Args:\n hidden_size (`int`):\n The size of the input/output activations, i.e., the number of input channels.\n inner_size (`int`):\n The size of the inner activations, i.e., the number of **query** channels in the attention block.\n intermediate_size (`int`):\n The number of intermediate channels in the feedforward network.\n intermediate_act_type (`str`):\n The activation function for the intermediate activations in the feedforward network.\n num_query_heads (`int`):\n The number of query heads.\n num_key_value_heads (`int`):\n The number of key-value heads.\n num_experts (`int`):\n The number of experts (for the feedforward network).\n with_qk_norm (`bool`, *optional*, defaults to `False`):\n Whether to apply normalization to queries and keys.\n with_rope (`bool`):\n Whether to use Rotary Positional Encoding (RoPE).\n vocab_size (`int`):\n The size of the vocabulary.\n num_hidden_layers (`int`):\n The number of hidden layers.\n tie_word_embeddings (`bool`):\n Whether to tie the word embeddings with the head weights.\n\n Attributes:\n head_size (`int`):\n The size of the head, equal to `num_query_channels // num_query_heads`.\n num_key_value_groups (`int`):\n The number of key-value groups, equal to `num_query_heads // num_key_value_heads`.\n intermediate_lowerbound (`float` or `None`):\n The lowerbound of the intermediate activations in feedforward network.\n \"\"\"\n\n tie_word_embeddings: bool = False\n\n\n@dataclass(kw_only=True)\nclass LlmSelfAttentionStruct(SelfAttentionStruct):\n \"\"\"Large Language Model Attention Block.\"\"\"\n\n # region relative keys\n q_rkey: tp.ClassVar[str] = \"attn_q\"\n k_rkey: tp.ClassVar[str] = \"attn_k\"\n v_rkey: tp.ClassVar[str] = \"attn_v\"\n # endregion\n\n parent: tp.Optional[\"LlmTransformerBlockStruct\"] = field(repr=False)\n kwargs: tuple[str, ...]\n\n def filter_kwargs(self, kwargs: dict) -> dict:\n \"\"\"Filter layer kwargs to attn kwargs.\"\"\"\n return {k: v for k, v in kwargs.items() if k in self.kwargs}\n\n @staticmethod\n def _default_construct(\n module: ATTENTION_CLS,\n /,\n parent: tp.Optional[\"LlmTransformerBlockStruct\"] = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"LlmSelfAttentionStruct\":\n if isinstance(module, T5Attention):\n with_rope, num_query_heads, num_key_value_heads = False, module.n_heads, module.n_heads\n q_proj, k_proj, v_proj, o_proj = module.q, module.k, module.v, module.o\n q_proj_rname, k_proj_rname, v_proj_rname, o_proj_rname = \"q\", \"k\", \"v\", \"o\"\n q, k, v = module.q, module.k, module.v\n q_rname, k_rname, v_rname = \"q\", \"k\", \"v\"\n kwargs = (\n \"mask\",\n \"key_value_states\",\n \"position_bias\",\n \"past_key_value\",\n \"layer_head_mask\",\n \"query_length\",\n \"use_cache\",\n \"output_attentions\",\n )\n elif isinstance(module, (LlamaAttention, MistralAttention, MixtralAttention, Qwen2Attention, Gemma2Attention)):\n with_rope = True\n num_query_heads = module.config.num_attention_heads\n num_key_value_heads = module.config.num_key_value_heads\n q_proj, k_proj, v_proj, o_proj = module.q_proj, module.k_proj, module.v_proj, module.o_proj\n q_proj_rname, k_proj_rname, v_proj_rname, o_proj_rname = \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\"\n if hasattr(module, \"q_rotary_emb\"):\n q, k = module.q_rotary_emb, module.k_rotary_emb\n q_rname, k_rname = \"q_rotary_emb\", \"k_rotary_emb\"\n assert isinstance(q, RotaryEmbedding)\n assert isinstance(k, RotaryEmbedding)\n else:\n q, k = module.q_proj, module.k_proj\n q_rname, k_rname = \"q_proj\", \"k_proj\"\n v, v_rname = module.v_proj, \"v_proj\"\n kwargs = (\n \"attention_mask\",\n \"position_ids\",\n \"past_key_value\",\n \"output_attentions\",\n \"use_cache\",\n \"position_embeddings\",\n \"cache_position\",\n )\n else:\n raise ValueError(f\"Unsupported attention type: {type(module)}\")\n config = AttentionConfigStruct(\n hidden_size=q_proj.weight.shape[1],\n inner_size=q_proj.weight.shape[0],\n num_query_heads=num_query_heads,\n num_key_value_heads=num_key_value_heads,\n with_qk_norm=False,\n with_rope=with_rope,\n )\n if parent is not None and parent.config is not None:\n assert parent.config.hidden_size == config.hidden_size\n assert parent.config.inner_size == config.inner_size\n assert parent.config.num_query_heads == config.num_query_heads\n assert parent.config.num_key_value_heads == config.num_key_value_heads\n assert parent.config.with_qk_norm == config.with_qk_norm\n assert parent.config.with_rope == config.with_rope\n return LlmSelfAttentionStruct(\n module=module,\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n config=config,\n q_proj=q_proj,\n k_proj=k_proj,\n v_proj=v_proj,\n o_proj=o_proj,\n q=q,\n k=k,\n v=v,\n q_proj_rname=q_proj_rname,\n k_proj_rname=k_proj_rname,\n v_proj_rname=v_proj_rname,\n o_proj_rname=o_proj_rname,\n q_rname=q_rname,\n k_rname=k_rname,\n v_rname=v_rname,\n kwargs=kwargs,\n )\n\n\n@dataclass(kw_only=True)\nclass LlmFeedForwardStruct(FeedForwardStruct):\n \"\"\"Large Language Model Feedforward Network.\"\"\"\n\n parent: tp.Optional[\"LlmTransformerBlockStruct\"] = field(repr=False)\n\n @staticmethod\n def _default_construct(\n module: FEEDFORWARD_CLS,\n /,\n parent: tp.Optional[\"LlmTransformerBlockStruct\"] = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"LlmFeedForwardStruct\":\n if isinstance(module, (LlamaMLP, MistralMLP, Qwen2MLP, Gemma2MLP)):\n if parent is not None:\n assert parent.config.intermediate_act_type.endswith(\"_glu\")\n act_type = parent.config.intermediate_act_type\n else:\n act_type = str(module.act_fn.__class__.__name__).removesuffix(\"activation\").lower() + \"_glu\"\n up_projs, down_projs = [module.up_proj, module.gate_proj], [module.down_proj]\n experts = [module]\n moe_gate = None\n up_proj_rnames = [\"up_proj\", \"gate_proj\"]\n down_proj_rnames = [\"down_proj\"]\n experts_rname = \"\"\n moe_gate_rname = \"\"\n elif isinstance(module, MixtralSparseMoeBlock):\n if parent is not None:\n assert parent.config.intermediate_act_type.endswith(\"_glu\")\n act_type = parent.config.intermediate_act_type\n else:\n act_type = str(module.experts[0].act_fn.__class__.__name__).removesuffix(\"activation\").lower() + \"_glu\"\n up_projs = [expert.w3 for expert in module.experts] + [expert.w1 for expert in module.experts]\n down_projs = [expert.w2 for expert in module.experts]\n experts = list(module.experts)\n moe_gate = module.gate\n up_proj_rnames = [\"w3\", \"w1\"]\n down_proj_rnames = [\"w2\"]\n experts_rname = \"experts\"\n moe_gate_rname = \"gate\"\n elif isinstance(module, T5DenseActDense):\n if parent is not None:\n assert not parent.config.intermediate_act_type.endswith(\"_glu\")\n act_type = parent.config.intermediate_act_type\n else:\n act_type = str(module.act.__class__.__name__).removesuffix(\"activation\").lower()\n up_projs = [module.wi]\n down_projs = [module.wo]\n experts = [module]\n moe_gate = None\n up_proj_rnames = [\"wi\"]\n down_proj_rnames = [\"wo\"]\n experts_rname = \"\"\n moe_gate_rname = \"\"\n elif isinstance(module, T5DenseGatedActDense):\n if parent is not None:\n assert parent.config.intermediate_act_type.endswith(\"_glu\")\n act_type = parent.config.intermediate_act_type\n else:\n act_type = str(module.act.__class__.__name__).removesuffix(\"activation\").lower() + \"_glu\"\n up_projs = [module.wi_1, module.wi_0]\n down_projs = [module.wo]\n experts = [module]\n moe_gate = None\n up_proj_rnames = [\"wi_1\", \"wi_0\"]\n down_proj_rnames = [\"wo\"]\n experts_rname = \"\"\n moe_gate_rname = \"\"\n else:\n raise ValueError(f\"Unsupported feed forward network type: {type(module)}\")\n config = FeedForwardConfigStruct(\n hidden_size=up_projs[0].weight.shape[1],\n intermediate_size=up_projs[0].weight.shape[0],\n intermediate_act_type=act_type,\n num_experts=len(experts),\n )\n if parent is not None and parent.config is not None:\n assert parent.config.hidden_size == config.hidden_size\n assert parent.config.intermediate_size == config.intermediate_size\n assert parent.config.intermediate_act_type == config.intermediate_act_type\n assert parent.config.num_experts == config.num_experts\n return LlmFeedForwardStruct(\n module=module,\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n config=config,\n up_projs=up_projs,\n down_projs=down_projs,\n moe_gate=moe_gate,\n experts=experts,\n up_proj_rnames=up_proj_rnames,\n down_proj_rnames=down_proj_rnames,\n moe_gate_rname=moe_gate_rname,\n experts_rname=experts_rname,\n )\n\n\n@dataclass(kw_only=True)\nclass LlmTransformerBlockStruct(TransformerBlockStruct):\n \"\"\"Large Language Model Transformer Block.\"\"\"\n\n # region relative keys\n attn_rkey: tp.ClassVar[str] = \"\"\n ffn_rkey: tp.ClassVar[str] = \"\"\n add_ffn_rkey: tp.ClassVar[str] = \"add\"\n attn_struct_cls: tp.ClassVar[tp.Type[LlmSelfAttentionStruct]] = LlmSelfAttentionStruct\n ffn_struct_cls: tp.ClassVar[tp.Type[LlmFeedForwardStruct]] = LlmFeedForwardStruct\n # endregion\n\n parent: tp.Optional[\"LlmTransformerStruct\"] = field(repr=False)\n parallel: bool = field(init=False, repr=False, default=False)\n config: LlmTransformerBlockConfigStruct = field(default=None)\n\n # region child modules\n pre_attn_add_norms: list[nn.LayerNorm] = field(init=False, repr=False, default_factory=list)\n post_attn_add_norms: list[nn.LayerNorm] = field(init=False, repr=False, default_factory=list)\n pre_add_ffn_norm: None = field(init=False, repr=False, default=None)\n add_ffn: None = field(init=False, repr=False, default=None)\n post_add_ffn_norm: None = field(init=False, repr=False, default=None)\n # endregion\n # region relative names\n pre_attn_add_norm_rnames: list[str] = field(init=False, repr=False, default_factory=list)\n post_attn_add_norm_rnames: list[str] = field(init=False, repr=False, default_factory=list)\n pre_add_ffn_norm_rname: str = field(init=False, repr=False, default=\"\")\n add_ffn_rname: str = field(init=False, repr=False, default=\"\")\n post_add_ffn_norm_rname: str = field(init=False, repr=False, default=\"\")\n # endregion\n # region child structs\n attn_structs: list[LlmSelfAttentionStruct] = field(init=False, repr=False)\n ffn_struct: LlmFeedForwardStruct = field(init=False, repr=False)\n add_ffn_struct: None = field(init=False, repr=False, default=None)\n # endregion\n\n # region aliases\n\n @property\n def pre_attn_norm(self) -> nn.LayerNorm | None:\n return self.pre_attn_norms[0] if self.pre_attn_norms else None\n\n @property\n def attn(self) -> nn.Module:\n return self.attns[0]\n\n @property\n def post_attn_norm(self) -> nn.LayerNorm | None:\n return self.post_attn_norms[0] if self.post_attn_norms else None\n\n @property\n def pre_attn_norm_rname(self) -> str:\n return self.pre_attn_norm_rnames[0] if self.pre_attn_norm_rnames else \"\"\n\n @property\n def attn_rname(self) -> str:\n return self.attn_rnames[0]\n\n @property\n def post_attn_norm_rname(self) -> str:\n return self.post_attn_norm_rnames[0] if self.post_attn_norm_rnames else \"\"\n\n @property\n def pre_attn_norm_name(self) -> str:\n return self.pre_attn_norm_names[0] if self.pre_attn_norm_names else \"\"\n\n @property\n def attn_name(self) -> str:\n return self.attn_names[0]\n\n @property\n def post_attn_norm_name(self) -> str:\n return self.post_attn_norm_names[0] if self.post_attn_norm_names else \"\"\n\n @property\n def attn_struct(self) -> LlmSelfAttentionStruct:\n return self.attn_structs[0]\n\n # endregion\n\n def __post_init__(self):\n super().__post_init__()\n assert len(self.attn_structs) == 1\n if self.config is None:\n self.config = LlmTransformerBlockConfigStruct(\n hidden_size=self.attn_struct.config.hidden_size,\n inner_size=self.attn_struct.config.inner_size,\n num_query_heads=self.attn_struct.config.num_query_heads,\n num_key_value_heads=self.attn_struct.config.num_key_value_heads,\n with_qk_norm=self.attn_struct.config.with_qk_norm,\n with_rope=self.attn_struct.config.with_rope,\n intermediate_size=self.ffn_struct.config.intermediate_size,\n intermediate_act_type=self.ffn_struct.config.intermediate_act_type,\n num_experts=self.ffn_struct.config.num_experts,\n )\n\n @staticmethod\n def _default_construct(\n module: TRANSFORMER_BLOCK_CLS,\n /,\n parent: tp.Optional[\"LlmTransformerStruct\"] = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"LlmTransformerBlockStruct\":\n if isinstance(\n module, (LlamaDecoderLayer, MistralDecoderLayer, Qwen2DecoderLayer, MixtralDecoderLayer, Gemma2DecoderLayer)\n ):\n pre_attn_norms, attns = [module.input_layernorm], [module.self_attn]\n pre_attn_norm_rnames, attn_rnames = [\"input_layernorm\"], [\"self_attn\"]\n if isinstance(module, Gemma2DecoderLayer):\n post_attn_norms, post_attn_norm_rnames = [module.post_attention_layernorm], [\"post_attention_layernorm\"]\n pre_ffn_norm, pre_ffn_norm_rname = (module.pre_feedforward_layernorm, \"pre_feedforward_layernorm\")\n post_ffn_norm, post_ffn_norm_rname = module.post_feedforward_layernorm, \"post_feedforward_layernorm\"\n else:\n post_attn_norms, post_attn_norm_rnames = [], []\n pre_ffn_norm, pre_ffn_norm_rname = module.post_attention_layernorm, \"post_attention_layernorm\"\n post_ffn_norm, post_ffn_norm_rname = None, \"\"\n if isinstance(module, MixtralDecoderLayer):\n ffn, ffn_rname = module.block_sparse_moe, \"block_sparse_moe\"\n else:\n ffn, ffn_rname = module.mlp, \"mlp\"\n elif isinstance(module, T5Block):\n pre_attn_norms, attns, pre_attn_norm_rnames, attn_rnames = [], [], [], []\n post_attn_norms, post_attn_norm_rnames = [], []\n post_ffn_norm, post_ffn_norm_rname = None, \"\"\n for i, layer in enumerate(module.layer):\n if isinstance(layer, T5LayerSelfAttention):\n pre_attn_norms.append(layer.layer_norm)\n attns.append(layer.SelfAttention)\n pre_attn_norm_rnames.append(f\"layer.{i}.layer_norm\")\n attn_rnames.append(f\"layer.{i}.SelfAttention\")\n else:\n assert isinstance(layer, T5LayerFF)\n pre_ffn_norm, ffn = layer.layer_norm, layer.DenseReluDense\n pre_ffn_norm_rname, ffn_rname = f\"layer.{i}.layer_norm\", f\"layer.{i}.DenseReluDense\"\n else:\n raise ValueError(f\"Unsupported layer type: {type(module)}\")\n config = parent.config if parent is not None and parent.config is not None else None\n return LlmTransformerBlockStruct(\n module=module,\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n config=config,\n pre_attn_norms=pre_attn_norms,\n attns=attns,\n post_attn_norms=post_attn_norms,\n pre_ffn_norm=pre_ffn_norm,\n ffn=ffn,\n post_ffn_norm=post_ffn_norm,\n pre_attn_norm_rnames=pre_attn_norm_rnames,\n attn_rnames=attn_rnames,\n post_attn_norm_rnames=post_attn_norm_rnames,\n pre_ffn_norm_rname=pre_ffn_norm_rname,\n ffn_rname=ffn_rname,\n post_ffn_norm_rname=post_ffn_norm_rname,\n )\n\n\n@dataclass(kw_only=True)\nclass LlmTransformerStruct(BaseTransformerStruct):\n \"\"\"Large Language Model Structure.\"\"\"\n\n # region relative keys\n layer_rkey: tp.ClassVar[str] = \"\"\n layer_struct_cls: tp.ClassVar[tp.Type[LlmTransformerBlockStruct]] = LlmTransformerBlockStruct\n # endregion\n\n parent: tp.Optional[\"LlmModelStruct\"] = field(repr=False)\n config: LlmTransformerConfigStruct = field(default=None)\n\n # region child modules\n # embeddings: list[nn.Embedding]\n # \"\"\"list of embeddings [embed_tokens, embed_positions]\"\"\"\n embed_tokens: nn.Embedding\n \"\"\"Token embedding module.\"\"\"\n embed_positions: nn.Embedding | None\n \"\"\"Position embedding module.\"\"\"\n layers: nn.ModuleList\n # endregion\n # region relative names\n embed_tokens_rname: str\n embed_positions_rname: str\n layers_rname: str\n # endregion\n # region absolute names\n embed_tokens_name: str = field(init=False, repr=False)\n embed_positions_name: str = field(init=False, repr=False)\n layers_name: str = field(init=False, repr=False)\n layer_names: list[str] = field(init=False, repr=False)\n # endregion\n # region child structs\n layer_structs: list[LlmTransformerBlockStruct] = field(init=False, repr=False)\n # endregion\n\n # region abstractmethod implementations\n\n @property\n def num_blocks(self) -> int:\n \"\"\"Get the number of transformer blocks.\"\"\"\n return len(self.layers)\n\n @property\n def block_structs(self) -> list[LlmTransformerBlockStruct]:\n return self.layer_structs\n\n @property\n def block_names(self) -> list[str]:\n \"\"\"Get the list of transformer block names.\"\"\"\n return self.layer_names\n\n # endregion\n\n def __post_init__(self) -> None:\n super().__post_init__()\n self.embed_tokens_name = join_name(self.name, self.embed_tokens_rname)\n if self.embed_positions is not None:\n self.embed_positions_name = join_name(self.name, self.embed_positions_rname)\n else:\n self.embed_positions_name = \"\"\n self.layers_name = join_name(self.name, self.layers_rname)\n layer_rnames = [f\"{self.layers_rname}.{idx}\" for idx in range(len(self.layers))]\n self.layer_names = [join_name(self.name, rname) for rname in layer_rnames]\n self.layer_structs = [\n self.layer_struct_cls.construct(\n layer, parent=self, fname=\"layer\", rname=rname, rkey=self.layer_rkey, idx=idx\n )\n for idx, (layer, rname) in enumerate(zip(self.layers, layer_rnames, strict=True))\n ]\n if self.config is None:\n assert all(block.config == self.block_structs[0].config for block in self.block_structs)\n ref_config = self.block_structs[0].config\n self.config = LlmTransformerConfigStruct(\n hidden_size=ref_config.hidden_size,\n inner_size=ref_config.inner_size,\n num_query_heads=ref_config.num_query_heads,\n num_key_value_heads=ref_config.num_key_value_heads,\n with_qk_norm=ref_config.with_qk_norm,\n with_rope=ref_config.with_rope,\n intermediate_size=ref_config.intermediate_size,\n intermediate_act_type=ref_config.intermediate_act_type,\n num_experts=ref_config.num_experts,\n vocab_size=self.embed_tokens.num_embeddings,\n num_hidden_layers=self.num_blocks,\n )\n else:\n assert self.config.vocab_size == self.embed_tokens.num_embeddings\n assert self.config.num_hidden_layers == self.num_blocks\n\n def get_iter_layer_activations_args(\n self, **kwargs\n ) -> tuple[list[nn.Module], list[LlmTransformerBlockStruct], list[bool], list[bool]]:\n \"\"\"\n Get the arguments for iterating over the layers and their activations.\n\n Args:\n skip_pre_modules (`bool`):\n Whether to skip the pre-modules\n skip_post_modules (`bool`):\n Whether to skip the post-modules\n\n Returns:\n `tuple[list[nn.Module], list[LlmTransformerBlockStruct], list[bool], list[bool]]`:\n the layers, the layer structs, the recomputes, and the use_prev_layer_outputs\n \"\"\"\n return self.layers, self.layer_structs, [False] * len(self.layers), [True] * len(self.layers)\n\n @staticmethod\n def _default_construct(\n module: TRANSFORMER_CLS,\n /,\n parent: tp.Optional[\"LlmModelStruct\"] = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"LlmTransformerStruct\":\n if isinstance(module, (LlamaModel, MistralModel, MixtralModel, Qwen2Model, Gemma2Model)):\n embed_tokens, embed_positions = module.embed_tokens, None\n layers = module.layers\n norm_in, norm_out = None, module.norm\n proj_in, proj_out = None, None\n embed_tokens_rname, embed_positions_rname = \"embed_tokens\", \"\"\n layers_rname = \"layers\"\n norm_in_rname, norm_out_rname = \"\", \"norm\"\n proj_in_rname, proj_out_rname = \"\", \"\"\n elif isinstance(module, T5Stack):\n embed_tokens, embed_positions = module.embed_tokens, None\n layers = module.block\n norm_in, norm_out = None, module.final_layer_norm\n proj_in, proj_out = None, None\n embed_tokens_rname, embed_positions_rname = \"embed_tokens\", \"\"\n layers_rname = \"block\"\n norm_in_rname, norm_out_rname = \"\", \"final_layer_norm\"\n proj_in_rname, proj_out_rname = \"\", \"\"\n else:\n raise ValueError(f\"Unsupported backbone type: {type(module)}\")\n config = parent.config if parent is not None and parent.config is not None else None\n return LlmTransformerStruct(\n module=module,\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n config=config,\n embed_tokens=embed_tokens,\n embed_positions=embed_positions,\n norm_in=norm_in,\n proj_in=proj_in,\n layers=layers,\n norm_out=norm_out,\n proj_out=proj_out,\n embed_tokens_rname=embed_tokens_rname,\n embed_positions_rname=embed_positions_rname,\n norm_in_rname=norm_in_rname,\n proj_in_rname=proj_in_rname,\n layers_rname=layers_rname,\n norm_out_rname=norm_out_rname,\n proj_out_rname=proj_out_rname,\n )\n\n\n@dataclass(kw_only=True)\nclass LlmModelStruct(BaseModuleStruct):\n \"\"\"Large Language Model Structure.\"\"\"\n\n # region relative keys\n backbone_rkey: tp.ClassVar[str] = \"\"\n head_rkey: tp.ClassVar[str] = \"head\"\n backbone_struct_cls: tp.ClassVar[tp.Type[LlmTransformerStruct]] = LlmTransformerStruct\n # endregion\n\n module: PreTrainedModel = field(repr=False, kw_only=False)\n config: LlmConfigStruct\n # region child modules\n backbone: nn.Module\n head: nn.Linear | None\n # endregion\n # region relative names\n backbone_rname: str\n head_rname: str\n # endregion\n # region absolute names\n backbone_name: str = field(init=False, repr=False)\n head_name: str = field(init=False, repr=False)\n # endregion\n # region absolute keys\n head_key: str = field(init=False, repr=False)\n # endregion\n # region child structs\n backbone_struct: LlmTransformerStruct = field(init=False, repr=False)\n # endregion\n\n def __post_init__(self) -> None:\n super().__post_init__()\n self.backbone_name = join_name(self.name, self.backbone_rname)\n if self.head is not None or self.head_rname:\n self.head_name = join_name(self.name, self.head_rname)\n else:\n self.head_name = self.head_rname = \"\"\n self.head_key = join_name(self.key, self.head_rkey, sep=\"_\")\n self.backbone_struct = self.backbone_struct_cls.construct(\n self.backbone, parent=self, fname=\"backbone\", rname=self.backbone_rname, rkey=self.backbone_rkey\n )\n\n def named_key_modules(self) -> tp.Generator[tp.Tuple[str, str, nn.Module, BaseModuleStruct, str], None, None]:\n yield from self.backbone_struct.named_key_modules()\n if self.head is not None:\n yield self.head_key, self.head_name, self.head, self, \"head\"\n\n def iter_attention_structs(self) -> tp.Generator[LlmSelfAttentionStruct, None, None]:\n yield from self.backbone_struct.iter_attention_structs()\n\n def iter_transformer_block_structs(self) -> tp.Generator[LlmTransformerBlockStruct, None, None]:\n yield from self.backbone_struct.iter_transformer_block_structs()\n\n def get_iter_layer_activations_args(\n self, **kwargs\n ) -> tuple[list[nn.Module], list[LlmTransformerBlockStruct], list[bool], list[bool]]:\n \"\"\"\n Get the arguments for iterating over the layers and their activations.\n\n Args:\n skip_pre_modules (`bool`):\n Whether to skip the pre-modules\n skip_post_modules (`bool`):\n Whether to skip the post-modules\n\n Returns:\n `tuple[list[nn.Module], list[LlmTransformerBlockStruct], list[bool], list[bool]]`:\n the layers, the layer structs, the recomputes, and the use_prev_layer_outputs\n \"\"\"\n return self.backbone_struct.get_iter_layer_activations_args(**kwargs)\n\n @staticmethod\n def _default_construct(\n model: nn.Module,\n /,\n parent: tp.Optional[BaseModuleStruct] = None,\n fname: str = \"\",\n rname: str = \"\",\n rkey: str = \"\",\n idx: int = 0,\n **kwargs,\n ) -> \"LlmModelStruct\":\n \"\"\"Build the Large Language Model Structure.\"\"\"\n if isinstance(model, CASUALLM_CLS) or isinstance(model, SEQCLSLM_CLS):\n backbone = model.model\n backbone_rname = \"model\"\n elif isinstance(model, T5EncoderModel):\n backbone = model.encoder\n backbone_rname = \"encoder\"\n elif isinstance(model, TRANSFORMER_CLS):\n backbone = model\n backbone_rname = \"\"\n else:\n raise ValueError(f\"Unsupported model type: {type(model)}\")\n if isinstance(model, CASUALLM_CLS):\n head = model.lm_head\n head_rname = \"lm_head\"\n elif isinstance(model, SEQCLSLM_CLS):\n head = model.score\n head_rname = \"score\"\n elif isinstance(model, T5EncoderModel):\n head = None\n head_rname = \"\"\n elif isinstance(model, TRANSFORMER_CLS):\n head = None\n head_rname = \"\"\n else:\n raise ValueError(f\"Unsupported model type: {type(model)}\")\n config = backbone.config\n if isinstance(config, (LlamaConfig, MistralConfig, MixtralConfig, Qwen2Config, Gemma2Config)):\n config_struct = LlmConfigStruct(\n hidden_size=config.hidden_size,\n inner_size=config.num_attention_heads * config.head_dim\n if isinstance(config, Gemma2Config)\n else config.hidden_size,\n num_query_heads=config.num_attention_heads,\n num_key_value_heads=config.num_key_value_heads,\n with_qk_norm=False,\n with_rope=True,\n intermediate_size=config.intermediate_size,\n intermediate_act_type=f\"{config.hidden_act}_glu\".lower(),\n num_experts=getattr(config, \"num_local_experts\", 1),\n vocab_size=config.vocab_size,\n num_hidden_layers=config.num_hidden_layers,\n tie_word_embeddings=config.tie_word_embeddings,\n )\n elif isinstance(config, T5Config):\n config_struct = LlmConfigStruct(\n hidden_size=config.d_model,\n inner_size=config.d_kv * config.num_heads,\n num_query_heads=config.num_heads,\n num_key_value_heads=config.num_heads,\n with_rope=False,\n intermediate_size=config.d_ff,\n intermediate_act_type=config.dense_act_fn.lower(),\n num_experts=1,\n vocab_size=config.vocab_size,\n num_hidden_layers=config.num_layers,\n tie_word_embeddings=False,\n )\n if config.is_gated_act:\n config_struct.intermediate_act_type += \"_glu\"\n else:\n raise ValueError(f\"Unsupported config type: {type(config)}\")\n return LlmModelStruct(\n module=model,\n parent=parent,\n fname=fname,\n idx=idx,\n rname=rname,\n rkey=rkey,\n config=config_struct,\n backbone=backbone,\n head=head,\n backbone_rname=backbone_rname,\n head_rname=head_rname,\n )\n\n\nLlmSelfAttentionStruct.register_factory(ATTENTION_CLS, LlmSelfAttentionStruct._default_construct)\n\nLlmFeedForwardStruct.register_factory(FEEDFORWARD_CLS, LlmFeedForwardStruct._default_construct)\n\nLlmTransformerBlockStruct.register_factory(TRANSFORMER_BLOCK_CLS, LlmTransformerBlockStruct._default_construct)\n\nLlmTransformerStruct.register_factory(TRANSFORMER_CLS, LlmTransformerStruct._default_construct)\n\nLlmModelStruct.register_factory(\n tp.Union[TRANSFORMER_CLS, CASUALLM_CLS, SEQCLSLM_CLS, T5EncoderModel], LlmModelStruct._default_construct\n)\n"
},
{
"path": "deepcompressor/app/llm/ptq.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Evaluate a large language model.\"\"\"\n\nimport gc\nimport json\nimport os\nimport pprint\nimport traceback\n\nimport torch\nfrom transformers import GenerationConfig, PreTrainedModel, PreTrainedTokenizer\n\nfrom deepcompressor.utils import tools\n\nfrom .config import LlmCacheConfig, LlmPtqRunConfig, LlmQuantCacheConfig, LlmQuantConfig\nfrom .nn import LlmModelStruct\nfrom .quant import quantize_llm_activations, quantize_llm_weights, reorder_llm, rotate_llm, smooth_llm\n\n__all__ = [\"ptq\"]\n\n\ndef ptq( # noqa: C901\n model: PreTrainedModel | LlmModelStruct,\n /,\n tokenizer: PreTrainedTokenizer,\n config: LlmQuantConfig,\n cache: LlmCacheConfig | None = None,\n load_dirpath: str = \"\",\n save_dirpath: str = \"\",\n copy_on_save: bool = False,\n save_model: bool = False,\n) -> PreTrainedModel:\n \"\"\"Post-training quantization of a large language model.\n\n Args:\n model (`PreTrainedModel` or `LlmStruct`):\n The large language model.\n tokenizer (`PreTrainedTokenizer`):\n The large language model tokenizer.\n config (`LlmQuantConfig`):\n The large language model post-training quantization configuration.\n cache (`LlmCacheConfig`, *optional*, defaults to `None`):\n The large language model quantization cache path configuration.\n load_dirpath (`str`, *optional*, defaults to `\"\"`):\n The directory path to load the quantization checkpoint.\n save_dirpath (`str`, *optional*, defaults to `\"\"`):\n The directory path to save the quantization checkpoint.\n copy_on_save (`bool`, *optional*, defaults to `False`):\n Whether to copy the cache to the save directory.\n save_model (`bool`, *optional*, defaults to `False`):\n Whether to save the quantized model checkpoint.\n\n Returns:\n `PreTrainedModel`:\n The quantized model.\n \"\"\"\n logger = tools.logging.getLogger(__name__)\n if not isinstance(model, LlmModelStruct):\n model = LlmModelStruct.construct(model)\n assert isinstance(model, LlmModelStruct)\n\n quant_wgts = config.enabled_wgts\n quant_ipts = config.enabled_ipts\n quant_opts = config.enabled_opts\n quant_acts = quant_ipts or quant_opts\n quant = quant_wgts or quant_acts\n needs_rotation = quant and config.enabled_rotation\n needs_reorder = quant and config.enabled_reorder\n needs_smooth = quant and config.enabled_smooth\n\n load_model_path, load_path, save_path = \"\", None, None\n if load_dirpath:\n load_path = LlmQuantCacheConfig(\n rotation=os.path.join(load_dirpath, \"rotation.pt\"),\n reorder=os.path.join(load_dirpath, \"reorder.pt\"),\n smooth=os.path.join(load_dirpath, \"smooth.pt\"),\n wgts=os.path.join(load_dirpath, \"wgts.pt\"),\n acts=os.path.join(load_dirpath, \"acts.pt\"),\n )\n load_model_path = os.path.join(load_dirpath, \"model.pt\")\n if os.path.exists(load_model_path):\n logger.info(f\"* Found the model from {load_model_path}\")\n load_model = True\n save_dirpath = \"\" # do not save the model if loading\n if needs_reorder and not config.reorder.dynamic:\n needs_reorder = False\n logger.info(\"* Safe to skip reordering the model\")\n if needs_smooth:\n needs_smooth = False\n logger.info(\"* Safe to skip smoothing the model\")\n else:\n logger.warning(f\"Model checkpoint {load_model_path} does not exist\")\n load_model, load_model_path = False, \"\"\n else:\n load_model = False\n if save_dirpath:\n os.makedirs(save_dirpath, exist_ok=True)\n save_path = LlmQuantCacheConfig(\n rotation=os.path.join(save_dirpath, \"rotation.pt\"),\n reorder=os.path.join(save_dirpath, \"reorder.pt\"),\n smooth=os.path.join(save_dirpath, \"smooth.pt\"),\n wgts=os.path.join(save_dirpath, \"wgts.pt\"),\n acts=os.path.join(save_dirpath, \"acts.pt\"),\n )\n else:\n save_model = False\n\n # region rotate model\n if needs_rotation:\n logger.info(\"* Rotating model\")\n tools.logging.Formatter.indent_inc()\n load_from = \"\"\n if load_path and os.path.exists(load_path.rotation):\n load_from = load_path.rotation\n elif cache and cache.path.rotation and os.path.exists(cache.path.rotation):\n load_from = cache.path.rotation\n elif os.path.exists(config.rotation.path):\n load_from = config.rotation.path\n if load_from:\n logger.info(f\"- Loading rotation from {load_from}\")\n rotation = torch.load(load_from).to(dtype=torch.float64)\n rotate_llm(model, config.rotation, rotation=rotation)\n else:\n logger.info(\"- Generating rotation\")\n rotation = rotate_llm(model, config.rotation)\n if cache and cache.path.rotation:\n logger.info(f\"- Saving rotation to {cache.path.rotation}\")\n os.makedirs(cache.dirpath.rotation, exist_ok=True)\n torch.save(rotation, cache.path.rotation)\n load_from = cache.path.rotation\n if save_path:\n if not copy_on_save and load_from:\n logger.info(f\"- Linking rotation to {save_path.rotation}\")\n os.symlink(os.path.relpath(load_from, save_dirpath), save_path.rotation)\n else:\n logger.info(f\"- Saving rotation to {save_path.rotation}\")\n torch.save(rotation, save_path.rotation)\n del rotation\n tools.logging.Formatter.indent_dec()\n gc.collect()\n torch.cuda.empty_cache()\n logger.info(f\"* Development dtype is {config.develop_dtype}\")\n # endregion\n # region reorder channels\n if needs_reorder:\n logger.info(\"* Reordering channels\")\n tools.logging.Formatter.indent_inc()\n load_from = \"\"\n if load_path and os.path.exists(load_path.reorder):\n load_from = load_path.reorder\n elif cache and cache.path.reorder and os.path.exists(cache.path.reorder):\n load_from = cache.path.reorder\n if load_from:\n logger.info(f\"- Loading reorder indices from {load_from}\")\n reorder_cache = torch.load(load_from)\n reorder_llm(model, config, tokenizer, reorder_cache=reorder_cache)\n else:\n logger.info(\"- Generating reorder indices\")\n reorder_cache = reorder_llm(model, config, tokenizer)\n if cache and cache.path.reorder:\n logger.info(f\"- Saving reorder indices to {cache.path.reorder}\")\n os.makedirs(cache.dirpath.reorder, exist_ok=True)\n torch.save(reorder_cache, cache.path.reorder)\n load_from = cache.path.reorder\n if save_path:\n if not copy_on_save and load_from:\n logger.info(f\"- Linking reorder indices to {save_path.reorder}\")\n os.symlink(os.path.relpath(load_from, save_dirpath), save_path.reorder)\n else:\n logger.info(f\"- Saving reorder indices to {save_path.reorder}\")\n torch.save(reorder_cache, save_path.reorder)\n del reorder_cache\n tools.logging.Formatter.indent_dec()\n gc.collect()\n torch.cuda.empty_cache()\n # endregion\n # region smooth quantization\n if needs_smooth:\n logger.info(\"* Smoothing model for quantization\")\n tools.logging.Formatter.indent_inc()\n load_from = \"\"\n if load_path and os.path.exists(load_path.smooth):\n load_from = load_path.smooth\n elif cache and cache.path.smooth and os.path.exists(cache.path.smooth):\n load_from = cache.path.smooth\n if load_from:\n logger.info(f\"- Loading smooth scales from {load_from}\")\n smooth_cache = torch.load(load_from)\n smooth_llm(model, config, smooth_cache=smooth_cache)\n else:\n logger.info(\"- Generating smooth scales\")\n smooth_cache = smooth_llm(model, config, tokenizer=tokenizer)\n if cache and cache.path.smooth:\n logger.info(f\"- Saving smooth scales to {cache.path.smooth}\")\n os.makedirs(cache.dirpath.smooth, exist_ok=True)\n torch.save(smooth_cache, cache.path.smooth)\n load_from = cache.path.smooth\n if save_path:\n if not copy_on_save and load_from:\n logger.info(f\"- Linking smooth scales to {save_path.smooth}\")\n os.symlink(os.path.relpath(load_from, save_dirpath), save_path.smooth)\n else:\n logger.info(f\"- Saving smooth scales to {save_path.smooth}\")\n torch.save(smooth_cache, save_path.smooth)\n del smooth_cache\n tools.logging.Formatter.indent_dec()\n gc.collect()\n torch.cuda.empty_cache()\n # endregion\n # region collect original state dict\n if config.needs_acts_quantizer_cache:\n if load_path and os.path.exists(load_path.acts):\n orig_state_dict = None\n elif cache and cache.path.acts and os.path.exists(cache.path.acts):\n orig_state_dict = None\n else:\n orig_state_dict: dict[str, torch.Tensor] = {\n name: param.detach().clone() for name, param in model.module.named_parameters() if param.ndim > 1\n }\n else:\n orig_state_dict = None\n # endregion\n if load_model:\n logger.info(f\"* Loading model checkpoint from {load_model_path}\")\n model.module.load_state_dict(torch.load(load_model_path))\n gc.collect()\n torch.cuda.empty_cache()\n elif quant_wgts:\n logger.info(\"* Quantizing weights\")\n tools.logging.Formatter.indent_inc()\n load_from = \"\"\n if load_path and os.path.exists(load_path.wgts):\n load_from = load_path.wgts\n elif cache and cache.path.wgts and os.path.exists(cache.path.wgts):\n load_from = cache.path.wgts\n if load_from:\n logger.info(f\"- Loading weight quantizer settings from {load_from}\")\n quantizer_state_dict = torch.load(load_from)\n _, scale_state_dict = quantize_llm_weights(\n model,\n config,\n tokenizer=tokenizer,\n quantizer_state_dict=quantizer_state_dict,\n return_with_scale_state_dict=save_model,\n )\n else:\n logger.info(\"- Generating weight quantizer settings\")\n quantizer_state_dict, scale_state_dict = quantize_llm_weights(\n model, config, tokenizer=tokenizer, return_with_scale_state_dict=save_model\n )\n if cache and cache.dirpath.wgts:\n logger.info(f\"- Saving weight quantizer settings to {cache.path.wgts}\")\n os.makedirs(cache.dirpath.wgts, exist_ok=True)\n torch.save(quantizer_state_dict, cache.path.wgts)\n load_from = cache.path.wgts\n if save_path:\n if not copy_on_save and load_from:\n logger.info(f\"- Linking weight quantizer settings to {save_path.wgts}\")\n os.symlink(os.path.relpath(load_from, save_dirpath), save_path.wgts)\n else:\n logger.info(f\"- Saving weight quantizer settings to {save_path.wgts}\")\n torch.save(quantizer_state_dict, save_path.wgts)\n if save_model:\n logger.info(f\"- Saving model checkpoint to {save_dirpath}\")\n torch.save(scale_state_dict, os.path.join(save_dirpath, \"scale.pt\"))\n torch.save(model.module.state_dict(), os.path.join(save_dirpath, \"model.pt\"))\n del quantizer_state_dict, scale_state_dict\n tools.logging.Formatter.indent_dec()\n gc.collect()\n torch.cuda.empty_cache()\n if quant_acts:\n logger.info(\" * Quantizing activations\")\n tools.logging.Formatter.indent_inc()\n if config.needs_acts_quantizer_cache:\n load_from = \"\"\n if load_path and os.path.exists(load_path.acts):\n load_from = load_path.acts\n elif cache and cache.path.acts and os.path.exists(cache.path.acts):\n load_from = cache.path.acts\n if load_from:\n logger.info(f\"- Loading activation quantizer settings from {load_from}\")\n quantizer_state_dict = torch.load(load_from)\n quantize_llm_activations(\n model,\n config,\n tokenizer=tokenizer,\n quantizer_state_dict=quantizer_state_dict,\n orig_state_dict=orig_state_dict,\n )\n else:\n logger.info(\"- Generating activation quantizer settings\")\n quantizer_state_dict = quantize_llm_activations(\n model, config, tokenizer=tokenizer, orig_state_dict=orig_state_dict\n )\n if cache and cache.dirpath.acts:\n logger.info(f\"- Saving activation quantizer settings to {cache.path.acts}\")\n os.makedirs(cache.dirpath.acts, exist_ok=True)\n torch.save(quantizer_state_dict, cache.path.acts)\n load_from = cache.path.acts\n if save_dirpath:\n if not copy_on_save and load_from:\n logger.info(f\"- Linking activation quantizer settings to {save_path.acts}\")\n os.symlink(os.path.relpath(load_from, save_dirpath), save_path.acts)\n else:\n logger.info(f\"- Saving activation quantizer settings to {save_path.acts}\")\n torch.save(quantizer_state_dict, save_path.acts)\n del quantizer_state_dict\n else:\n logger.info(\"- No need to generate/load activation quantizer settings\")\n quantize_llm_activations(model, config, tokenizer=tokenizer, orig_state_dict=orig_state_dict)\n tools.logging.Formatter.indent_dec()\n del orig_state_dict\n gc.collect()\n torch.cuda.empty_cache()\n return model.module\n\n\ndef main(config: LlmPtqRunConfig, logging_level: int = tools.logging.DEBUG) -> None: # noqa: C901\n \"\"\"Post-training quantization and evaluation of a large language model.\n\n Args:\n config (`LlmPtqConfig`):\n The large language model post-training quantization configuration.\n logging_level (`int`, *optional*, defaults to `logging.DEBUG`):\n The logging level.\n \"\"\"\n config.output.lock()\n config.dump(path=config.output.get_running_job_path(\"config.yaml\"))\n tools.logging.setup(path=config.output.get_running_job_path(\"run.log\"), level=logging_level)\n logger = tools.logging.getLogger(__name__)\n # region log configurations\n logger.info(\"=== Configurations ===\")\n tools.logging.info(config.formatted_str(), logger=logger)\n logger.info(\"=== Dumped Configurations ===\")\n tools.logging.info(pprint.pformat(config.dump(), indent=2, width=120), logger=logger)\n logger.info(\"=== Output Directory ===\")\n logger.info(config.output.job_dirpath)\n # endregion\n logger.info(\"=== Start Evaluating ===\")\n logger.info(f\"* Building model {config.model.name} from {config.model.path}\")\n tools.logging.Formatter.indent_inc()\n model, tokenizer = config.model.build()\n tools.logging.Formatter.indent_dec()\n save_dirpath = os.path.join(config.output.running_job_dirpath, \"cache\")\n if config.save_model:\n if config.save_model.lower() in (\"false\", \"none\", \"null\", \"nil\"):\n save_model = False\n elif config.save_model.lower() in (\"true\", \"default\"):\n save_dirpath, save_model = os.path.join(config.output.running_job_dirpath, \"model\"), True\n else:\n save_dirpath, save_model = config.save_model, True\n else:\n save_model = False\n model = ptq(\n model,\n tokenizer=tokenizer,\n config=config.quant,\n cache=config.cache,\n load_dirpath=config.load_from,\n save_dirpath=save_dirpath,\n copy_on_save=config.copy_on_save,\n save_model=save_model,\n )\n # region evaluate model\n if not config.skip_eval:\n logger.info(\"* Evaluating model\")\n eos_token_ids = GenerationConfig.from_pretrained(config.model.path).eos_token_id\n if not isinstance(eos_token_ids, list):\n eos_token_ids = [eos_token_ids]\n tools.logging.Formatter.indent_inc()\n results = config.eval.evaluate(\n model,\n tokenizer,\n model_name=config.model.name,\n eos_token_ids=eos_token_ids,\n output_dirpath=config.output.get_running_job_path(\"eval\"),\n )\n tools.logging.Formatter.indent_dec()\n logger.info(f\"* Saving results to {config.output.job_dirpath}\")\n # dump results\n with open(os.path.join(config.output.get_running_job_path(\"results.json\")), \"w\") as f:\n json.dump(results, f, indent=2)\n # endregion\n config.output.unlock()\n\n\nif __name__ == \"__main__\":\n config, _, unused_cfgs, unused_args, unknown_args = LlmPtqRunConfig.get_parser().parse_known_args()\n if len(unused_cfgs) > 0:\n tools.logging.warning(f\"Unused configurations: {unused_cfgs}\")\n if unused_args is not None:\n tools.logging.warning(f\"Unused arguments: {unused_args}\")\n assert len(unknown_args) == 0, f\"Unknown arguments: {unknown_args}\"\n try:\n main(config, logging_level=tools.logging.DEBUG)\n except Exception as e:\n tools.logging.Formatter.indent_reset()\n tools.logging.error(\"=== Error ===\")\n tools.logging.error(traceback.format_exc())\n tools.logging.shutdown()\n traceback.print_exc()\n config.output.unlock(error=True)\n raise e\n"
},
{
"path": "deepcompressor/app/llm/quant/__init__.py",
"content": "# -*- coding: utf-8 -*-\n\nfrom .activation import quantize_llm_activations\nfrom .config import LlmQuantCacheConfig, LlmQuantConfig\nfrom .quantizer import LlmActivationQuantizer, LlmWeightQuantizer\nfrom .reorder import reorder_llm\nfrom .rotate import rotate_llm\nfrom .smooth import smooth_llm\nfrom .weight import quantize_llm_weights\n"
},
{
"path": "deepcompressor/app/llm/quant/activation.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"LLM activation quantization calibration module.\"\"\"\n\nimport gc\nimport typing as tp\n\nimport torch\nimport torch.nn as nn\nfrom tqdm import tqdm\nfrom transformers import PreTrainedTokenizer\n\nfrom deepcompressor.data.cache import IOTensorsCache\nfrom deepcompressor.data.common import TensorType\nfrom deepcompressor.utils import tools\n\nfrom ..nn import LlmModelStruct, LlmTransformerBlockStruct\nfrom .config import LlmQuantConfig\nfrom .quantizer import LlmActivationQuantizer\nfrom .utils import get_needs_inputs_fn, get_needs_outputs_fn\n\n__all__ = [\"quantize_llm_activations\"]\n\n\n@torch.inference_mode()\ndef quantize_llm_layer_activations( # noqa: C901\n layer: LlmTransformerBlockStruct,\n config: LlmQuantConfig,\n quantizer_state_dict: dict[str, tp.Any],\n layer_cache: dict[str, IOTensorsCache] | None = None,\n layer_kwargs: dict[str, tp.Any] | None = None,\n orig_state_dict: dict[str, torch.Tensor] | None = None,\n) -> None:\n \"\"\"Calibrate the activation quantization ranges of modules in a layer.\n\n Args:\n layer (`LlmTransformerBlockStruct`):\n Layer.\n config (`LlmQuantConfig`):\n Quantization configuration.\n quantizer_state_dict (`dict[str, Any]`):\n Activation quantizer state dict.\n layer_cache (`dict[str, IOTensorsCache]` or `None`, *optional*, defaults to `None`):\n Layer activations cache.\n layer_kwargs (`dict[str, tp.Any]` or `None`, *optional*, defaults to `None`):\n Keyword arguments for the layer.\n orig_state_dict (`dict[str, torch.Tensor]` or `None`, *optional*, defaults to `None`):\n Original weight state dict.\n \"\"\"\n logger = tools.logging.getLogger(f\"{__name__}.ActivationQuant\")\n logger.debug(\"- Quantizing layer %s\", layer.name)\n layer_cache = layer_cache or {}\n layer_kwargs = layer_kwargs or {}\n orig_state_dict = orig_state_dict or {}\n args_caches: list[\n tuple[\n str, # key\n TensorType,\n list[nn.Linear], # modules\n str, # module name\n nn.Module, # eval module\n str, # eval name\n dict[str, tp.Any], # eval kwargs\n list[tuple[nn.Parameter, torch.Tensor]], # original wgts\n ]\n ] = []\n In, Out = TensorType.Inputs, TensorType.Outputs\n\n attn, ffn = layer.attn_struct, layer.ffn_struct\n # region attn\n attn_kwargs = attn.filter_kwargs(layer_kwargs)\n if orig_state_dict:\n orig_wgts = [\n (module.weight, orig_state_dict[f\"{module_name}.weight\"])\n for module_name, module in zip(attn.qkv_proj_names, attn.qkv_proj, strict=True)\n ] + [(attn.out_proj.weight, orig_state_dict[f\"{attn.out_proj_name}.weight\"])]\n else:\n orig_wgts = None\n # region qkv_proj (Inputs)\n module_name = attn.v_proj_name\n module_key, cache_key, modules = attn.qkv_proj_key, f\"{module_name}.input\", attn.qkv_proj\n args_caches.append((module_key, In, modules, module_name, attn, attn.name, attn_kwargs, orig_wgts))\n # endregion\n # region qkv_attn (Outputs)\n orig_proj_wgts = (orig_wgts + orig_wgts) if orig_wgts else None\n for idx, module_key in enumerate((attn.q_key, attn.k_key, attn.v_key)):\n module = getattr(attn, \"qkv\"[idx])\n module_name = getattr(attn, f\"{'qkv'[idx]}_name\")\n cache_key = f\"{module_name}.output\"\n orig_wgts = orig_proj_wgts[idx : idx + 4] if orig_proj_wgts else None\n args_caches.append((module_key, Out, [module], module_name, attn, attn.name, attn_kwargs, orig_wgts))\n # endregion\n # region out_proj (Inputs)\n module_name, module = attn.out_proj_name, attn.out_proj\n module_key, cache_key = attn.out_proj_key, f\"{module_name}.input\"\n orig_wgts = [(module.weight, orig_state_dict[f\"{module_name}.weight\"])] if orig_state_dict else None\n args_caches.append((module_key, In, [module], module_name, module, module_name, None, orig_wgts))\n # endregion\n del orig_wgts\n # endregion\n # region ffn\n # region ffn block projections\n for expert_idx in range(ffn.config.num_experts):\n expert = ffn.experts[expert_idx]\n expert_name = ffn.expert_names[expert_idx]\n # region proj 1st in expert (Inputs)\n module_name = ffn.up_proj_names[expert_idx]\n modules = ffn.up_projs[expert_idx :: ffn.config.num_experts]\n module_key, cache_key = ffn.up_proj_key, f\"{module_name}.input\"\n if orig_state_dict:\n orig_wgts = [\n (module.weight, orig_state_dict[f\"{expert_name}.{ffn.up_proj_rnames[module_idx]}.weight\"])\n for module_idx, module in enumerate(modules)\n ]\n else:\n orig_wgts = None\n args_caches.append((module_key, In, modules, module_name, expert, module_name, None, orig_wgts))\n # endregion\n # region proj 2nd in expert (Inputs)\n module_name, module = ffn.down_proj_names[expert_idx], ffn.down_projs[expert_idx]\n module_key, cache_key = ffn.down_proj_key, f\"{module_name}.input\"\n if orig_state_dict:\n orig_wgts = [(module.weight, orig_state_dict[f\"{module_name}.weight\"])]\n else:\n orig_wgts = None\n args_caches.append((module_key, In, [module], module_name, module, module_name, None, orig_wgts))\n # endregion\n # endregion\n # endregion\n quantizers: dict[str, LlmActivationQuantizer] = {}\n tools.logging.Formatter.indent_inc()\n for module_key, tensor_type, modules, module_name, eval_module, eval_name, eval_kwargs, orig_wgts in args_caches:\n if tensor_type == TensorType.Inputs:\n cache_key = f\"{module_name}.input\"\n quantizer_config = config.ipts\n activations = layer_cache.get(module_name, IOTensorsCache()).inputs\n device = modules[0].weight.device\n else:\n cache_key = f\"{module_name}.output\"\n quantizer_config = config.opts\n activations = layer_cache.get(module_name, IOTensorsCache()).outputs\n device = attn.out_proj.weight.device\n quantizer = LlmActivationQuantizer(\n quantizer_config,\n channels_dim=-1,\n develop_dtype=config.develop_dtype,\n key=module_key,\n tensor_type=tensor_type,\n )\n if quantizer.is_enabled():\n quantizers[cache_key] = quantizer\n if cache_key not in quantizer_state_dict:\n logger.debug(\"- Calibrating %s\", cache_key)\n quantizer.calibrate_dynamic_range(\n modules=modules,\n activations=activations,\n eval_module=eval_module,\n eval_inputs=layer_cache[eval_name].inputs if layer_cache else None,\n eval_kwargs=eval_kwargs,\n orig_weights=orig_wgts,\n )\n quantizer_state_dict[cache_key] = quantizer.state_dict()\n gc.collect()\n torch.cuda.empty_cache()\n else:\n quantizer.load_state_dict(quantizer_state_dict[cache_key], device=device)\n if tensor_type == TensorType.Inputs:\n if attn.v_proj_rname in cache_key:\n for proj_name in [attn.q_proj_rname, attn.k_proj_rname]:\n quantizers[cache_key.replace(attn.v_proj_rname, proj_name)] = quantizer\n if ffn.up_proj_rnames[0] in cache_key:\n for proj_name in ffn.up_proj_rnames[1:]:\n quantizers[cache_key.replace(ffn.up_proj_rnames[0], proj_name)] = quantizer\n del quantizer\n for name, module in layer.module.named_modules():\n module_name = f\"{layer.name}.{name}\"\n ipts_quantizer = quantizers.get(f\"{module_name}.input\", None)\n opts_quantizer = quantizers.get(f\"{module_name}.output\", None)\n needs_quant_ipts = ipts_quantizer is not None and ipts_quantizer.is_enabled()\n needs_quant_opts = opts_quantizer is not None and opts_quantizer.is_enabled()\n if needs_quant_ipts or needs_quant_opts:\n logger.debug(\n \"- Quantizing %s (%s)\",\n module_name,\n (\"inputs\" if needs_quant_ipts else \"\")\n + (\" and \" if needs_quant_ipts and needs_quant_opts else \"\")\n + (\"outputs\" if needs_quant_opts else \"\"),\n )\n if needs_quant_ipts:\n ipts_quantizer.as_hook(is_output=False).register(module)\n if needs_quant_opts:\n opts_quantizer.as_hook(is_output=True).register(module)\n tools.logging.Formatter.indent_dec()\n\n\n@torch.inference_mode()\ndef quantize_llm_activations(\n model: nn.Module | LlmModelStruct,\n config: LlmQuantConfig,\n tokenizer: PreTrainedTokenizer | None = None,\n quantizer_state_dict: dict[str, tp.Any] | None = None,\n orig_state_dict: dict[str, torch.Tensor] | None = None,\n) -> dict[str, tp.Any]:\n \"\"\"Quantize the large foundation model activations.\n\n Args:\n model (`nn.Module` or `LlmStruct`):\n Model to be quantized.\n config (`LlmQuantConfig`):\n Quantization configuration.\n tokenizer (`PreTrainedTokenizer`, *optional*, defaults to `None`):\n Tokenizer.\n quantizer_state_dict (`dict[str, Any]`, *optional*, defaults to `None`):\n Activation quantizer state dict cache.\n orig_state_dict (`dict[str, torch.Tensor]`, *optional*, defaults to `None`):\n Original weight state dict\n\n Returns:\n `dict[str, Any]`:\n Activation quantizer state dict cache.\n \"\"\"\n if not isinstance(model, LlmModelStruct):\n model = LlmModelStruct.construct(model)\n assert isinstance(model, LlmModelStruct)\n quantizer_state_dict = quantizer_state_dict or {}\n with tools.logging.redirect_tqdm():\n if not quantizer_state_dict and config.needs_acts_quantizer_cache:\n for _, (layer, layer_cache, layer_kwargs) in tqdm(\n config.calib.build_loader(tokenizer).iter_layer_activations(\n model,\n needs_inputs_fn=get_needs_inputs_fn(model=model, config=config),\n needs_outputs_fn=get_needs_outputs_fn(model=model, config=config),\n ),\n desc=\"quantizing activations\",\n leave=False,\n total=len(model.backbone_struct.layer_structs),\n dynamic_ncols=True,\n ):\n quantize_llm_layer_activations(\n layer=layer,\n config=config,\n quantizer_state_dict=quantizer_state_dict,\n layer_cache=layer_cache,\n layer_kwargs=layer_kwargs,\n orig_state_dict=orig_state_dict,\n )\n else:\n for layer in tqdm(\n model.backbone_struct.layer_structs,\n desc=\"quantizing activations\",\n leave=False,\n dynamic_ncols=True,\n ):\n quantize_llm_layer_activations(\n layer=layer,\n config=config,\n quantizer_state_dict=quantizer_state_dict,\n orig_state_dict=orig_state_dict,\n )\n return quantizer_state_dict\n"
},
{
"path": "deepcompressor/app/llm/quant/config.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Quantization config.\"\"\"\n\nimport os\nfrom dataclasses import dataclass, field\n\nimport torch\nfrom omniconfig import configclass\n\nfrom deepcompressor.calib.config import (\n QuantRotationConfig,\n SearchBasedCalibGranularity,\n SearchBasedCalibObjective,\n SearchBasedCalibStrategy,\n SkipBasedChannelOrderConfig,\n SmoothTransfomerConfig,\n)\nfrom deepcompressor.data.utils.dtype import eval_dtype\nfrom deepcompressor.utils.common import num2str\n\nfrom ..cache.config import LlmQuantCacheConfig\nfrom ..nn.struct import LlmFeedForwardStruct, LlmSelfAttentionStruct\nfrom .dataset import LlmCalibDataLoaderConfig\nfrom .quantizer import LlmModuleQuantizerConfig\n\n__all__ = [\"LlmQuantConfig\"]\n\n\n@configclass\n@dataclass\nclass LlmQuantConfig(LlmModuleQuantizerConfig):\n \"\"\"Large Language Model Module quantization configuration.\n\n Args:\n wgts (`LlmWeightQuantizerConfig`):\n The weight quantization configuration.\n ipts (`LlmActivationQuantizerConfig`):\n The input activation quantization configuration.\n opts (`LlmActivationQuantizerConfig`):\n The output activation quantization configuration.\n calib (`LlmCalibDataLoaderConfig`):\n The calibration dataset configuration.\n rotation (`QuantRotationConfig` or `None`, *optional*, defaults to `None`):\n The quantization rotation configuration.\n reorder (`SkipBasedChannelOrderConfig` or `None`, *optional*, defaults to `None`):\n The quantization reordering configuration.\n smooth (`SmoothTransfomerConfig`, *optional*, defaults to `None`):\n The quantization smoothing configuration.\n develop_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The development data type during quantization.\n \"\"\"\n\n calib: LlmCalibDataLoaderConfig\n rotation: QuantRotationConfig | None = None\n reorder: SkipBasedChannelOrderConfig | None = None\n smooth: SmoothTransfomerConfig | None = None\n develop_dtype: torch.dtype = field(default_factory=lambda s=torch.float32: eval_dtype(s, with_quant_dtype=False))\n\n def __post_init__(self) -> None: # noqa: C901\n if self.smooth is not None:\n if not self.smooth.enabled_proj and not self.smooth.enabled_attn:\n self.smooth = None\n if self.rotation is not None and self.reorder is not None:\n self.reorder.skips.append(\"residual\")\n if self.rotation.transforms:\n self.reorder.skips.extend(self.rotation.transforms)\n self.reorder.skips = sorted(set(self.reorder.skips))\n if self.enabled_ipts:\n if self.ipts.enabled_calib_range and self.ipts.calib_range.granularity == SearchBasedCalibGranularity.Group:\n self.ipts.calib_range.granularity = SearchBasedCalibGranularity.ChannelGroup\n if self.ipts.static:\n assert self.ipts.smallest_group_shape[0] == -1, \"static quantization requires batch group size to be -1\"\n if self.enabled_opts:\n if self.opts.enabled_calib_range and self.opts.calib_range.granularity == SearchBasedCalibGranularity.Group:\n self.opts.calib_range.granularity = SearchBasedCalibGranularity.ChannelGroup\n if self.opts.static:\n assert self.opts.smallest_group_shape[0] == -1, \"static quantization requires batch group size to be -1\"\n if self.enabled_reorder:\n if not self.reorder.dynamic:\n qkv_proj_rkey, up_proj_rkey = LlmSelfAttentionStruct.qkv_proj_rkey, LlmFeedForwardStruct.up_proj_rkey\n skips_to_remove = []\n for skip in self.reorder.skips:\n if skip.startswith(qkv_proj_rkey) or skip.endswith(f\"_{qkv_proj_rkey}\"):\n self.reorder.skips.append(\"residual\")\n skips_to_remove.append(skip)\n elif skip.startswith(up_proj_rkey) or skip.endswith(f\"_{up_proj_rkey}\"):\n self.reorder.skips.append(\"residual\")\n skips_to_remove.append(skip)\n self.reorder.skips = sorted(set(self.reorder.skips))\n for skip in skips_to_remove:\n self.reorder.skips.remove(skip)\n self.reorder.skips = sorted(set(self.reorder.skips))\n\n @property\n def enabled_smooth(self) -> bool:\n \"\"\"Whether to enable smooth quantization.\"\"\"\n return self.smooth is not None\n\n @property\n def enabled_smooth_proj(self) -> bool:\n \"\"\"Whether to enable xw smooth quantization.\"\"\"\n return self.enabled_smooth and self.smooth.enabled_proj\n\n @property\n def enabled_smooth_attn(self) -> bool:\n \"\"\"Whether to enable yy smooth quantization.\"\"\"\n return self.enabled_smooth and self.smooth.enabled_attn\n\n @property\n def enabled_reorder(self) -> bool:\n \"\"\"Whether to enable channel reorder.\"\"\"\n return self.reorder is not None and self.reorder.is_enabled()\n\n @property\n def enabled_rotation(self) -> bool:\n \"\"\"Whether to enable rotation.\"\"\"\n return self.rotation is not None\n\n @property\n def needs_acts_quantizer_cache(self) -> bool:\n \"\"\"Whether to cache the activations quantizer settings.\"\"\"\n if self.enabled_ipts and self.ipts.needs_calib_data:\n return True\n if self.enabled_opts and self.opts.needs_calib_data:\n return True\n return False\n\n def generate_calib_dirname(self) -> str:\n name = \"\"\n if self.enabled_rotation:\n name += \"-rotate\"\n if self.rotation.random:\n name += \".rnd\"\n if self.enabled_reorder:\n name += \"-reorder\"\n if self.reorder.dynamic:\n name += \".dyn\"\n if self.enabled_smooth:\n name += \"-smooth\"\n if self.enabled_smooth_proj:\n name += \".proj\"\n if self.enabled_smooth_attn:\n name += \".attn\"\n calib_name = super().generate_calib_dirname()\n if calib_name:\n name += f\"-{calib_name}\"\n return name[1:] if name else name\n\n def generate_default_dirname(self) -> str: # noqa: C901\n \"\"\"Generate directory name for a large language model quantization configuration.\"\"\"\n w_names = x_names = {\"qkv_proj\": \"qkv\", \"out_proj\": \"out\", \"up_proj\": \"fc1\", \"down_proj\": \"fc2\"}\n y_names = {\"attn_q\": \"q\", \"attn_k\": \"k\", \"attn_v\": \"v\"}\n skip_name = \"\"\n if self.enabled_opts:\n skip_y_name = \"+\".join(y_names[y] for y in self.opts.skips if y in y_names)\n if skip_y_name:\n skip_name += f\".y.[{skip_y_name}]\"\n if self.enabled_wgts:\n skip_w_name = \"+\".join(w_names[w] for w in self.wgts.skips if w in w_names)\n if skip_w_name:\n skip_name += f\".w.[{skip_w_name}]\"\n if self.enabled_ipts:\n skip_x_name = \"+\".join(x_names[x] for x in self.ipts.skips if x in x_names)\n if skip_x_name:\n skip_name += f\".x.[{skip_x_name}]\"\n if skip_name:\n skip_name = \"-skip\" + skip_name\n if self.enabled_wgts and self.wgts.enabled_gptq:\n skip_name += \"-gptq\"\n rotation_name = \"\"\n if self.enabled_rotation:\n rotation_name = \"-rot\"\n if self.rotation.path:\n rotation_name += f\".{self.rotation.name}\"\n elif self.rotation.random:\n rotation_name += \".rnd\"\n if self.rotation.transforms:\n rotation_name += \".[+{}]\".format(\"+\".join(w_names[w] for w in self.rotation.transforms))\n reorder_name = \"\"\n if self.enabled_reorder:\n reorder_name = \"-rodr\"\n if self.reorder.strategy == SearchBasedCalibStrategy.Manual:\n if self.reorder.channel_metric.value != \"xMax\":\n reorder_name += f\".{self.reorder.channel_metric.value}\"\n if self.reorder.channel_index.value != \"Seq\":\n reorder_name += f\".{self.reorder.channel_index.value}\"\n else:\n reorder_name += f\".{self.reorder.strategy.name}\"\n reorders, skips = [], []\n for k in w_names.keys() if self.reorder.dynamic else (\"residual\", \"out_proj\", \"down_proj\"):\n v = w_names.get(k, \"res\")\n if k in self.reorder.skips:\n skips.append(v)\n else:\n reorders.append(v)\n if len(reorders) <= len(skips):\n reorder_name += \".[{}]\".format(\"+\".join(reorders))\n elif skips:\n reorder_name += \".skip.[{}]\".format(\"+\".join(skips))\n smooth_name = \"\"\n if self.enabled_smooth:\n smooth_name = \"-smth\"\n if self.smooth.enabled_proj:\n smooth_name += \".proj\"\n if self.smooth.proj.granularity != SearchBasedCalibGranularity.Layer:\n smooth_name += f\".{self.smooth.proj.granularity.name}\"\n if self.smooth.proj.strategy != SearchBasedCalibStrategy.Manual:\n smooth_name += f\".{self.smooth.proj.strategy.name}\"\n if self.smooth.proj.alpha <= 0:\n smooth_name += f\".a{num2str(self.smooth.proj.alpha)}\"\n if self.smooth.proj.beta <= 0:\n smooth_name += f\".b{num2str(self.smooth.proj.beta)}\"\n else:\n smooth_name += f\".a{num2str(self.smooth.proj.alpha)}\"\n smooth_name += f\".b{num2str(self.smooth.proj.beta)}\"\n xspan_eq_wspan = True\n for xspan, wspan in self.smooth.proj.spans:\n if xspan != wspan:\n xspan_eq_wspan = False\n break\n if xspan_eq_wspan:\n smooth_name += \".[{}]\".format(\"+\".join(xspan.name for xspan, _ in self.smooth.proj.spans))\n else:\n smooth_name += \".[{}]\".format(\n \"+\".join(f\"x.{xspan.name}.w.{wspan.name}\" for xspan, wspan in self.smooth.proj.spans)\n )\n smooths, skips = [], []\n for k, v in w_names.items():\n if k in self.smooth.proj.skips:\n skips.append(v)\n else:\n smooths.append(v)\n if len(smooths) <= len(skips):\n smooth_name += \".[{}]\".format(\"+\".join(smooths))\n elif skips:\n smooth_name += \".skip.[{}]\".format(\"+\".join(skips))\n if self.smooth.enabled_attn:\n smooth_name += \".attn\"\n if self.smooth.attn.granularity != SearchBasedCalibGranularity.Layer:\n smooth_name += f\".{self.smooth.attn.granularity.name}\"\n if self.smooth.attn.strategy != SearchBasedCalibStrategy.Manual:\n smooth_name += f\".{self.smooth.attn.strategy.name}\"\n if self.smooth.attn.alpha <= 0:\n smooth_name += f\".a{num2str(self.smooth.attn.alpha)}\"\n if self.smooth.attn.beta <= 0:\n smooth_name += f\".b{num2str(self.smooth.attn.beta)}\"\n else:\n smooth_name += f\".a{num2str(self.smooth.attn.alpha)}\"\n smooth_name += f\".b{num2str(self.smooth.attn.beta)}\"\n xspan_eq_yspan = True\n for xspan, yspan in self.smooth.attn.spans:\n if xspan != yspan:\n xspan_eq_yspan = False\n break\n if xspan_eq_yspan:\n smooth_name += \".[{}]\".format(\"+\".join(xspan.name for xspan, _ in self.smooth.attn.spans))\n else:\n smooth_name += \".[{}]\".format(\n \"+\".join(f\"x.{xspan.name}.y.{yspan.name}\" for xspan, yspan in self.smooth.attn.spans)\n )\n wrange_name = \"\"\n if (\n self.enabled_wgts\n and self.wgts.enabled_calib_range\n and (self.wgts.calib_range.needs_search or self.wgts.calib_range.ratio != 1)\n ):\n wrange_name = \"-w.range\"\n if self.wgts.calib_range.needs_search:\n if self.wgts.calib_range.granularity != SearchBasedCalibGranularity.Group:\n wrange_name += f\".{self.wgts.calib_range.granularity.name}\"\n if self.wgts.calib_range.objective != SearchBasedCalibObjective.OutputsError:\n wrange_name += f\".{self.wgts.calib_range.objective.name}\"\n if self.wgts.calib_range.degree != 2:\n wrange_name += f\".d{num2str(self.wgts.calib_range.degree)}\"\n wrange_name += f\".[{num2str(self.wgts.calib_range.max_shrink)}\"\n wrange_name += f\".{num2str(self.wgts.calib_range.max_expand)}\"\n wrange_name += f\".g{self.wgts.calib_range.num_grids}]\"\n else:\n wrange_name += f\".r{num2str(self.wgts.calib_range.ratio)}\"\n if self.wgts.calib_range.skips:\n wrange_name += \".skip.[{}]\".format(\"+\".join(w_names[w] for w in self.wgts.calib_range.skips))\n xrange_name = \"\"\n if (\n self.enabled_ipts\n and self.ipts.enabled_calib_range\n and (self.ipts.calib_range.needs_search or self.ipts.calib_range.ratio != 1)\n ):\n xrange_name = \"-x.range\"\n if self.ipts.calib_range.needs_search:\n if self.ipts.calib_range.granularity != SearchBasedCalibGranularity.Group:\n xrange_name += f\".{self.ipts.calib_range.granularity.name}\"\n if self.ipts.calib_range.objective != SearchBasedCalibObjective.OutputsError:\n xrange_name += f\".{self.ipts.calib_range.objective.name}\"\n if self.ipts.calib_range.degree != 2:\n xrange_name += f\".d{num2str(self.ipts.calib_range.degree)}\"\n xrange_name += f\".[{num2str(self.ipts.calib_range.max_shrink)}\"\n xrange_name += f\".{num2str(self.ipts.calib_range.max_expand)}\"\n xrange_name += f\".g{self.ipts.calib_range.num_grids}]\"\n else:\n xrange_name += f\".r{num2str(self.ipts.calib_range.ratio)}\"\n if self.ipts.calib_range.skips:\n xrange_name += \".skip.[{}]\".format(\"+\".join(w_names[w] for w in self.ipts.calib_range.skips))\n yrange_name = \"\"\n if (\n self.enabled_opts\n and self.opts.enabled_calib_range\n and (self.opts.calib_range.needs_search or self.opts.calib_range.ratio != 1)\n ):\n yrange_name = \"-y.range\"\n if self.opts.calib_range.needs_search:\n if self.opts.calib_range.granularity != SearchBasedCalibGranularity.Group:\n yrange_name += f\".{self.opts.calib_range.granularity.name}\"\n if self.opts.calib_range.objective != SearchBasedCalibObjective.OutputsError:\n yrange_name += f\".{self.opts.calib_range.objective.name}\"\n if self.opts.calib_range.degree != 2:\n yrange_name += f\".d{num2str(self.opts.calib_range.degree)}\"\n yrange_name += f\".[{num2str(self.opts.calib_range.max_shrink)}\"\n yrange_name += f\".{num2str(self.opts.calib_range.max_expand)}\"\n yrange_name += f\".g{self.opts.calib_range.num_grids}]\"\n else:\n yrange_name += f\".r{num2str(self.opts.calib_range.ratio)}\"\n if self.opts.calib_range.skips:\n yrange_name += \".skip.[{}]\".format(\"+\".join(y_names[y] for y in self.opts.calib_range.skips))\n name = skip_name + rotation_name + reorder_name + smooth_name + wrange_name + xrange_name + yrange_name\n name = name[1:] if name else \"default\"\n name += f\"-{self.calib.generate_dirnames()[0]}\"\n return name\n\n def generate_cache_dirpath(\n self, *, root: str, seed: int, default_dtype: torch.dtype = torch.float16\n ) -> LlmQuantCacheConfig: # noqa: C901\n \"\"\"Generate the cache paths for the module quantization configuration.\"\"\"\n quant_names = self.generate_dirnames(default_dtype=default_dtype)\n w_kernel_names = []\n if self.enabled_wgts and self.wgts.enabled_gptq:\n w_kernel_names = self.wgts.kernel_gptq.generate_dirnames(prefix=\"w.kernel\")\n if self.enabled_rotation:\n quant_names.extend(self.rotation.generate_dirnames(prefix=\"rotate\"))\n reorder_dirpath = \"\"\n if self.enabled_reorder:\n reorder_names = self.reorder.generate_dirnames(prefix=\"reorder\")\n quant_names.extend(reorder_names)\n reorder_dirpath = os.path.join(\"reorder\", *quant_names)\n smooth_dirpath = \"\"\n if self.enabled_smooth:\n smooth_names = self.smooth.generate_dirnames(prefix=\"smooth\")\n quant_names.extend(smooth_names)\n smooth_dirpath = os.path.join(\"smooth\", *quant_names)\n quant_names.extend(w_kernel_names)\n wgts_dirpath = \"\"\n if self.enabled_wgts and self.wgts.enabled_calib_range:\n quant_names.extend(self.wgts.calib_range.generate_dirnames(prefix=\"w.range\"))\n wgts_dirpath = os.path.join(\"wgts\", *quant_names)\n acts_dirpath = \"\"\n if self.needs_acts_quantizer_cache:\n if self.enabled_ipts and self.ipts.enabled_calib_range:\n quant_names.extend(self.ipts.calib_range.generate_dirnames(prefix=\"x.range\"))\n if self.enabled_opts and self.opts.enabled_calib_range:\n quant_names.extend(self.opts.calib_range.generate_dirnames(prefix=\"y.range\"))\n acts_dirpath = os.path.join(\"acts\", *quant_names)\n cache_dirpath = LlmQuantCacheConfig(\n reorder=reorder_dirpath,\n smooth=smooth_dirpath,\n wgts=wgts_dirpath,\n acts=acts_dirpath,\n ).add_parent_dirs(*self.calib.generate_dirnames())\n if self.enabled_rotation:\n if self.rotation.path:\n cache_dirpath.rotation = \"\"\n else:\n cache_dirpath.rotation = os.path.join(\n \"rotation\",\n f\"seed.{seed}\" if self.rotation.random else \"hadamard\",\n )\n cache_dirpath.add_parent_dirs(root, \"llm\", \"cache\", \"quant\")\n return cache_dirpath\n"
},
{
"path": "deepcompressor/app/llm/quant/dataset.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Functions for collecting calibration dataset for quantization.\"\"\"\n\nimport os\nimport random\nimport typing as tp\nfrom dataclasses import MISSING, dataclass, field\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.data\nfrom datasets import load_dataset\nfrom omniconfig import configclass\nfrom transformers import PreTrainedTokenizer\nfrom transformers.cache_utils import Cache\nfrom transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock\nfrom transformers.models.t5.modeling_t5 import T5DenseActDense, T5DenseGatedActDense\n\nfrom deepcompressor.data.cache import IOTensorsCache, ModuleForwardInput, TensorCache\nfrom deepcompressor.data.utils.reshape import LinearReshapeFn\nfrom deepcompressor.dataset.action import CacheAction, ConcatCacheAction\nfrom deepcompressor.dataset.cache import BaseCalibCacheLoader\nfrom deepcompressor.dataset.config import BaseDataLoaderConfig\n\nfrom ..nn.patch import RotaryEmbedding\nfrom ..nn.struct import LlmModelStruct, LlmTransformerBlockStruct\n\n__all__ = [\"LlmCalibDataLoaderConfig\", \"LlmCalibCacheLoader\"]\n\n\n@configclass\n@dataclass(kw_only=True)\nclass LlmCalibDataLoaderConfig(BaseDataLoaderConfig):\n \"\"\"Configuration for collecting calibration dataset for quantization.\n\n Args:\n data (`str`):\n Dataset name.\n num_samples (`int`):\n Number of dataset samples.\n path (`str`):\n Path to the dataset.\n seq_length (`int`):\n Sequence length of each sample.\n min_seq_length (`int`, *optional*, defaults to `0`):\n Minimum sequence length of each sample.\n max_seq_length (`int`, *optional*, defaults to `0`):\n Maximum sequence length of each sample.\n local_path (`str`, *optional*, defaults to `\"\"`):\n Local path to the dataset.\n \"\"\"\n\n path: str\n seq_length: int\n min_seq_length: int = 0\n max_seq_length: int = 0\n local_path: str = \"\"\n batch_size: int = field(init=False, default=1)\n\n def __post_init__(self) -> None:\n self.min_seq_length = max(0, self.min_seq_length)\n self.max_seq_length = max(0, self.max_seq_length)\n self.path = os.path.expanduser(self.path)\n self.local_path = os.path.expanduser(self.local_path)\n if os.path.exists(self.local_path):\n self.path = self.local_path\n\n def generate_dirnames(self, *, prefix: str = \"\", **kwargs) -> list[str]:\n \"\"\"Get the names of the configuration fields.\"\"\"\n name = f\"{self.data}.{self.num_samples}x{self.seq_length}.[{self.min_seq_length}-{self.max_seq_length}]\"\n return [f\"{prefix}.{name}\" if prefix else name]\n\n def build_dataset(self, tokenizer: PreTrainedTokenizer) -> \"LlmCalibDataset\":\n \"\"\"Build calibration dataset.\n\n Args:\n tokenizer (`PreTrainedTokenizer`):\n Tokenizer for encoding text.\n\n Returns:\n `LlmCalibDataset`:\n Calibration dataset.\n \"\"\"\n return LlmCalibDataset(\n tokenizer,\n data=self.data,\n path=self.path,\n num_samples=self.num_samples,\n seq_length=self.seq_length,\n max_seq_length=self.max_seq_length,\n min_seq_length=self.min_seq_length,\n )\n\n def build_loader(self, tokenizer: PreTrainedTokenizer) -> \"LlmCalibCacheLoader\":\n \"\"\"Build calibration data cache.\n\n Args:\n tokenizer (`PreTrainedTokenizer`):\n Tokenizer for encoding text.\n\n Returns:\n `LlmCalibDataCache`:\n Calibration data cache.\n \"\"\"\n return LlmCalibCacheLoader(config=self, tokenizer=tokenizer)\n\n\nclass LlmCalibDataset(torch.utils.data.Dataset):\n data: list[torch.Tensor]\n\n def __init__(\n self,\n tokenizer: PreTrainedTokenizer,\n data: str,\n path: str,\n num_samples: int,\n seq_length: int,\n max_seq_length: int = -1,\n min_seq_length: int = -1,\n seed: int = 42,\n ) -> None:\n assert num_samples > 0, \"num_samples should be positive\"\n assert seq_length > 0, \"seq_length should be positive\"\n num_tokens = num_samples * seq_length\n assert tokenizer is not None, \"tokenizer is required\"\n if data == \"pileval\":\n dataset = load_dataset(path, split=\"validation\")\n else:\n raise NotImplementedError(f\"Calibration dataset {data} is not supported\")\n dataset = dataset.shuffle(seed=seed)\n rng = random.Random(seed)\n seqs, toks = [], 0\n for sample in dataset:\n line = tokenizer.encode(sample[\"text\"].strip())\n length = len(line)\n if length == 0:\n continue\n if min_seq_length > 0 and length < min_seq_length:\n continue\n if max_seq_length > 0 and length > max_seq_length:\n continue\n # sample is a tensor of shape (seq_length, )\n seq = torch.tensor(line)\n if length > seq_length:\n tok = rng.randint(0, length - seq_length)\n seq = seq[tok : tok + seq_length]\n seqs.append(seq)\n toks += seq.numel()\n if len(seqs) >= num_samples and toks >= num_tokens:\n break\n # now concatenate all samples and split according to seq_length\n seqs = torch.cat(seqs).split(seq_length)\n if toks > num_tokens:\n seqs = seqs[:-1]\n seqs = seqs[:num_samples]\n self.data = seqs\n\n def __len__(self) -> int:\n return len(self.data)\n\n def __getitem__(self, idx: int) -> torch.Tensor:\n return self.data[idx]\n\n\nclass LlmCalibCacheLoader(BaseCalibCacheLoader):\n \"\"\"Cache for collecting calibration dataset for quantizing large language models.\"\"\"\n\n config: LlmCalibDataLoaderConfig\n dataset: LlmCalibDataset\n\n def __init__(self, config: LlmCalibDataLoaderConfig, tokenizer: PreTrainedTokenizer) -> None:\n \"\"\"Initialize large language model calibration cache loader.\n\n Args:\n config (`LlmCalibDataLoaderConfig`):\n Configuration for loading calibration dataset.\n tokenizer (`PreTrainedTokenizer`):\n Tokenizer for encoding text.\n \"\"\"\n super().__init__(dataset=config.build_dataset(tokenizer=tokenizer), batch_size=config.batch_size)\n self.batch_size = min(self.batch_size, len(self.dataset))\n self.config = config\n\n def _init_cache(self, name: str, module: nn.Module) -> IOTensorsCache:\n \"\"\"Initialize cache.\n\n Args:\n name (`str`):\n Module name.\n module (`nn.Module`):\n Module.\n\n Returns:\n `IOTensorsCache`:\n Input and output tensors cache.\n \"\"\"\n if isinstance(\n module, (nn.Linear, RotaryEmbedding, MixtralSparseMoeBlock, T5DenseActDense, T5DenseGatedActDense)\n ) or module.__class__.__name__.endswith((\"DecoderLayer\", \"Attention\", \"MLP\")):\n return IOTensorsCache(\n inputs=TensorCache(channels_dim=-1, reshape=LinearReshapeFn()),\n outputs=TensorCache(channels_dim=-1, reshape=LinearReshapeFn()),\n )\n else:\n super()._init_cache(name, module)\n\n def _convert_layer_inputs(\n self, m: nn.Module, args: tuple[tp.Any, ...], kwargs: dict[str, tp.Any], save_all: bool = False\n ) -> ModuleForwardInput:\n \"\"\"Convert layer inputs to module forward input.\n\n Args:\n m (`nn.Module`):\n Layer.\n args (`tuple[Any, ...]`):\n Layer input arguments.\n kwargs (`dict[str, Any]`):\n Layer input keyword arguments.\n save_all (`bool`, *optional*, defaults to `False`):\n Whether to save all inputs.\n\n Returns:\n `ModuleForwardInput`:\n Module forward input.\n \"\"\"\n x = args[0].detach().cpu() if save_all else MISSING\n return ModuleForwardInput(\n args=[x, *args[1:]], kwargs={k: None if isinstance(v, Cache) else v for k, v in kwargs.items()}\n )\n\n def iter_samples(self) -> tp.Generator[ModuleForwardInput, None, None]:\n \"\"\"Iterate over model input samples.\n\n Args:\n tokenizer (`nn.Module`):\n Tokenizer for encoding text.\n\n Yields:\n `ModuleForwardInput`:\n Module forward input.\n \"\"\"\n dataloader = torch.utils.data.DataLoader(\n self.dataset, batch_size=self.batch_size, shuffle=False, drop_last=True\n )\n for data in dataloader:\n yield ModuleForwardInput(args=(data,))\n\n def iter_layer_activations( # noqa: C901\n self,\n model: nn.Module | LlmModelStruct,\n *args,\n action: CacheAction | None = None,\n needs_inputs_fn: tp.Callable[[str, nn.Module], bool] | bool | None = True,\n needs_outputs_fn: tp.Callable[[str, nn.Module], bool] | bool | None = None,\n **kwargs,\n ) -> tp.Generator[\n tuple[\n str,\n tuple[\n LlmTransformerBlockStruct,\n dict[str, IOTensorsCache],\n dict[str, tp.Any],\n ],\n ],\n None,\n None,\n ]:\n \"\"\"Iterate over model activations for each layer.\n\n Args:\n model (`nn.Module`):\n Model.\n action (`CacheAction`, *optional*, defaults to `None`):\n Action for caching activations. If ``None``, ``ConcatCacheAction(\"cpu\")`` is used.\n needs_inputs_fn (`Callable[[str, nn.Module], bool]` or `bool` or `None`, *optional*, defaults to `True`):\n Function for determining whether to cache inputs for a module given its name and itself.\n needs_outputs_fn (`Callable[[str, nn.Module], bool]` or `bool` or `None`, *optional*, defaults to `None`):\n Function for determining whether to cache outputs for a module given its name and itself.\n *args: Arguments for ``_iter_samples``.\n **kwargs: Keyword arguments for ``_iter_samples``.\n\n Yields:\n Generator[\n tuple[str, tuple[LlmTransformerBlockStruct, dict[str, IOTensorsCache], dict[str, Any]]],\n None,\n None\n ]:\n Generator of tuple of\n - layer name\n - a tuple of\n - layer struct,\n - input and output caches for each module in the layer,\n - layer input keyword arguments.\n \"\"\"\n if isinstance(model, LlmModelStruct):\n model_struct = model\n model = model_struct.module\n else:\n model_struct = LlmModelStruct.construct(model)\n layers, layer_structs, recomputes, use_prev_layer_outputs = model_struct.get_iter_layer_activations_args()\n action = ConcatCacheAction(\"cpu\") if action is None else action\n for layer_idx, (layer_name, (layer, layer_cache, layer_inputs)) in enumerate(\n self._iter_layer_activations(\n model,\n *args,\n action=action,\n layers=layers,\n needs_inputs_fn=needs_inputs_fn,\n needs_outputs_fn=needs_outputs_fn,\n recomputes=recomputes,\n use_prev_layer_outputs=use_prev_layer_outputs,\n **kwargs,\n )\n ):\n layer_kwargs = layer_inputs[0].kwargs\n for layer_input in layer_inputs:\n for key, value in layer_input.kwargs.items():\n if isinstance(value, torch.Tensor):\n assert torch.equal(value, layer_kwargs[key])\n else:\n assert value == layer_kwargs[key]\n layer_struct = layer_structs[layer_idx]\n assert layer_name == layer_struct.name, f\"Expected {layer_struct.name}, got {layer_name}\"\n assert layer is layer_struct.module\n for transformer_block_struct in layer_struct.iter_transformer_block_structs():\n for attn_struct in transformer_block_struct.iter_attention_structs():\n if attn_struct.v_proj_name in layer_cache:\n cache = layer_cache[attn_struct.v_proj_name]\n layer_cache[attn_struct.q_proj_name] = cache\n layer_cache[attn_struct.k_proj_name] = cache\n ffn_struct = transformer_block_struct.ffn_struct\n up_proj_names = ffn_struct.up_proj_names\n if up_proj_names[0] in layer_cache:\n for expert_idx in range(ffn_struct.config.num_experts):\n cache = layer_cache[up_proj_names[expert_idx]]\n for name in up_proj_names[expert_idx :: ffn_struct.config.num_experts]:\n layer_cache[name] = cache\n if ffn_struct.config.num_experts == 1 and ffn_struct.name not in layer_cache:\n layer_cache[ffn_struct.name] = layer_cache[up_proj_names[0]]\n if ffn_struct.config.num_experts > 1 and ffn_struct.name in layer_cache:\n layer_cache[ffn_struct.moe_gate_name] = layer_cache[ffn_struct.name]\n yield layer_name, (layer_struct, layer_cache, layer_kwargs)\n"
},
{
"path": "deepcompressor/app/llm/quant/quantizer/__init__.py",
"content": "# -*- coding: utf-8 -*-\n\nfrom .config import LlmModuleQuantizerConfig\nfrom .quantizer import LlmActivationQuantizer, LlmWeightQuantizer\n"
},
{
"path": "deepcompressor/app/llm/quant/quantizer/config.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Quantizatizer config.\"\"\"\n\nimport typing as tp\nfrom dataclasses import dataclass, field\n\nimport torch\nfrom omniconfig import configclass\n\nfrom deepcompressor.calib.config import SkipBasedDynamicRangeCalibConfig\nfrom deepcompressor.data.dtype import QuantDataType\nfrom deepcompressor.quantizer.config import ProgressiveQuantizerConfig\nfrom deepcompressor.quantizer.kernel import QuantGptqConfig\nfrom deepcompressor.utils.config import EnableConfig, SkipBasedConfig\n\n__all__ = [\"LlmQuantizerConfig\", \"LlmWeightQuantizerConfig\", \"LlmActivationQuantizerConfig\", \"LlmModuleQuantizerConfig\"]\n\n\n@configclass\n@dataclass\nclass LlmQuantizerConfig(SkipBasedConfig, ProgressiveQuantizerConfig):\n \"\"\"Llm Quantizer Configuration.\n\n Args:\n dtype (`QuantDataType` or `None`, *optional*, defaults to `None`):\n The quantization data type.\n zero_point (`ZeroPointDomain` or `None`, *optional*, defaults to `None`):\n The zero-point domain.\n group_shapes (`Sequence[Sequence[int]]`, *optional*, defaults to `((-1, -1, -1),)`):\n The shapes for per-group quantization.\n scale_dtypes (`Sequence[torch.dtype | QuantDataType | None]`, *optional*, defaults to `(None,)`):\n The quantization scale data type for per-group quantization.\n intermediate_dtypes (`Sequence[QuantDataType]`, *optional*, defaults to `()`):\n The intermediate quantization data types.\n intermediate_levels (Sequence[int], *optional*, defaults to `()`):\n The intermediate quantization levels.\n needs_dequant_saturation (`bool`, *optional*, defaults to `False`):\n Whether the dequantization needs saturation.\n skips (`Sequence[str]`, *optional*, defaults to `[]`):\n The keys of the modules to skip.\n static (`bool`, *optional*, defaults to `False`):\n Whether to use static quantization.\n kernel_gptq (`QuantGptqConfig` or `None`, *optional*, defaults to `None`):\n The GPTQ kernel configuration.\n calib_range (`SkipBasedDynamicRangeCalibConfig` or `None`, *optional*, defaults to `None`):\n The dynamic range calibration configuration.\n \"\"\"\n\n static: bool = False\n kernel_gptq: QuantGptqConfig | None = None\n calib_range: SkipBasedDynamicRangeCalibConfig | None = None\n\n def __post_init__(self) -> None:\n super().__post_init__()\n if self.quant_dtype is None:\n self.static = False\n self.kernel_gptq = None\n self.calib_range = None\n if self.static and self.calib_range is None:\n self.calib_range = SkipBasedDynamicRangeCalibConfig()\n\n @property\n def enabled_gptq(self) -> bool:\n \"\"\"Whether quantization kernel calibration is enabled.\"\"\"\n return self.kernel_gptq is not None\n\n @property\n def enabled_calib_range(self) -> bool:\n \"\"\"Whether quantization dynamic range calibration is enabled.\"\"\"\n return self.calib_range is not None\n\n @property\n def needs_calib_data(self) -> bool:\n return self.enabled_calib_range and (self.calib_range.needs_search or self.static)\n\n def generate_calib_dirname(self) -> str:\n \"\"\"Generate the name for quantization calibration.\n\n Returns:\n str: The name.\n \"\"\"\n name = \"\"\n if self.static:\n name += \".static\"\n if self.enabled_gptq:\n name += \".gptq\"\n if self.enabled_calib_range and (self.calib_range.needs_search or self.calib_range.ratio != 1):\n name += \".range\"\n return name[1:] if name else \"\"\n\n\n@configclass\n@dataclass\nclass LlmWeightQuantizerConfig(LlmQuantizerConfig):\n \"\"\"Llm Weight Quantizer Configuration.\n\n Args:\n dtype (`QuantDataType` or `None`, *optional*, defaults to `None`):\n The quantization data type.\n zero_point (`ZeroPointDomain` or `None`, *optional*, defaults to `None`):\n The zero-point domain.\n group_shapes (`Sequence[Sequence[int]]`, *optional*, defaults to `((-1, -1, -1),)`):\n The shapes for per-group quantization.\n scale_dtypes (`Sequence[torch.dtype | QuantDataType | None]`, *optional*, defaults to `(None,)`):\n The quantization scale data type for per-group quantization.\n intermediate_dtypes (`Sequence[QuantDataType]`, *optional*, defaults to `()`):\n The intermediate quantization data types.\n intermediate_levels (Sequence[int], *optional*, defaults to `()`):\n The intermediate quantization levels.\n needs_dequant_saturation (`bool`, *optional*, defaults to `False`):\n Whether the dequantization needs saturation.\n skips (`Sequence[str]`, *optional*, defaults to `[]`):\n The keys of the modules to skip.\n kernel_gptq (`QuantGptqConfig` or `None`, *optional*, defaults to `None`):\n The GPTQ kernel configuration.\n calib_range (`SkipBasedDynamicRangeCalibConfig` or `None`, *optional*, defaults to `None`):\n The dynamic range calibration configuration.\n \"\"\"\n\n static: bool = field(init=False, default=True)\n\n\n@configclass\n@dataclass\nclass LlmActivationQuantizerConfig(LlmQuantizerConfig):\n \"\"\"Llm Activation quantization configuration.\n\n Args:\n dtype (`QuantDataType` or `None`, *optional*, defaults to `None`):\n The quantization data type.\n zero_point (`ZeroPointDomain` or `None`, *optional*, defaults to `None`):\n The zero-point domain.\n group_shapes (`Sequence[Sequence[int]]`, *optional*, defaults to `((-1, -1, -1),)`):\n The shapes for per-group quantization.\n scale_dtypes (`Sequence[torch.dtype | QuantDataType | None]`, *optional*, defaults to `(None,)`):\n The quantization scale data type for per-group quantization.\n skips (`Sequence[str]`, *optional*, defaults to `[]`):\n The keys of the modules to skip.\n static (`bool`, *optional*, defaults to `False`):\n Whether to use static quantization.\n calib_range (`SkipBasedDynamicRangeCalibConfig` or `None`, *optional*, defaults to `None`):\n The dynamic range calibration configuration.\n \"\"\"\n\n intermediate_dtypes: tp.Sequence[QuantDataType] = field(init=False, default=())\n intermediate_levels: tp.Sequence[int] = field(init=False, default=())\n needs_dequant_saturation: bool = field(init=False, default=False)\n kernel_gptq: None = field(init=False, default=None)\n\n\n@configclass\n@dataclass\nclass LlmModuleQuantizerConfig(EnableConfig):\n \"\"\"Llm Module quantization configuration.\n\n Args:\n wgts (`LlmWeightQuantizerConfig`):\n The weight quantization configuration.\n ipts (`LlmActivationQuantizerConfig`):\n The input activation quantization configuration.\n opts (`LlmActivationQuantizerConfig`):\n The output activation quantization configuration.\n \"\"\"\n\n wgts: LlmWeightQuantizerConfig\n ipts: LlmActivationQuantizerConfig\n opts: LlmActivationQuantizerConfig\n\n def is_enabled(self) -> bool:\n \"\"\"Whether the quantization is enabled.\"\"\"\n return self.enabled_wgts or self.enabled_ipts or self.enabled_opts\n\n @property\n def enabled_wgts(self) -> bool:\n \"\"\"Whether to enable weight quantization.\"\"\"\n return self.wgts is not None and self.wgts.is_enabled()\n\n @property\n def enabled_ipts(self) -> bool:\n \"\"\"Whether to enable activation quantization.\"\"\"\n return self.ipts is not None and self.ipts.is_enabled()\n\n @property\n def enabled_opts(self) -> bool:\n \"\"\"Whether to enable activation quantization.\"\"\"\n return self.opts is not None and self.opts.is_enabled()\n\n def generate_dirnames(\n self,\n *,\n prefix: str = \"\",\n shape: torch.Size | tuple[int, ...] = (1024, 1024, 16, 16),\n default_dtype: torch.dtype = torch.float16,\n **kwargs,\n ) -> list[str]:\n \"\"\"Get the directory names of the quantization configuration.\n\n Args:\n prefix (`str`, *optional*, defaults to `\"\"`):\n The prefix for the directory names.\n shape (`torch.Size` or `tuple[int, ...]`, *optional*, defaults to `(1024, 1024, 16, 16)`):\n The shape of the tensor to be quantized.\n default_dtype (`torch.dtype`, *optional*, defaults to `torch.float16`):\n The dtype of the tensor to be quantized.\n\n Returns:\n `list[str]`:\n The directory names of the quantization configuration.\n - The number of effective bits.\n - The name of the quantization data type.\n - The name of the group shapes.\n - The name of the modules to skip.\n \"\"\"\n wgts_names = self.wgts.generate_dirnames(prefix=\"w\", shape=shape, default_dtype=default_dtype)\n ipts_names = self.ipts.generate_dirnames(prefix=\"x\", shape=shape, default_dtype=default_dtype)\n opts_names = self.opts.generate_dirnames(prefix=\"y\", shape=shape, default_dtype=default_dtype)\n names = [\n f\"{wgts_name}-{ipts_name}-{opts_name}\"\n for wgts_name, ipts_name, opts_name in zip(wgts_names, ipts_names, opts_names, strict=True)\n ]\n if prefix:\n names = [f\"{prefix}.[{name}]\" for name in names]\n return names\n\n def generate_calib_dirname(self) -> str:\n \"\"\"Generate the name for quantization calibration.\n\n Returns:\n `str`:\n The name.\n \"\"\"\n name = \"\"\n if self.enabled_wgts:\n calib_name = self.wgts.generate_calib_dirname()\n if calib_name:\n name += f\"-w.{calib_name}\"\n if self.enabled_ipts:\n calib_name = self.ipts.generate_calib_dirname()\n if calib_name:\n name += f\"-x.{calib_name}\"\n if self.enabled_opts:\n calib_name = self.opts.generate_calib_dirname()\n if calib_name:\n name += f\"-y.{calib_name}\"\n return name[1:] if name else name\n"
},
{
"path": "deepcompressor/app/llm/quant/quantizer/quantizer.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Tensor Quantizer module.\"\"\"\n\nimport typing as tp\nfrom dataclasses import dataclass, field\n\nimport torch\nimport torch.nn as nn\n\nfrom deepcompressor.calib.range import calibrate_dynamic_range\nfrom deepcompressor.data.cache import TensorsCache\nfrom deepcompressor.data.common import TensorType\nfrom deepcompressor.data.range import DynamicRange\nfrom deepcompressor.quantizer.kernel import QuantGptqConfig\nfrom deepcompressor.quantizer.processor import Quantizer\n\nfrom .config import LlmActivationQuantizerConfig, LlmQuantizerConfig, LlmWeightQuantizerConfig\n\n__all__ = [\"LlmQuantizer\", \"LlmWeightQuantizer\", \"LlmActivationQuantizer\"]\n\n\n@dataclass\nclass LlmQuantizer(Quantizer):\n \"\"\"Llm quantizer class.\n\n Args:\n config (`LlmQuantizerConfig`):\n The quantizer configuration.\n key (`str`, *optional*, defaults to `\"\"`):\n The key of the quantizer.\n tensor_type (`TensorType`, *optional*, defaults to `TensorType.Weights`):\n The type of the tensor to quantize.\n channels_dim (`int` or `None`, *optional*, defaults to `None`):\n The dimension of channels.\n scale (`torch.Tensor` or `Sequence[torch.Tensor]` or `None`, *optional*, defaults to `None`):\n The scale tensor.\n zero (`torch.Tensor` or `None`, *optional*, defaults to `None`):\n The zero point tensor.\n dynamic_range (`DynamicRange` or `Sequence[DynamicRange]` or `None`, *optional*, defaults to `None`):\n The dynamic range.\n range_bound (`RangeBound` or `None`, *optional*, defaults to `None`):\n The dynamic range bound.\n quant_range (`QuantRange` or `None`, *optional*, defaults to `None`):\n The quantization range.\n default_dtype (`torch.dtype` or `None`, *optional*, defaults to `None`):\n The default scale dtype\n develop_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The quantization development dtype.\n\n low_rank (`QuantLowRankConfig` or `None`, *optional*, defaults to `None`):\n The quantization low-rank branch configuration.\n input_packager (`BaseInputPackager` or `None`, *optional*, defaults to `None`):\n The input packager, used for unpacking and repacking the input tensor(s).\n output_packager (`BaseOutputPackager` or `None`, *optional*, defaults to `None`):\n The output packager, used for unpacking and repacking the output tensor(s).\n \"\"\"\n\n config: LlmQuantizerConfig\n kernel: QuantGptqConfig | None = field(init=False)\n tensor_type: TensorType = TensorType.Weights\n\n def __post_init__(self) -> None:\n self.kernel = self.config.kernel_gptq\n\n def calibrate_dynamic_range(\n self,\n modules: tp.Sequence[nn.Module],\n activations: TensorsCache,\n weights: tp.Sequence[nn.Parameter] = None,\n eval_inputs: TensorsCache | None = None,\n eval_module: nn.Module | None = None,\n eval_kwargs: dict[str, tp.Any] | None = None,\n orig_weights: tp.Sequence[tuple[nn.Parameter, torch.Tensor]] | None = None,\n orig_activations: TensorsCache | None = None,\n orig_eval_inputs: TensorsCache | None = None,\n ) -> tp.Sequence[DynamicRange] | None:\n \"\"\"Calibrate the dynamic range.\n\n Args:\n modules (`Sequence[nn.Module]`):\n The modules to calibrate.\n activations (`TensorsCache`):\n The inputs cache if the tensor type is not outputs, or the outputs cache if the tensor type is outputs.\n weights (`Sequence[nn.Parameter]` or `None`, *optional*, defaults to `None`):\n The weights to calibrate.\n If not provided, the weights of the modules will be used.\n eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The cache of the inputs for evaluation.\n If not provided, the `activations` cache will be used.\n eval_module (`nn.Module` or `None`, *optional*, defaults to `None`):\n The module to evaluate the quantization error.\n If not provided, the module to calibrate will be used.\n eval_kwargs (`dict[str, tp.Any]` or `None`, *optional*, defaults to `None`):\n The keyword arguments for evaluation.\n orig_weights (`Sequence[tuple[nn.Parameter, torch.Tensor]]` or `None`, *optional*, defaults to `None`):\n The original weights.\n orig_activations (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The original activations.\n orig_eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The original evaluation inputs.\n\n Returns:\n `Sequence[DynamicRange]` or `None`:\n The dynamic ranges of each quantization step.\n \"\"\"\n if (\n not self.is_enabled()\n or self.config.calib_range is None\n or not self.config.calib_range.is_enabled_for(self.key)\n ):\n self.dynamic_range = None\n else:\n self.dynamic_range = calibrate_dynamic_range(\n tensor_type=self.tensor_type,\n config=self.config.calib_range,\n static=self.config.static,\n quantizer=self,\n modules=modules,\n activations=activations,\n weights=weights,\n eval_inputs=eval_inputs,\n eval_module=eval_module,\n eval_kwargs=eval_kwargs,\n orig_weights=orig_weights,\n orig_activations=orig_activations,\n orig_eval_inputs=orig_eval_inputs,\n )\n return self.dynamic_range\n\n\n@dataclass\nclass LlmWeightQuantizer(LlmQuantizer):\n \"\"\"Llm Weight Quantizer class.\n\n Args:\n config (`LlmWeightQuantizerConfig`):\n The quantizer configuration.\n key (`str`, *optional*, defaults to `\"\"`):\n The key of the quantizer.\n scale (`torch.Tensor` or `Sequence[torch.Tensor]` or `None`, *optional*, defaults to `None`):\n The scale tensor.\n zero (`torch.Tensor` or `None`, *optional*, defaults to `None`):\n The zero point tensor.\n dynamic_range (`DynamicRange` or `Sequence[DynamicRange]` or `None`, *optional*, defaults to `None`):\n The dynamic range.\n range_bound (`RangeBound` or `None`, *optional*, defaults to `None`):\n The dynamic range bound.\n quant_range (`QuantRange` or `None`, *optional*, defaults to `None`):\n The quantization range.\n default_dtype (`torch.dtype` or `None`, *optional*, defaults to `None`):\n The default scale dtype\n develop_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The quantization development dtype.\n\n low_rank (`QuantLowRankConfig` or `None`, *optional*, defaults to `None`):\n The quantization low-rank branch configuration.\n input_packager (`BaseInputPackager` or `None`, *optional*, defaults to `None`):\n The input packager, used for unpacking and repacking the input tensor(s).\n output_packager (`BaseOutputPackager` or `None`, *optional*, defaults to `None`):\n The output packager, used for unpacking and repacking the output tensor(s).\n \"\"\"\n\n config: LlmWeightQuantizerConfig\n channels_dim: None = field(init=False, default=None)\n tensor_type: TensorType = field(init=False, default=TensorType.Weights)\n\n def calibrate_dynamic_range(\n self,\n module: nn.Module,\n inputs: TensorsCache,\n weight: nn.Parameter | None = None,\n eval_inputs: TensorsCache | None = None,\n eval_module: nn.Module | None = None,\n eval_kwargs: dict[str, tp.Any] | None = None,\n orig_inputs: TensorsCache | None = None,\n orig_eval_inputs: TensorsCache | None = None,\n ) -> DynamicRange | tuple[DynamicRange, ...]:\n \"\"\"Calibrate the dynamic range.\n\n Args:\n module (`nn.Module`):\n The module to calibrate.\n inputs (`TensorsCache`):\n The inputs cache.\n weight (`nn.Parameter` or `None`, *optional*, defaults to `None`):\n The weight parameter to calibrate.\n If not provided, the weight of the `module` will be used.\n eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The cache of the inputs for evaluation.\n If not provided, the `activations` cache will be used.\n eval_module (`nn.Module` or `None`, *optional*, defaults to `None`):\n The module to evaluate the quantization error.\n If not provided, the module to calibrate will be used.\n eval_kwargs (`dict[str, tp.Any]` or `None`, *optional*, defaults to `None`):\n The keyword arguments for evaluation.\n orig_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The original inputs.\n orig_eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The original evaluation inputs.\n\n Returns:\n `Sequence[DynamicRange]` or `None`:\n The dynamic ranges of each quantization step.\n \"\"\"\n return super().calibrate_dynamic_range(\n modules=[module],\n weights=[weight] if weight is not None else [module.weight],\n activations=inputs,\n eval_inputs=eval_inputs,\n eval_module=eval_module,\n eval_kwargs=eval_kwargs,\n orig_activations=orig_inputs,\n orig_eval_inputs=orig_eval_inputs,\n )\n\n\n@dataclass\nclass LlmActivationQuantizer(LlmQuantizer):\n \"\"\"Llm Activation Quantizer class.\n\n Args:\n config (`LlmActivationQuantizerConfig`):\n The quantizer configuration.\n key (`str`, *optional*, defaults to `\"\"`):\n The key of the quantizer.\n tensor_type (`TensorType`, *optional*, defaults to `TensorType.Inputs`):\n The type of the tensor to quantize.\n channels_dim (`int` or `None`, *optional*, defaults to `None`):\n The dimension of channels.\n scale (`torch.Tensor` or `Sequence[torch.Tensor]` or `None`, *optional*, defaults to `None`):\n The scale tensor.\n zero (`torch.Tensor` or `None`, *optional*, defaults to `None`):\n The zero point tensor.\n dynamic_range (`DynamicRange` or `Sequence[DynamicRange]` or `None`, *optional*, defaults to `None`):\n The dynamic range.\n range_bound (`RangeBound` or `None`, *optional*, defaults to `None`):\n The dynamic range bound.\n quant_range (`QuantRange` or `None`, *optional*, defaults to `None`):\n The quantization range.\n default_dtype (`torch.dtype` or `None`, *optional*, defaults to `None`):\n The default scale dtype\n develop_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The quantization development dtype.\n\n low_rank (`QuantLowRankConfig` or `None`, *optional*, defaults to `None`):\n The quantization low-rank branch configuration.\n input_packager (`BaseInputPackager` or `None`, *optional*, defaults to `None`):\n The input packager, used for unpacking and repacking the input tensor(s).\n output_packager (`BaseOutputPackager` or `None`, *optional*, defaults to `None`):\n The output packager, used for unpacking and repacking the output tensor(s).\n \"\"\"\n\n config: LlmActivationQuantizerConfig\n tensor_type: TensorType = TensorType.Inputs\n\n def __post_init__(self) -> None:\n super().__post_init__()\n assert self.tensor_type != TensorType.Weights, \"The tensor type cannot be weights.\"\n assert isinstance(self.channels_dim, int), \"The channels dimension must be provided.\"\n"
},
{
"path": "deepcompressor/app/llm/quant/reorder.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"LLM quantization channel reordering module.\"\"\"\n\nimport gc\nimport typing as tp\n\nimport torch\nimport torch.nn as nn\nimport torch.utils\nfrom tqdm import tqdm\nfrom transformers import PreTrainedTokenizer\n\nfrom deepcompressor.calib.reorder import ChannelOrderCalibrator, ChannelReorderer\nfrom deepcompressor.data.cache import IOTensorsCache, TensorCache, TensorsCache\nfrom deepcompressor.quantizer.processor import Quantizer\nfrom deepcompressor.utils import tools\n\nfrom ..nn import LlmModelStruct, LlmTransformerBlockStruct\nfrom .config import LlmQuantConfig\nfrom .utils import get_needs_inputs_fn\n\n__all__ = [\"reorder_llm\"]\n\n\ndef _extend_params_(\n params: list[tuple[nn.Parameter, int]],\n modules: list[nn.Linear | nn.Embedding, nn.LayerNorm],\n out_channels_dim: int | None = None,\n in_channels_dim: int | None = None,\n) -> list[tuple[nn.Parameter, int]]:\n \"\"\"Extend the parameters to be reordered.\"\"\"\n if out_channels_dim is not None:\n assert in_channels_dim is None\n else:\n assert in_channels_dim is not None\n for module in modules:\n if module is None:\n continue\n if out_channels_dim is not None:\n params.append((module.weight, out_channels_dim))\n if hasattr(module, \"bias\") and module.bias is not None:\n params.append((module.bias, 0))\n else:\n params.append((module.weight, in_channels_dim))\n return params\n\n\n@torch.inference_mode()\ndef reorder_llm_layer( # noqa: C901\n layer: LlmTransformerBlockStruct,\n config: LlmQuantConfig,\n reorder_cache: dict[str, torch.Tensor],\n residual_calibrator: ChannelOrderCalibrator | None = None,\n layer_cache: dict[str, IOTensorsCache] | None = None,\n layer_kwargs: dict[str, tp.Any] | None = None,\n) -> ChannelOrderCalibrator | None:\n \"\"\"Calibrate the channel order in a layer.\n\n Args:\n layer (`LlmTransformerBlockStruct`):\n Large language model layer to be reordered.\n config (`LlmQuantConfig`):\n Quantization config.\n reorder_cache (`dict[str, torch.Tensor]`):\n Reorder indexes cache.\n residual_calibrator (`ChannelOrderCalibrator` or `None`, *optional*, defaults to `None`):\n Channel order calibrator for residual modules.\n layer_cache (`dict[str, IOTensorsCache]`, *optional*, defaults to `None`):\n Layer activations cache.\n layer_kwargs (`dict[str, tp.Any]`, *optional*, defaults to `None`):\n Layer keyword arguments.\n\n Returns:\n `ChannelOrderCalibrator` or `None`:\n Channel order calibrator for residual modules.\n \"\"\"\n logger = tools.logging.getLogger(f\"{__name__}.Reorder\")\n layer_cache = layer_cache or {}\n\n attn = layer.attn_struct\n qkv_proj, out_proj = attn.qkv_proj, attn.out_proj\n num_heads, num_head_repeats = attn.config.num_query_heads, attn.config.num_head_repeats\n # region reorder in attention module\n if config.reorder.dynamic and config.reorder.is_enabled_for(attn.qkv_proj_key):\n logger.debug(\"- Reordering %s\", attn.qkv_proj_names)\n cache_key = attn.name\n if cache_key not in reorder_cache:\n index = ChannelOrderCalibrator(\n config=config.reorder,\n weight_quantizer=Quantizer(config.wgts, key=attn.qkv_proj_key),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=attn.qkv_proj_key),\n develop_dtype=config.develop_dtype,\n ).calibrate(\n x_wgts=[m.weight for m in qkv_proj],\n x_acts=layer_cache[attn.v_proj_name].inputs if layer_cache else None,\n x_mods=qkv_proj,\n eval_inputs=layer_cache[cache_key].inputs if layer_cache else None,\n eval_module=attn.module,\n eval_kwargs=attn.filter_kwargs(layer_kwargs),\n reorder_wgts=[(m.weight, 1) for m in qkv_proj],\n reorder_ipt_mods=[(attn.module, -1, None)],\n reorder_opt_mods=[],\n )\n reorder_cache[cache_key] = index.to(device=torch.device(\"cpu\"))\n index = reorder_cache[cache_key]\n for proj in qkv_proj:\n index = index.to(proj.weight.device)\n proj.weight.data = proj.weight.data.index_select(1, index)\n ChannelReorderer(index, channels_dim=-1).as_hook().register(attn.module)\n gc.collect()\n torch.cuda.empty_cache()\n if config.reorder.is_enabled_for(attn.out_proj_key):\n logger.debug(\"- Reordering %s\", attn.out_proj_name)\n cache_key = attn.out_proj_name\n if cache_key not in reorder_cache:\n index = ChannelOrderCalibrator(\n config=config.reorder,\n weight_quantizer=Quantizer(config.wgts, key=attn.out_proj_key),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=attn.out_proj_key),\n num_heads=num_heads,\n num_head_repeats=num_head_repeats,\n develop_dtype=config.develop_dtype,\n ).calibrate(\n x_wgts=[out_proj.weight],\n x_acts=layer_cache[cache_key].inputs if layer_cache else None,\n x_mods=[out_proj],\n eval_inputs=layer_cache[cache_key].inputs if layer_cache else None,\n eval_module=out_proj,\n reorder_wgts=[(out_proj.weight, 1)],\n reorder_ipt_mods=[(out_proj, -1, None)],\n reorder_opt_mods=[],\n )\n reorder_cache[cache_key] = index.to(device=torch.device(\"cpu\"))\n index = reorder_cache[cache_key]\n index = index.to(out_proj.weight.device)\n out_proj.weight.data = out_proj.weight.data.index_select(1, index)\n v_proj = qkv_proj[2]\n if num_heads > 1 and num_head_repeats > 1:\n num_channels = index.numel()\n head_channels = num_channels // num_heads\n index = index.view(num_heads, head_channels)\n delta = torch.arange(0, num_channels, head_channels, device=index.device).view(num_heads, 1)\n index = index - delta\n num_v_channels = num_channels // num_head_repeats\n num_v_heads = num_heads // num_head_repeats\n index = index.view(num_v_heads, num_head_repeats, head_channels)[:, 0, :]\n delta = torch.arange(0, num_v_channels, head_channels, device=index.device).view(num_v_heads, 1)\n index = index + delta\n index = index.view(-1)\n v_proj.weight.data = v_proj.weight.data.index_select(0, index.to(v_proj.weight.device))\n if v_proj.bias is not None:\n v_proj.bias.data = v_proj.bias.data[index.to(v_proj.bias.device)].contiguous()\n gc.collect()\n torch.cuda.empty_cache()\n # endregion\n ffn = layer.ffn_struct\n num_experts = ffn.config.num_experts\n up_proj, down_proj = ffn.up_projs, ffn.down_projs\n # region reorder in feed-forward module\n if config.reorder.dynamic and config.reorder.is_enabled_for(ffn.up_proj_key):\n logger.debug(\"- Reordering %s\", ffn.name)\n cache_key = ffn.name\n if cache_key not in reorder_cache:\n index = ChannelOrderCalibrator(\n config=config.reorder,\n weight_quantizer=Quantizer(config.wgts, key=ffn.up_proj_key),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=ffn.up_proj_key),\n develop_dtype=config.develop_dtype,\n ).calibrate(\n x_wgts=[m.weight for m in up_proj],\n x_acts=layer_cache[cache_key].inputs if layer_cache else None,\n x_mods=up_proj,\n eval_inputs=layer_cache[cache_key].inputs if layer_cache else None,\n eval_module=ffn.module,\n reorder_wgts=[(m.weight, 1) for m in up_proj],\n reorder_ipt_mods=[(ffn.module, -1, None)],\n reorder_opt_mods=[],\n )\n reorder_cache[cache_key] = index.to(device=torch.device(\"cpu\"))\n index = reorder_cache[cache_key]\n index = index.to(device=up_proj[0].weight.device)\n for fc in up_proj:\n fc.weight.data = fc.weight.data.index_select(1, index.to(fc.weight.device))\n moe_gate = ffn.moe_gate\n if moe_gate is not None:\n moe_gate.weight.data = moe_gate.weight.data.index_select(1, index.to(moe_gate.weight.device))\n ChannelReorderer(index, channels_dim=-1).as_hook().register(ffn.module)\n if config.reorder.is_enabled_for(ffn.down_proj_key):\n for expert_idx, (fc2_name, fc2) in enumerate(zip(ffn.down_proj_names, down_proj, strict=True)):\n logger.debug(\"- Reordering module %s\", fc2_name)\n cache_key = fc2_name\n if cache_key not in reorder_cache:\n index = ChannelOrderCalibrator(\n config=config.reorder,\n weight_quantizer=Quantizer(config.wgts, key=ffn.down_proj_key),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=ffn.down_proj_key),\n develop_dtype=config.develop_dtype,\n ).calibrate(\n x_wgts=[fc2.weight],\n x_acts=layer_cache[cache_key].inputs if layer_cache else None,\n x_mods=[fc2],\n eval_inputs=layer_cache[cache_key].inputs if layer_cache else None,\n eval_module=fc2,\n reorder_wgts=[(fc2.weight, 1)],\n reorder_ipt_mods=[(fc2, -1, None)],\n reorder_opt_mods=[],\n )\n reorder_cache[cache_key] = index.to(device=torch.device(\"cpu\"))\n index = reorder_cache[cache_key]\n index = index.to(fc2.weight.device)\n fc2.weight.data = fc2.weight.data.index_select(1, index.to(fc2.weight.device))\n for fc1 in up_proj[expert_idx::num_experts]:\n fc1.weight.data = fc1.weight.data.index_select(0, index.to(fc1.weight.device))\n if fc1.bias is not None:\n fc1.bias.data = fc1.bias.data[index.to(fc1.bias.device)].contiguous()\n gc.collect()\n torch.cuda.empty_cache()\n # endregion\n\n if residual_calibrator is not None and (\n config.reorder.dynamic\n or not config.reorder.is_enabled_for(attn.qkv_proj_key)\n or not config.reorder.is_enabled_for(ffn.up_proj_key)\n ):\n residual_calibrator = None\n if residual_calibrator is not None and \"residual\" not in reorder_cache:\n residual_calibrator.update_channel_metrics(\n weights=[m.weight for m in qkv_proj],\n inputs=layer_cache[attn.v_proj_name].inputs if layer_cache else None,\n )\n for expert_idx in range(num_experts):\n residual_calibrator.update_channel_metrics(\n weights=[m.weight for m in up_proj[expert_idx::num_experts]],\n inputs=layer_cache[ffn.up_proj_names[expert_idx]].inputs if layer_cache else None,\n )\n return residual_calibrator\n\n\n@torch.inference_mode()\ndef reorder_llm( # noqa: C901\n model: nn.Module | LlmModelStruct,\n config: LlmQuantConfig,\n tokenizer: PreTrainedTokenizer | None = None,\n reorder_cache: dict[str, torch.Tensor] | None = None,\n) -> dict[str, torch.Tensor]:\n \"\"\"Quantize the large foundation model weights.\n\n Args:\n model (`nn.Module` or `LlmStruct`):\n Model to be reordered.\n config (`LlmQuantConfig`):\n Quantization config.\n tokenizer (`PreTrainedTokenizer` or `None`, *optional*, defaults to `None`):\n Tokenizer.\n reorder_cache (`dict[str, torch.Tensor]`, *optional*, defaults to `None`):\n Reorder indexes cache.\n\n Returns:\n `dict[str, torch.Tensor]`:\n Reorder indexes cache.\n \"\"\"\n if not isinstance(model, LlmModelStruct):\n model = LlmModelStruct.construct(model)\n assert isinstance(model, LlmModelStruct)\n logger = tools.logging.getLogger(f\"{__name__}.Reorder\")\n reorder_cache = {} if reorder_cache is None else reorder_cache\n residual_calibrator = None\n if \"residual\" not in reorder_cache and not config.reorder.dynamic and config.reorder.is_enabled_for(\"residual\"):\n residual_calibrator = ChannelOrderCalibrator(\n config=config.reorder,\n weight_quantizer=Quantizer(config.wgts),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1),\n develop_dtype=config.develop_dtype,\n )\n with tools.logging.redirect_tqdm():\n if not reorder_cache:\n calib_cache = config.calib.build_loader(tokenizer)\n for _, (layer, layer_cache, layer_kwargs) in tqdm(\n calib_cache.iter_layer_activations(\n model,\n needs_inputs_fn=get_needs_inputs_fn(model=model, config=config),\n ),\n desc=\"reordering\",\n leave=False,\n total=len(model.backbone_struct.layer_structs),\n dynamic_ncols=True,\n ):\n residual_calibrator = reorder_llm_layer(\n layer=layer,\n config=config,\n reorder_cache=reorder_cache,\n residual_calibrator=residual_calibrator,\n layer_cache=layer_cache,\n layer_kwargs=layer_kwargs,\n )\n gc.collect()\n torch.cuda.empty_cache()\n else:\n calib_cache = None\n for layer in tqdm(\n model.backbone_struct.layer_structs,\n desc=\"reordering\",\n leave=False,\n dynamic_ncols=True,\n ):\n residual_calibrator = reorder_llm_layer(\n layer=layer,\n config=config,\n reorder_cache=reorder_cache,\n residual_calibrator=residual_calibrator,\n )\n if residual_calibrator is None:\n return reorder_cache\n # region add extra params to be reordered\n backbone = model.backbone_struct\n x_mods: list[nn.Linear] = []\n reorder_wgts: list[tuple[nn.Parameter, int]] = []\n for layer in backbone.layer_structs:\n x_mods.extend(layer.attn_struct.qkv_proj)\n x_mods.extend(layer.ffn_struct.up_projs)\n _extend_params_(\n reorder_wgts,\n [\n layer.pre_attn_norm,\n layer.attn_struct.out_proj,\n layer.post_attn_norm,\n layer.pre_ffn_norm,\n *layer.ffn_struct.down_projs,\n layer.post_ffn_norm,\n ],\n out_channels_dim=0,\n )\n _extend_params_(reorder_wgts, [layer.ffn_struct.moe_gate], in_channels_dim=1)\n need_reorder_head = model.head is not None\n if backbone.proj_in is not None:\n _extend_params_(reorder_wgts, [backbone.proj_in], out_channels_dim=0)\n _extend_params_(reorder_wgts, [backbone.embed_positions], out_channels_dim=1)\n else:\n _extend_params_(reorder_wgts, [backbone.embed_tokens, backbone.embed_positions], out_channels_dim=1)\n _extend_params_(reorder_wgts, [backbone.norm_in, backbone.norm_out], out_channels_dim=0)\n if backbone.proj_out is not None:\n _extend_params_(reorder_wgts, [backbone.proj_out], in_channels_dim=1)\n need_reorder_head = False\n logger.debug(\"- Reordering residual modules\")\n _extend_params_(reorder_wgts, x_mods, in_channels_dim=1)\n if \"residual\" not in reorder_cache:\n calib_cache = calib_cache or config.calib.build_loader(tokenizer)\n residual_calibrator.init_channel_indexes()\n index = residual_calibrator.calibrate(\n x_wgts=[m.weight for m in x_mods],\n x_acts=None,\n eval_inputs=TensorsCache(TensorCache(calib_cache.dataset.data, channels_dim=-1, orig_device=\"cuda\")),\n eval_module=model.backbone,\n x_mods=x_mods,\n reorder_wgts=reorder_wgts,\n reorder_ipt_mods=[],\n reorder_opt_mods=[(model.backbone, -1, None)] if need_reorder_head else [],\n )\n reorder_cache[\"residual\"] = index.to(device=torch.device(\"cpu\"))\n del x_mods, residual_calibrator, calib_cache\n gc.collect()\n torch.cuda.empty_cache()\n index = reorder_cache[\"residual\"]\n for wgt, dim in reorder_wgts:\n wgt.data = wgt.data.index_select(dim=dim, index=index.to(wgt.data.device))\n if need_reorder_head and not model.config.tie_word_embeddings:\n model.head.weight.data = model.head.weight.data.index_select(dim=1, index=index.to(model.head.weight.device))\n gc.collect()\n torch.cuda.empty_cache()\n return reorder_cache\n"
},
{
"path": "deepcompressor/app/llm/quant/rotate.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Large Language Model Rotation module.\"\"\"\n\nimport gc\n\nimport torch\nfrom tqdm import tqdm\nfrom transformers import PreTrainedModel\n\nfrom deepcompressor.calib.config import QuantRotationConfig\nfrom deepcompressor.calib.rotate import (\n get_rotation_matrix,\n hadamard_in_channels,\n rotate_in_channels,\n rotate_out_channels,\n transform_norm_and_linear,\n)\nfrom deepcompressor.utils import tools\n\nfrom ..nn import LlmModelStruct\n\n__all__ = [\"rotate_llm\"]\n\n\n@torch.inference_mode()\ndef rotate_llm( # noqa: C901\n model: PreTrainedModel | LlmModelStruct,\n /,\n config: QuantRotationConfig,\n rotation: torch.Tensor | None = None,\n) -> torch.Tensor:\n \"\"\"Rotate the weights of the large language model.\n\n Args:\n model (`PreTrainedModel` or `LlmStruct`):\n Model to be rotated.\n config (`QuantRotationConfig`):\n Rotation configuration.\n rotation (`torch.Tensor` or `None`, *optional*, defaults to `None`):\n Rotation matrix.\n\n Returns:\n `torch.Tensor`:\n The rotation matrix.\n \"\"\"\n if not isinstance(model, LlmModelStruct):\n model = LlmModelStruct.construct(model)\n assert isinstance(model, LlmModelStruct)\n devices: list[torch.device] = []\n dtypes: list[torch.dtype] = []\n linears: list[torch.nn.Linear] = []\n size: float = 0\n for m in model.module.modules():\n if isinstance(m, torch.nn.Linear):\n devices.append(m.weight.device)\n dtypes.append(m.weight.dtype)\n linears.append(m)\n size += m.weight.numel() / 1e9\n for linear in linears:\n linear.to(dtype=torch.float32, device=\"cpu\" if size > 30 else None)\n\n for block in model.iter_transformer_block_structs():\n assert not block.post_attn_norms, \"Rotation is only supported for models without post-attention norms.\"\n assert not block.post_ffn_norm, \"Rotation is only supported for models without post-FFN norms.\"\n\n logger = tools.logging.getLogger(f\"{__name__}.Rotate\")\n backbone = model.backbone_struct\n layers = backbone.layer_structs\n # region transform norm and linear\n if backbone.norm_in is None:\n if backbone.proj_in is None:\n prev_modules = [backbone.embed_tokens]\n prev_out_channels_dims = 1\n if backbone.embed_positions is not None:\n prev_modules.append(backbone.embed_positions)\n elif backbone.embed_positions is None:\n prev_modules = [backbone.proj_in]\n prev_out_channels_dims = 0\n else:\n prev_modules = [backbone.proj_in, backbone.embed_positions]\n prev_out_channels_dims = [0, 1]\n else:\n prev_modules = [backbone.norm_in]\n prev_out_channels_dims = 0\n with tools.logging.redirect_tqdm():\n for layer in tqdm(layers, desc=\"Transforming norm and linear\", dynamic_ncols=True):\n logger.debug(f\"- Transforming norm and linear in {layer.name}\")\n transform_norm_and_linear(\n parent=layer.module,\n norm_name=layer.pre_attn_norm_rname,\n next_modules=layer.attn_struct.qkv_proj,\n prev_modules=prev_modules,\n prev_out_channels_dims=prev_out_channels_dims,\n )\n prev_modules = [layer.attn_struct.out_proj]\n prev_out_channels_dims = 0\n transform_norm_and_linear(\n parent=layer.module,\n norm_name=layer.pre_ffn_norm_rname,\n next_modules=layer.ffn_struct.up_projs\n + ([layer.ffn_struct.moe_gate] if layer.ffn_struct.moe_gate is not None else []),\n prev_modules=prev_modules,\n prev_out_channels_dims=prev_out_channels_dims,\n )\n prev_modules = layer.ffn_struct.down_projs\n prev_out_channels_dims = 0\n gc.collect()\n torch.cuda.empty_cache()\n logger.debug(f\"- Transforming {backbone.norm_out_name}\")\n transform_norm_and_linear(\n parent=backbone.module,\n norm_name=backbone.norm_out_rname,\n next_modules=[model.head if backbone.proj_out is None else backbone.proj_out],\n prev_modules=prev_modules,\n prev_out_channels_dims=prev_out_channels_dims,\n )\n # endregion\n if rotation is None:\n rotation = get_rotation_matrix(backbone.config.num_channels, random=config.random)\n # region rotate embeddings\n if backbone.proj_in is None:\n logger.debug(f\"- Rotating {backbone.embed_tokens_name}\")\n weight = backbone.embed_tokens.weight\n rotation = rotation.to(weight.device)\n rotate_in_channels(weight, rotation=rotation)\n else:\n logger.debug(f\"- Rotating {backbone.proj_in_name} (out)\")\n weight = backbone.proj_in.weight\n rotation = rotation.to(weight.device)\n rotate_out_channels(weight, rotation=rotation, bias=backbone.proj_in.bias)\n if backbone.embed_positions is not None:\n logger.debug(f\"- Rotating {backbone.embed_positions_name}\")\n weight = backbone.embed_positions.weight\n rotation = rotation.to(weight.device)\n rotate_in_channels(weight, rotation=rotation)\n # endregion\n down_proj = []\n # region rotate backbone layers\n head_rotation = get_rotation_matrix(model.config.num_head_channels, random=config.random)\n with tools.logging.redirect_tqdm():\n for layer in tqdm(layers, desc=\"Rotating backbone layers\", dynamic_ncols=True):\n logger.debug(f\"- Rotating {layer.name}\")\n tools.logging.Formatter.indent_inc()\n attn, ffn = layer.attn_struct, layer.ffn_struct\n for proj_name, proj in zip(attn.qkv_proj_names, attn.qkv_proj, strict=True):\n logger.debug(f\"- Rotating {proj_name} (in)\")\n rotation = rotation.to(proj.weight.device)\n rotate_in_channels(proj.weight, rotation=rotation)\n logger.debug(f\"- Rotating {attn.out_proj_name} (out)\")\n rotation = rotation.to(attn.out_proj.weight.device)\n rotate_out_channels(attn.out_proj.weight, rotation=rotation, bias=attn.out_proj.bias)\n if attn.out_proj_key in config.transforms:\n logger.debug(f\"- Rotating {attn.v_proj_name} (out)\")\n rotate_out_channels(attn.v_proj.weight, rotation=head_rotation, bias=attn.v_proj.bias)\n logger.debug(f\"- Rotating {attn.o_proj_name} (in)\")\n rotate_in_channels(attn.o_proj.weight, rotation=head_rotation)\n for fc_name, fc in zip(ffn.up_proj_names, ffn.up_projs, strict=True):\n logger.debug(f\"- Rotating {fc_name} (in)\")\n rotation = rotation.to(fc.weight.device)\n rotate_in_channels(fc.weight, rotation=rotation)\n if ffn.moe_gate is not None:\n logger.debug(f\"- Rotating {ffn.moe_gate_name} (in)\")\n rotation = rotation.to(ffn.moe_gate.weight.device)\n rotate_in_channels(ffn.moe_gate.weight, rotation=rotation)\n for fc_name, fc in zip(ffn.down_proj_names, ffn.down_projs, strict=True):\n logger.debug(f\"- Rotating {fc_name} (out)\")\n rotation = rotation.to(fc.weight.device)\n rotate_out_channels(fc.weight, rotation=rotation, bias=fc.bias)\n if ffn.down_proj_key in config.transforms:\n down_proj.extend(ffn.down_projs)\n tools.logging.Formatter.indent_dec()\n gc.collect()\n torch.cuda.empty_cache()\n if backbone.proj_out is not None:\n logger.debug(f\"- Rotating {backbone.proj_out_name} (in)\")\n weight = backbone.proj_out.weight\n rotation = rotation.to(weight.device)\n rotate_in_channels(weight, rotation=rotation)\n logger.debug(f\"- Rotating {backbone.proj_out_name} (out)\")\n rotation = rotation.to(weight.device)\n rotate_out_channels(weight, rotation=rotation, bias=backbone.proj_out.bias)\n # endregion\n if down_proj:\n logger.debug(f\"- Applying Hadamard transform on {backbone.name}.down_proj (in)\")\n hadamard_in_channels(down_proj)\n if backbone.proj_out is not None:\n logger.debug(f\"- Rotating {backbone.proj_out_name} (in)\")\n weight = backbone.proj_out.weight\n else:\n logger.debug(f\"- Rotating {model.head_name} (in)\")\n weight = model.head.weight\n rotation = rotation.to(weight.device)\n rotate_in_channels(weight, rotation=rotation)\n for device, dtype, linear in zip(devices, dtypes, linears, strict=True):\n linear.to(device=device, dtype=dtype)\n return rotation.cpu()\n"
},
{
"path": "deepcompressor/app/llm/quant/smooth.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"LLM smooth quantization module.\"\"\"\n\nimport typing as tp\n\nimport torch\nimport torch.nn as nn\nfrom tqdm import tqdm\nfrom transformers import PreTrainedTokenizer\n\nfrom deepcompressor.calib.smooth import ActivationSmoother, smooth_attention, smooth_linear_modules\nfrom deepcompressor.data.cache import IOTensorsCache\nfrom deepcompressor.quantizer.processor import Quantizer\nfrom deepcompressor.utils import tools\n\nfrom ..nn.struct import LlmModelStruct, LlmTransformerBlockStruct\nfrom .config import LlmQuantConfig\nfrom .utils import get_needs_inputs_fn, get_needs_outputs_fn\n\n__all__ = [\"smooth_llm\"]\n\n\n@torch.inference_mode()\ndef smooth_llm_layer( # noqa: C901\n layer: LlmTransformerBlockStruct,\n config: LlmQuantConfig,\n smooth_cache: dict[str, torch.Tensor],\n layer_cache: dict[str, IOTensorsCache] | None = None,\n layer_kwargs: dict[str, tp.Any] | None = None,\n) -> None:\n \"\"\"Smooth a large language model layer.\n\n Args:\n layer (`LlmTransformerBlockStruct`):\n Large language model layer to smooth.\n config (`LlmQuantConfig`):\n Quantization configuration.\n smooth_cache (`dict[str, torch.Tensor]`):\n Smoothing scale caches.\n layer_caches (`dict[str, IOTensorsCache]` or `None`, *optional*, defaults to `None`):\n Activation caches of the layer.\n layer_kwargs (`dict[str, tp.Any]` or `None`, *optional*, defaults to `None`):\n Keyword arguments for the layer.\n \"\"\"\n logger = tools.logging.getLogger(f\"{__name__}.SmoothQuant\")\n logger.debug(\"- Smoothing %s\", layer.name)\n tools.logging.Formatter.indent_inc()\n layer_cache = layer_cache or {}\n layer_kwargs = layer_kwargs or {}\n attn, ffn = layer.attn_struct, layer.ffn_struct\n # region attention qk\n needs_quant = config.enabled_opts\n needs_quant = needs_quant and (config.opts.is_enabled_for(attn.q_key) or config.opts.is_enabled_for(attn.k_key))\n if config.smooth.enabled_attn and needs_quant:\n logger.debug(\"- %s.%s\", attn.name, attn.k_rkey)\n cache_key = f\"{attn.name}.{attn.k_rkey}\"\n smooth_cache[cache_key] = smooth_attention(\n k_proj=attn.k_proj,\n q_proj=attn.q_proj,\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.attn,\n query_quantizer=Quantizer(config.opts, channels_dim=-1, key=attn.q_key),\n key_quantizer=Quantizer(config.opts, channels_dim=-1, key=attn.k_key),\n queries=layer_cache[attn.q_name].outputs if layer_cache else None,\n keys=layer_cache[attn.k_name].outputs if layer_cache else None,\n attn_q=attn.q,\n attn_k=attn.k,\n eval_inputs=layer_cache[attn.name].inputs if layer_cache else None,\n eval_module=attn,\n eval_kwargs=attn.filter_kwargs(layer_kwargs),\n num_heads=attn.config.num_query_heads,\n num_head_repeats=attn.config.num_head_repeats,\n with_rope=attn.config.with_rope,\n develop_dtype=config.develop_dtype,\n )\n # endregion\n # region qkv projection\n needs_quant = config.enabled_ipts and config.ipts.is_enabled_for(attn.qkv_proj_key)\n needs_quant = needs_quant or (config.enabled_wgts and config.wgts.is_enabled_for(attn.qkv_proj_key))\n if config.smooth.enabled_proj and config.smooth.proj.is_enabled_for(attn.qkv_proj_key) and needs_quant:\n logger.debug(\"- %s.%s\", attn.name, attn.qkv_proj_rkey)\n cache_key = attn.v_proj_name\n smooth_cache[cache_key] = smooth_linear_modules(\n attn.parent.pre_attn_norm,\n attn.qkv_proj,\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config.wgts, key=attn.qkv_proj_key),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=attn.qkv_proj_key),\n inputs=layer_cache[attn.q_proj_name].inputs if layer_cache else None,\n eval_inputs=layer_cache[attn.name].inputs if layer_cache else None,\n eval_module=attn,\n eval_kwargs=attn.filter_kwargs(layer_kwargs),\n develop_dtype=config.develop_dtype,\n )\n if not attn.parent.pre_attn_norm:\n ActivationSmoother(smooth_cache[cache_key], channels_dim=-1).as_hook().register(attn.qkv_proj)\n # endregion\n # region output projection\n needs_quant = config.enabled_ipts and config.ipts.is_enabled_for(attn.out_proj_key)\n needs_quant = needs_quant or (config.enabled_wgts and config.wgts.is_enabled_for(attn.out_proj_key))\n if config.smooth.enabled_proj and config.smooth.proj.is_enabled_for(attn.out_proj_key) and needs_quant:\n logger.debug(\"- %s.%s\", attn.name, attn.out_proj_rkey)\n cache_key = attn.o_proj_name\n smooth_cache[cache_key] = smooth_linear_modules(\n None if attn.config.linear_attn else attn.v_proj,\n attn.o_proj,\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config.wgts, key=attn.out_proj_key),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=attn.out_proj_key),\n inputs=layer_cache[cache_key].inputs if layer_cache else None,\n eval_inputs=layer_cache[cache_key].inputs if layer_cache else None,\n eval_module=attn.o_proj,\n num_heads=attn.config.num_query_heads,\n num_head_repeats=attn.config.num_head_repeats,\n develop_dtype=config.develop_dtype,\n )\n if attn.config.linear_attn:\n ActivationSmoother(smooth_cache[cache_key], channels_dim=-1).as_hook().register(attn.o_proj)\n # endregion\n num_experts = ffn.config.num_experts\n # region up projection\n needs_quant = config.enabled_ipts and config.ipts.is_enabled_for(ffn.up_proj_key)\n needs_quant = needs_quant or (config.enabled_wgts and config.wgts.is_enabled_for(ffn.up_proj_key))\n if config.smooth.enabled_proj and config.smooth.proj.is_enabled_for(ffn.up_proj_key) and needs_quant:\n logger.debug(\"- %s.%s\", ffn.name, ffn.up_proj_rkey)\n cache_key = ffn.name\n smooth_cache[cache_key] = smooth_linear_modules(\n ffn.parent.pre_ffn_norm,\n ffn.up_projs,\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config.wgts, key=ffn.up_proj_key),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=ffn.up_proj_key),\n inputs=layer_cache[ffn.name].inputs if layer_cache else None,\n eval_inputs=layer_cache[ffn.name].inputs if layer_cache else None,\n eval_module=ffn,\n extra_modules=[ffn.moe_gate] if num_experts > 1 else None,\n develop_dtype=config.develop_dtype,\n )\n if not ffn.parent.pre_ffn_norm:\n hook = ActivationSmoother(smooth_cache[cache_key], channels_dim=-1).as_hook().register(ffn.up_projs)\n if num_experts > 1:\n hook.register(ffn.moe_gate)\n # endregion\n # region down projection\n needs_quant = config.enabled_ipts and config.ipts.is_enabled_for(ffn.down_proj_key)\n needs_quant = needs_quant or (config.enabled_wgts and config.wgts.is_enabled_for(ffn.down_proj_key))\n if config.smooth.enabled_proj and config.smooth.proj.is_enabled_for(ffn.down_proj_key) and needs_quant:\n for expert_idx in range(num_experts):\n logger.debug(\"- %s.%s\", ffn.expert_names[expert_idx], ffn.down_proj_rkey)\n cache_key = ffn.down_proj_names[expert_idx]\n smooth_cache[cache_key] = smooth_linear_modules(\n ffn.up_projs[expert_idx],\n ffn.down_projs[expert_idx],\n scale=smooth_cache.get(cache_key, None),\n config=config.smooth.proj,\n weight_quantizer=Quantizer(config.wgts, key=ffn.down_proj_key),\n input_quantizer=Quantizer(config.ipts, channels_dim=-1, key=ffn.down_proj_key),\n inputs=layer_cache[ffn.down_proj_names[expert_idx]].inputs if layer_cache else None,\n eval_inputs=layer_cache[ffn.down_proj_names[expert_idx]].inputs if layer_cache else None,\n eval_module=ffn.down_projs[expert_idx],\n develop_dtype=config.develop_dtype,\n )\n # endregion\n tools.logging.Formatter.indent_dec()\n\n\n@torch.inference_mode()\ndef smooth_llm(\n model: nn.Module | LlmModelStruct,\n /,\n config: LlmQuantConfig,\n tokenizer: PreTrainedTokenizer | None = None,\n smooth_cache: dict[str, torch.Tensor] | None = None,\n) -> dict[str, torch.Tensor]:\n \"\"\"Smooth the large language model.\n\n Args:\n model (`nn.Module` or `LlmStruct`):\n Model to be smoothed.\n config (`LlmQuantConfig`):\n Quantization configuration.\n tokenizer (`PreTrainedTokenizer`, *optional*, defaults to `None`):\n Tokenizer.\n smooth_cache (`dict[str, torch.Tensor]`, *optional*, defaults to `None`):\n Smoothing scale caches.\n\n Returns:\n `dict[str, torch.Tensor]`:\n Dictionary mapping module names to smoothing scales.\n \"\"\"\n if not isinstance(model, LlmModelStruct):\n model = LlmModelStruct.construct(model)\n assert isinstance(model, LlmModelStruct)\n smooth_cache = smooth_cache or {}\n if not smooth_cache:\n with tools.logging.redirect_tqdm():\n for _, (layer, layer_cache, layer_kwargs) in tqdm(\n config.calib.build_loader(tokenizer).iter_layer_activations(\n model,\n needs_inputs_fn=get_needs_inputs_fn(model=model, config=config),\n needs_outputs_fn=get_needs_outputs_fn(model=model, config=config),\n ),\n desc=\"smoothing\",\n leave=False,\n total=len(model.backbone_struct.layer_structs),\n dynamic_ncols=True,\n ):\n smooth_llm_layer(\n layer=layer,\n config=config,\n smooth_cache=smooth_cache,\n layer_cache=layer_cache,\n layer_kwargs=layer_kwargs,\n )\n else:\n for layer in model.backbone_struct.layer_structs:\n smooth_llm_layer(layer=layer, config=config, smooth_cache=smooth_cache)\n return smooth_cache\n"
},
{
"path": "deepcompressor/app/llm/quant/utils.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"LLM quantization utils module.\"\"\"\n\nimport typing as tp\n\nimport torch.nn as nn\n\nfrom ..nn.struct import LlmModelStruct\nfrom .quantizer.config import LlmModuleQuantizerConfig\n\n__all__ = [\"get_needs_inputs_fn\", \"get_needs_outputs_fn\"]\n\n\ndef get_needs_inputs_fn(model: LlmModelStruct, config: LlmModuleQuantizerConfig) -> tp.Callable[[str, nn.Module], bool]:\n \"\"\"Get function that checks if the module needs to cache the inputs.\n\n Args:\n model (`LlmStruct`):\n Model struct.\n config (`LlmModuleQuantizerConfig`):\n Module quantization config.\n\n Returns:\n `Callable[[str, nn.Module], bool]`:\n Function to check if the module needs to cache the inputs.\n \"\"\"\n\n needs_inputs_names = set()\n\n example_layer = model.backbone_struct.layer_structs[0]\n attn, ffn = example_layer.attn_struct, example_layer.ffn_struct\n if (config.enabled_wgts and config.wgts.is_enabled_for(attn.qkv_proj_key)) or (\n config.enabled_ipts and config.ipts.is_enabled_for(attn.qkv_proj_key)\n ):\n needs_inputs_names.add(attn.rname)\n needs_inputs_names.add(attn.v_proj_rname)\n if (config.enabled_wgts and config.wgts.is_enabled_for(attn.out_proj_key)) or (\n config.enabled_ipts and config.ipts.is_enabled_for(attn.out_proj_key)\n ):\n needs_inputs_names.add(attn.o_proj_rname)\n if (config.enabled_wgts and config.wgts.is_enabled_for(ffn.up_proj_key)) or (\n config.enabled_ipts and config.ipts.is_enabled_for(ffn.up_proj_key)\n ):\n needs_inputs_names.add(ffn.rname)\n needs_inputs_names.add(ffn.up_proj_rnames[0])\n if (config.enabled_wgts and config.wgts.is_enabled_for(ffn.down_proj_key)) or (\n config.enabled_ipts and config.ipts.is_enabled_for(ffn.down_proj_key)\n ):\n needs_inputs_names.add(ffn.down_proj_rnames[0])\n if config.enabled_opts:\n needs_inputs_names.add(attn.rname)\n\n needs_inputs_names = tuple(needs_inputs_names)\n\n def needs_inputs(name: str, module: nn.Module) -> bool:\n return name.endswith(needs_inputs_names)\n\n return needs_inputs\n\n\ndef get_needs_outputs_fn(\n model: LlmModelStruct, config: LlmModuleQuantizerConfig\n) -> tp.Callable[[str, nn.Module], bool]:\n \"\"\"Get function that checks if the module needs to cache the outputs.\n\n Args:\n model (`LlmStruct`):\n Model struct.\n config (`LlmModuleQuantizerConfig`):\n Module quantization config.\n\n Returns:\n `Callable[[str, nn.Module], bool]`:\n Function to check if the module needs to cache the outputs.\n \"\"\"\n\n attn = model.backbone_struct.layer_structs[0].attn_struct\n needs_outputs_names = set()\n if config.enabled_opts:\n needs_outputs_names.add(attn.q_rname)\n needs_outputs_names.add(attn.k_rname)\n needs_outputs_names.add(attn.v_rname)\n needs_outputs_names = tuple(needs_outputs_names)\n\n def needs_outputs(name: str, module: nn.Module) -> bool:\n return name.endswith(needs_outputs_names)\n\n return needs_outputs\n"
},
{
"path": "deepcompressor/app/llm/quant/weight.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"LLM weight quantization calibration module.\"\"\"\n\nimport gc\nimport typing as tp\n\nimport torch\nimport torch.nn as nn\nfrom tqdm import tqdm\nfrom transformers import PreTrainedTokenizer\n\nfrom deepcompressor.data.cache import IOTensorsCache\nfrom deepcompressor.data.zero import ZeroPointDomain\nfrom deepcompressor.utils import tools\n\nfrom ..nn.struct import LlmModelStruct, LlmSelfAttentionStruct, LlmTransformerBlockStruct\nfrom .config import LlmQuantConfig\nfrom .quantizer import LlmWeightQuantizer\nfrom .utils import get_needs_inputs_fn\n\n__all__ = [\"quantize_llm_weights\"]\n\n\n@torch.inference_mode()\ndef quantize_llm_layer_weights( # noqa: C901\n layer: LlmTransformerBlockStruct,\n config: LlmQuantConfig,\n quantizer_state_dict: dict[str, tp.Any],\n layer_cache: dict[str, IOTensorsCache] | None = None,\n layer_kwargs: dict[str, tp.Any] | None = None,\n return_with_scale_state_dict: bool = False,\n) -> dict[str, torch.Tensor | float | None]:\n \"\"\"Calibrate the weight quantization ranges of modules in a layer.\n\n Args:\n layer (`LlmTransformerBlockStruct`):\n Layer.\n config (`LlmQuantConfig`):\n Quantization config.\n quantizer_state_dict (`dict[str, Any]`):\n Weight quantizer.\n layer_cache (`dict[str, IOTensorsCache]` or `None`, *optional*, defaults to `None`):\n Layer activations cache.\n layer_kwargs (`dict[str, tp.Any]` or `None`, *optional*, defaults to `None`):\n Keyword arguments for the layer.\n return_with_scale_state_dict (bool, *optional*, defaults to `False`):\n Whether to return with scale state dict.\n\n Returns:\n `dict[str, torch.Tensor | float | None]`:\n Scale state dict.\n \"\"\"\n logger = tools.logging.getLogger(f\"{__name__}.WeightQuant\")\n logger.debug(\"- Quantizing layer %s\", layer.name)\n tools.logging.Formatter.indent_inc()\n layer_cache = layer_cache or {}\n layer_kwargs = layer_kwargs or {}\n for module_key, module_name, module, parent, field_name in layer.named_key_modules():\n assert isinstance(module, nn.Linear)\n if field_name in (\"q_proj\", \"k_proj\"):\n assert isinstance(parent, LlmSelfAttentionStruct)\n eval_name, eval_module, eval_kwargs = parent.name, parent.module, parent.filter_kwargs(layer_kwargs)\n else:\n eval_name, eval_module, eval_kwargs = module_name, module, None\n quantizer = LlmWeightQuantizer(config.wgts, develop_dtype=config.develop_dtype, key=module_key)\n if quantizer.is_enabled():\n if module_name not in quantizer_state_dict:\n logger.debug(\"- Calibrating %s.weight\", module_name)\n quantizer.calibrate_dynamic_range(\n module=module,\n inputs=layer_cache[module_name].inputs if layer_cache else None,\n eval_inputs=layer_cache[eval_name].inputs if layer_cache else None,\n eval_module=eval_module,\n eval_kwargs=eval_kwargs,\n )\n quantizer_state_dict[module_name] = quantizer.state_dict()\n gc.collect()\n torch.cuda.empty_cache()\n scale_state_dict: dict[str, torch.Tensor | float | None] = {}\n for module_key, module_name, module, _, _ in layer.named_key_modules():\n assert isinstance(module, nn.Linear)\n quantizer = LlmWeightQuantizer(config.wgts, develop_dtype=config.develop_dtype, key=module_key)\n param_name = f\"{module_name}.weight\"\n if quantizer.is_enabled():\n logger.debug(\"- Quantizing %s\", param_name)\n quantizer.load_state_dict(quantizer_state_dict[module_name], device=module.weight.device)\n result = quantizer.quantize(\n module.weight.data,\n inputs=layer_cache[module_name].inputs.front() if layer_cache else None,\n return_with_dequant=True,\n return_with_quant=return_with_scale_state_dict,\n )\n module.weight.data = result.data\n if return_with_scale_state_dict:\n scale_state_dict.update(result.scale.state_dict(f\"{param_name}.scale\"))\n zero_name = \"scaled_zero\" if config.wgts.zero_domain is ZeroPointDomain.PostScale else \"zero\"\n if isinstance(result.zero, torch.Tensor):\n scale_state_dict[f\"{param_name}.{zero_name}\"] = result.zero.to(\"cpu\")\n else:\n scale_state_dict[f\"{param_name}.{zero_name}\"] = result.zero\n del result\n gc.collect()\n torch.cuda.empty_cache()\n tools.logging.Formatter.indent_dec()\n return scale_state_dict\n\n\n@torch.inference_mode()\ndef quantize_llm_weights(\n model: nn.Module | LlmModelStruct,\n config: LlmQuantConfig,\n tokenizer: PreTrainedTokenizer | None = None,\n quantizer_state_dict: dict[str, tp.Any] | None = None,\n return_with_scale_state_dict: bool = False,\n) -> tuple[dict[str, tp.Any], dict[str, torch.Tensor | float | None]]:\n \"\"\"Quantize the large language model weights.\n\n Args:\n model (`nn.Module` or `LlmStruct`):\n Model to be quantized.\n config (`LlmQuantConfig`):\n Quantization configuration.\n tokenizer (`PreTrainedTokenizer`, *optional*, defaults to `None`):\n Tokenizer.\n quantizer_state_dict (`dict[str, Any]`, *optional*, defaults to `None`):\n Weight quantizer state dict.\n return_with_scale_state_dict (bool, *optional*, defaults to `False`):\n Whether to return with scale state dict.\n\n Returns:\n `tuple[dict[str, Any], dict[str, torch.Tensor | float | None]`:\n Weight quantizer cache and scale state dict.\n \"\"\"\n if not isinstance(model, LlmModelStruct):\n model = LlmModelStruct.construct(model)\n assert isinstance(model, LlmModelStruct)\n quantizer_state_dict = quantizer_state_dict or {}\n scale_state_dict: dict[str, torch.Tensor | float | None] = {}\n with tools.logging.redirect_tqdm():\n if config.wgts.enabled_gptq or (not quantizer_state_dict and config.wgts.needs_calib_data):\n for _, (layer, layer_cache, layer_kwargs) in tqdm(\n config.calib.build_loader(tokenizer).iter_layer_activations(\n model,\n needs_inputs_fn=get_needs_inputs_fn(model=model, config=config),\n ),\n desc=\"quantizing weights\",\n leave=False,\n total=len(model.backbone_struct.layer_structs),\n dynamic_ncols=True,\n ):\n scale_state_dict.update(\n quantize_llm_layer_weights(\n layer=layer,\n config=config,\n quantizer_state_dict=quantizer_state_dict,\n layer_cache=layer_cache,\n layer_kwargs=layer_kwargs,\n return_with_scale_state_dict=return_with_scale_state_dict,\n )\n )\n else:\n for layer in tqdm(\n model.backbone_struct.layer_structs, desc=\"quantizing weights\", leave=False, dynamic_ncols=True\n ):\n scale_state_dict.update(\n quantize_llm_layer_weights(\n layer=layer,\n config=config,\n quantizer_state_dict=quantizer_state_dict,\n return_with_scale_state_dict=return_with_scale_state_dict,\n )\n )\n return quantizer_state_dict, scale_state_dict\n"
},
{
"path": "deepcompressor/backend/__init__.py",
"content": ""
},
{
"path": "deepcompressor/backend/nunchaku/__init__.py",
"content": ""
},
{
"path": "deepcompressor/backend/nunchaku/convert.py",
"content": "\"\"\"Converts a DeepCompressor state dict to a Nunchaku state dict.\"\"\"\n\nimport argparse\nimport os\n\nimport safetensors.torch\nimport torch\nimport tqdm\n\nfrom .utils import convert_to_nunchaku_w4x4y16_linear_weight, convert_to_nunchaku_w4x16_linear_weight\n\n\ndef convert_to_nunchaku_w4x4y16_linear_state_dict(\n weight: torch.Tensor,\n scale: torch.Tensor,\n bias: torch.Tensor | None = None,\n smooth: torch.Tensor | None = None,\n lora: tuple[torch.Tensor, torch.Tensor] | None = None,\n shift: torch.Tensor | None = None,\n smooth_fused: bool = False,\n float_point: bool = False,\n subscale: torch.Tensor | None = None,\n) -> dict[str, torch.Tensor]:\n if weight.ndim > 2: # pointwise conv\n assert weight.numel() == weight.shape[0] * weight.shape[1]\n weight = weight.view(weight.shape[0], weight.shape[1])\n if scale.numel() > 1:\n assert scale.ndim == weight.ndim * 2\n assert scale.numel() == scale.shape[0] * scale.shape[2]\n scale = scale.view(scale.shape[0], 1, scale.shape[2], 1)\n scale_key = \"wcscales\" if scale.shape[2] == 1 else \"wscales\"\n else:\n scale_key = \"wtscale\"\n if subscale is None:\n subscale_key = \"\"\n else:\n assert subscale.ndim == weight.ndim * 2\n assert subscale.numel() == subscale.shape[0] * subscale.shape[2]\n assert subscale.numel() > 1\n subscale = subscale.view(subscale.shape[0], 1, subscale.shape[2], 1)\n subscale_key = \"wcscales\" if subscale.shape[2] == 1 else \"wscales\"\n if lora is not None and (smooth is not None or shift is not None):\n # unsmooth lora down projection\n dtype = weight.dtype\n lora_down, lora_up = lora\n lora_down = lora_down.to(dtype=torch.float64)\n if smooth is not None and not smooth_fused:\n lora_down = lora_down.div_(smooth.to(torch.float64).unsqueeze(0))\n if shift is not None:\n bias = torch.zeros([lora_up.shape[0]], dtype=torch.float64) if bias is None else bias.to(torch.float64)\n if shift.numel() == 1:\n shift = shift.view(1, 1).expand(lora_down.shape[1], 1).to(torch.float64)\n else:\n shift = shift.view(-1, 1).to(torch.float64)\n bias = bias.add_((lora_up.to(dtype=torch.float64) @ lora_down @ shift).view(-1))\n bias = bias.to(dtype=dtype)\n lora = (lora_down.to(dtype=dtype), lora_up)\n weight, scale, bias, smooth, lora, subscale = convert_to_nunchaku_w4x4y16_linear_weight(\n weight, scale=scale, bias=bias, smooth=smooth, lora=lora, float_point=float_point, subscale=subscale\n )\n state_dict: dict[str, torch.Tensor] = {}\n state_dict[\"qweight\"] = weight\n state_dict[scale_key] = scale\n if subscale is not None:\n state_dict[subscale_key] = subscale\n state_dict[\"bias\"] = bias\n state_dict[\"smooth_orig\"] = smooth\n state_dict[\"smooth\"] = torch.ones_like(smooth) if smooth_fused else smooth.clone()\n if lora is not None:\n state_dict[\"lora_down\"] = lora[0]\n state_dict[\"lora_up\"] = lora[1]\n return state_dict\n\n\ndef convert_to_nunchaku_w4x16_adanorm_single_state_dict(\n weight: torch.Tensor,\n scale: torch.Tensor,\n bias: torch.Tensor,\n) -> dict[str, torch.Tensor]:\n weight, scale, zero, bias = convert_to_nunchaku_w4x16_linear_weight(\n weight, scale=scale, bias=bias, adanorm_splits=3\n )\n state_dict: dict[str, torch.Tensor] = {}\n state_dict = {}\n state_dict[\"qweight\"] = weight\n state_dict[\"wscales\"] = scale\n state_dict[\"wzeros\"] = zero\n state_dict[\"bias\"] = bias\n return state_dict\n\n\ndef convert_to_nunchaku_w4x16_adanorm_zero_state_dict(\n weight: torch.Tensor,\n scale: torch.Tensor,\n bias: torch.Tensor,\n) -> dict[str, torch.Tensor]:\n weight, scale, zero, bias = convert_to_nunchaku_w4x16_linear_weight(\n weight, scale=scale, bias=bias, adanorm_splits=6\n )\n state_dict: dict[str, torch.Tensor] = {}\n state_dict = {}\n state_dict[\"qweight\"] = weight\n state_dict[\"wscales\"] = scale\n state_dict[\"wzeros\"] = zero\n state_dict[\"bias\"] = bias\n return state_dict\n\n\ndef update_state_dict(\n lhs: dict[str, torch.Tensor], rhs: dict[str, torch.Tensor], prefix: str = \"\"\n) -> dict[str, torch.Tensor]:\n for rkey, value in rhs.items():\n lkey = f\"{prefix}.{rkey}\" if prefix else rkey\n assert lkey not in lhs, f\"Key {lkey} already exists in the state dict.\"\n lhs[lkey] = value\n return lhs\n\n\ndef convert_to_nunchaku_transformer_block_state_dict(\n state_dict: dict[str, torch.Tensor],\n scale_dict: dict[str, torch.Tensor],\n smooth_dict: dict[str, torch.Tensor],\n branch_dict: dict[str, torch.Tensor],\n block_name: str,\n local_name_map: dict[str, str | list[str]],\n smooth_name_map: dict[str, str],\n branch_name_map: dict[str, str],\n convert_map: dict[str, str],\n float_point: bool = False,\n) -> dict[str, torch.Tensor]:\n print(f\"Converting block {block_name}...\")\n converted: dict[str, torch.Tensor] = {}\n candidates: dict[str, torch.Tensor] = {\n param_name: param for param_name, param in state_dict.items() if param_name.startswith(block_name)\n }\n for converted_local_name, candidate_local_names in tqdm.tqdm(\n local_name_map.items(), desc=f\"Converting {block_name}\", dynamic_ncols=True\n ):\n if isinstance(candidate_local_names, str):\n candidate_local_names = [candidate_local_names]\n candidate_names = [f\"{block_name}.{candidate_local_name}\" for candidate_local_name in candidate_local_names]\n weight = [candidates[f\"{candidate_name}.weight\"] for candidate_name in candidate_names]\n bias = [candidates.get(f\"{candidate_name}.bias\", None) for candidate_name in candidate_names]\n scale = [scale_dict.get(f\"{candidate_name}.weight.scale.0\", None) for candidate_name in candidate_names]\n subscale = [scale_dict.get(f\"{candidate_name}.weight.scale.1\", None) for candidate_name in candidate_names]\n if len(weight) > 1:\n bias = None if all(b is None for b in bias) else torch.concat(bias, dim=0)\n if all(s is None for s in scale):\n scale = None\n else:\n if scale[0].numel() == 1: # switch from per-tensor to per-channel scale\n assert all(s.numel() == 1 for s in scale)\n scale = torch.concat(\n [\n s.view(-1).expand(weight[i].shape[0]).reshape(weight[i].shape[0], 1, 1, 1)\n for i, s in enumerate(scale)\n ],\n dim=0,\n )\n else:\n scale = torch.concat(scale, dim=0)\n subscale = None if all(s is None for s in subscale) else torch.concat(subscale, dim=0)\n weight = torch.concat(weight, dim=0)\n else:\n weight, bias, scale, subscale = weight[0], bias[0], scale[0], subscale[0]\n smooth = smooth_dict.get(f\"{block_name}.{smooth_name_map.get(converted_local_name, '')}\", None)\n branch = branch_dict.get(f\"{block_name}.{branch_name_map.get(converted_local_name, '')}\", None)\n if branch is not None:\n branch = (branch[\"a.weight\"], branch[\"b.weight\"])\n if scale is None:\n assert smooth is None and branch is None and subscale is None\n print(f\" - Copying {block_name} weights of {candidate_local_names} as {converted_local_name}.weight\")\n converted[f\"{converted_local_name}.weight\"] = weight.clone().cpu()\n if bias is not None:\n print(f\" - Copying {block_name} biases of {candidate_local_names} as {converted_local_name}.bias\")\n converted[f\"{converted_local_name}.bias\"] = bias.clone().cpu()\n continue\n if convert_map[converted_local_name] == \"adanorm_single\":\n print(f\" - Converting {block_name} weights of {candidate_local_names} to {converted_local_name}.\")\n update_state_dict(\n converted,\n convert_to_nunchaku_w4x16_adanorm_single_state_dict(weight=weight, scale=scale, bias=bias),\n prefix=converted_local_name,\n )\n elif convert_map[converted_local_name] == \"adanorm_zero\":\n print(f\" - Converting {block_name} weights of {candidate_local_names} to {converted_local_name}.\")\n update_state_dict(\n converted,\n convert_to_nunchaku_w4x16_adanorm_zero_state_dict(weight=weight, scale=scale, bias=bias),\n prefix=converted_local_name,\n )\n elif convert_map[converted_local_name] == \"linear\":\n smooth_fused = \"out_proj\" in converted_local_name and smooth_dict.get(\"proj.fuse_when_possible\", True)\n shift = [candidates.get(f\"{candidate_name[:-7]}.shift\", None) for candidate_name in candidate_names]\n assert all(s == shift[0] for s in shift)\n shift = shift[0]\n print(\n f\" - Converting {block_name} weights of {candidate_local_names} to {converted_local_name}.\"\n f\" (smooth_fused={smooth_fused}, shifted={shift is not None}, float_point={float_point})\"\n )\n update_state_dict(\n converted,\n convert_to_nunchaku_w4x4y16_linear_state_dict(\n weight=weight,\n scale=scale,\n bias=bias,\n smooth=smooth,\n lora=branch,\n shift=shift,\n smooth_fused=smooth_fused,\n float_point=float_point,\n subscale=subscale,\n ),\n prefix=converted_local_name,\n )\n else:\n raise NotImplementedError(f\"Conversion of {convert_map[converted_local_name]} is not implemented.\")\n return converted\n\n\ndef convert_to_nunchaku_flux_single_transformer_block_state_dict(\n state_dict: dict[str, torch.Tensor],\n scale_dict: dict[str, torch.Tensor],\n smooth_dict: dict[str, torch.Tensor],\n branch_dict: dict[str, torch.Tensor],\n block_name: str,\n float_point: bool = False,\n) -> dict[str, torch.Tensor]:\n down_proj_local_name = \"proj_out.linears.1.linear\"\n if f\"{block_name}.{down_proj_local_name}.weight\" not in state_dict:\n down_proj_local_name = \"proj_out.linears.1\"\n assert f\"{block_name}.{down_proj_local_name}.weight\" in state_dict\n\n return convert_to_nunchaku_transformer_block_state_dict(\n state_dict=state_dict,\n scale_dict=scale_dict,\n smooth_dict=smooth_dict,\n branch_dict=branch_dict,\n block_name=block_name,\n local_name_map={\n \"norm.linear\": \"norm.linear\",\n \"qkv_proj\": [\"attn.to_q\", \"attn.to_k\", \"attn.to_v\"],\n \"norm_q\": \"attn.norm_q\",\n \"norm_k\": \"attn.norm_k\",\n \"out_proj\": \"proj_out.linears.0\",\n \"mlp_fc1\": \"proj_mlp\",\n \"mlp_fc2\": down_proj_local_name,\n },\n smooth_name_map={\n \"qkv_proj\": \"attn.to_q\",\n \"out_proj\": \"proj_out.linears.0\",\n \"mlp_fc1\": \"attn.to_q\",\n \"mlp_fc2\": down_proj_local_name,\n },\n branch_name_map={\n \"qkv_proj\": \"attn.to_q\",\n \"out_proj\": \"proj_out.linears.0\",\n \"mlp_fc1\": \"proj_mlp\",\n \"mlp_fc2\": down_proj_local_name,\n },\n convert_map={\n \"norm.linear\": \"adanorm_single\",\n \"qkv_proj\": \"linear\",\n \"out_proj\": \"linear\",\n \"mlp_fc1\": \"linear\",\n \"mlp_fc2\": \"linear\",\n },\n float_point=float_point,\n )\n\n\ndef convert_to_nunchaku_flux_transformer_block_state_dict(\n state_dict: dict[str, torch.Tensor],\n scale_dict: dict[str, torch.Tensor],\n smooth_dict: dict[str, torch.Tensor],\n branch_dict: dict[str, torch.Tensor],\n block_name: str,\n float_point: bool = False,\n) -> dict[str, torch.Tensor]:\n down_proj_local_name = \"ff.net.2.linear\"\n if f\"{block_name}.{down_proj_local_name}.weight\" not in state_dict:\n down_proj_local_name = \"ff.net.2\"\n assert f\"{block_name}.{down_proj_local_name}.weight\" in state_dict\n context_down_proj_local_name = \"ff_context.net.2.linear\"\n if f\"{block_name}.{context_down_proj_local_name}.weight\" not in state_dict:\n context_down_proj_local_name = \"ff_context.net.2\"\n assert f\"{block_name}.{context_down_proj_local_name}.weight\" in state_dict\n\n return convert_to_nunchaku_transformer_block_state_dict(\n state_dict=state_dict,\n scale_dict=scale_dict,\n smooth_dict=smooth_dict,\n branch_dict=branch_dict,\n block_name=block_name,\n local_name_map={\n \"norm1.linear\": \"norm1.linear\",\n \"norm1_context.linear\": \"norm1_context.linear\",\n \"qkv_proj\": [\"attn.to_q\", \"attn.to_k\", \"attn.to_v\"],\n \"qkv_proj_context\": [\"attn.add_q_proj\", \"attn.add_k_proj\", \"attn.add_v_proj\"],\n \"norm_q\": \"attn.norm_q\",\n \"norm_k\": \"attn.norm_k\",\n \"norm_added_q\": \"attn.norm_added_q\",\n \"norm_added_k\": \"attn.norm_added_k\",\n \"out_proj\": \"attn.to_out.0\",\n \"out_proj_context\": \"attn.to_add_out\",\n \"mlp_fc1\": \"ff.net.0.proj\",\n \"mlp_fc2\": down_proj_local_name,\n \"mlp_context_fc1\": \"ff_context.net.0.proj\",\n \"mlp_context_fc2\": context_down_proj_local_name,\n },\n smooth_name_map={\n \"qkv_proj\": \"attn.to_q\",\n \"qkv_proj_context\": \"attn.add_k_proj\",\n \"out_proj\": \"attn.to_out.0\",\n \"out_proj_context\": \"attn.to_out.0\",\n \"mlp_fc1\": \"ff.net.0.proj\",\n \"mlp_fc2\": down_proj_local_name,\n \"mlp_context_fc1\": \"ff_context.net.0.proj\",\n \"mlp_context_fc2\": context_down_proj_local_name,\n },\n branch_name_map={\n \"qkv_proj\": \"attn.to_q\",\n \"qkv_proj_context\": \"attn.add_k_proj\",\n \"out_proj\": \"attn.to_out.0\",\n \"out_proj_context\": \"attn.to_add_out\",\n \"mlp_fc1\": \"ff.net.0.proj\",\n \"mlp_fc2\": down_proj_local_name,\n \"mlp_context_fc1\": \"ff_context.net.0.proj\",\n \"mlp_context_fc2\": context_down_proj_local_name,\n },\n convert_map={\n \"norm1.linear\": \"adanorm_zero\",\n \"norm1_context.linear\": \"adanorm_zero\",\n \"qkv_proj\": \"linear\",\n \"qkv_proj_context\": \"linear\",\n \"out_proj\": \"linear\",\n \"out_proj_context\": \"linear\",\n \"mlp_fc1\": \"linear\",\n \"mlp_fc2\": \"linear\",\n \"mlp_context_fc1\": \"linear\",\n \"mlp_context_fc2\": \"linear\",\n },\n float_point=float_point,\n )\n\n\ndef convert_to_nunchaku_flux_state_dicts(\n state_dict: dict[str, torch.Tensor],\n scale_dict: dict[str, torch.Tensor],\n smooth_dict: dict[str, torch.Tensor],\n branch_dict: dict[str, torch.Tensor],\n float_point: bool = False,\n) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:\n block_names: set[str] = set()\n other: dict[str, torch.Tensor] = {}\n for param_name in state_dict.keys():\n if param_name.startswith((\"transformer_blocks.\", \"single_transformer_blocks.\")):\n block_names.add(\".\".join(param_name.split(\".\")[:2]))\n else:\n other[param_name] = state_dict[param_name]\n block_names = sorted(block_names, key=lambda x: (x.split(\".\")[0], int(x.split(\".\")[-1])))\n print(f\"Converting {len(block_names)} transformer blocks...\")\n converted: dict[str, torch.Tensor] = {}\n for block_name in block_names:\n convert_fn = convert_to_nunchaku_flux_single_transformer_block_state_dict\n if block_name.startswith(\"transformer_blocks\"):\n convert_fn = convert_to_nunchaku_flux_transformer_block_state_dict\n update_state_dict(\n converted,\n convert_fn(\n state_dict=state_dict,\n scale_dict=scale_dict,\n smooth_dict=smooth_dict,\n branch_dict=branch_dict,\n block_name=block_name,\n float_point=float_point,\n ),\n prefix=block_name,\n )\n return converted, other\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--quant-path\", type=str, required=True, help=\"path to the quantization checkpoint directory.\")\n parser.add_argument(\"--output-root\", type=str, default=\"\", help=\"root to the output checkpoint directory.\")\n parser.add_argument(\"--model-name\", type=str, default=None, help=\"name of the model.\")\n parser.add_argument(\"--float-point\", action=\"store_true\", help=\"use float-point 4-bit quantization.\")\n args = parser.parse_args()\n if not args.output_root:\n args.output_root = args.quant_path\n if args.model_name is None:\n assert args.model_path is not None, \"model name or path is required.\"\n model_name = args.model_path.rstrip(os.sep).split(os.sep)[-1]\n print(f\"Model name not provided, using {model_name} as the model name.\")\n else:\n model_name = args.model_name\n assert model_name, \"Model name must be provided.\"\n assert \"flux\" in model_name.lower(), \"Only Flux models are supported.\"\n state_dict_path = os.path.join(args.quant_path, \"model.pt\")\n scale_dict_path = os.path.join(args.quant_path, \"scale.pt\")\n smooth_dict_path = os.path.join(args.quant_path, \"smooth.pt\")\n branch_dict_path = os.path.join(args.quant_path, \"branch.pt\")\n map_location = \"cuda\" if torch.cuda.is_available() and torch.cuda.device_count() > 0 else \"cpu\"\n state_dict = torch.load(state_dict_path, map_location=map_location)\n scale_dict = torch.load(scale_dict_path, map_location=\"cpu\")\n smooth_dict = torch.load(smooth_dict_path, map_location=map_location) if os.path.exists(smooth_dict_path) else {}\n branch_dict = torch.load(branch_dict_path, map_location=map_location) if os.path.exists(branch_dict_path) else {}\n converted_state_dict, other_state_dict = convert_to_nunchaku_flux_state_dicts(\n state_dict=state_dict,\n scale_dict=scale_dict,\n smooth_dict=smooth_dict,\n branch_dict=branch_dict,\n float_point=args.float_point,\n )\n output_dirpath = os.path.join(args.output_root, model_name)\n os.makedirs(output_dirpath, exist_ok=True)\n safetensors.torch.save_file(converted_state_dict, os.path.join(output_dirpath, \"transformer_blocks.safetensors\"))\n safetensors.torch.save_file(other_state_dict, os.path.join(output_dirpath, \"unquantized_layers.safetensors\"))\n print(f\"Quantized model saved to {output_dirpath}.\")\n"
},
{
"path": "deepcompressor/backend/nunchaku/convert_lora.py",
"content": "\"\"\"Convert LoRA weights to Nunchaku format.\"\"\"\n\nimport argparse\nimport os\n\nimport safetensors\nimport safetensors.torch\nimport torch\nimport tqdm\n\nfrom ..utils import load_state_dict_in_safetensors, pad\nfrom .convert import update_state_dict\nfrom .utils import NunchakuWeightPacker\n\n\ndef reorder_adanorm_lora_up(lora_up: torch.Tensor, splits: int) -> torch.Tensor:\n c, r = lora_up.shape\n assert c % splits == 0\n return lora_up.view(splits, c // splits, r).transpose(0, 1).reshape(c, r).contiguous()\n\n\ndef convert_to_nunchaku_transformer_block_lowrank_dict( # noqa: C901\n orig_state_dict: dict[str, torch.Tensor],\n extra_lora_dict: dict[str, torch.Tensor],\n converted_block_name: str,\n candidate_block_name: str,\n local_name_map: dict[str, str | list[str]],\n convert_map: dict[str, str],\n default_dtype: torch.dtype = torch.bfloat16,\n) -> dict[str, torch.Tensor]:\n print(f\"Converting LoRA branch for block {candidate_block_name}...\")\n converted: dict[str, torch.Tensor] = {}\n packer = NunchakuWeightPacker(bits=4)\n for converted_local_name, candidate_local_names in tqdm.tqdm(\n local_name_map.items(), desc=f\"Converting {candidate_block_name}\", dynamic_ncols=True\n ):\n if isinstance(candidate_local_names, str):\n candidate_local_names = [candidate_local_names]\n # region original LoRA\n orig_lora = (\n orig_state_dict.get(f\"{converted_block_name}.{converted_local_name}.lora_down\", None),\n orig_state_dict.get(f\"{converted_block_name}.{converted_local_name}.lora_up\", None),\n )\n if orig_lora[0] is None or orig_lora[1] is None:\n assert orig_lora[0] is None and orig_lora[1] is None\n orig_lora = None\n else:\n assert orig_lora[0] is not None and orig_lora[1] is not None\n orig_lora = (\n packer.unpack_lowrank_weight(orig_lora[0], down=True),\n packer.unpack_lowrank_weight(orig_lora[1], down=False),\n )\n print(f\" - Found {converted_block_name} LoRA of {converted_local_name} (rank: {orig_lora[0].shape[0]})\")\n # endregion\n # region extra LoRA\n extra_lora = [\n (\n extra_lora_dict.get(f\"{candidate_block_name}.{candidate_local_name}.lora_A.weight\", None),\n extra_lora_dict.get(f\"{candidate_block_name}.{candidate_local_name}.lora_B.weight\", None),\n )\n for candidate_local_name in candidate_local_names\n ]\n # if any of the extra LoRA is None, all of them should be None\n if any(lora[0] is not None or lora[1] is not None for lora in extra_lora):\n # merge extra LoRAs into one LoRA\n if len(extra_lora) > 1:\n first_lora = None\n for lora in extra_lora:\n if lora[0] is not None:\n assert lora[1] is not None\n first_lora = lora\n break\n assert first_lora is not None\n for lora_index in range(len(extra_lora)):\n if extra_lora[lora_index][0] is None:\n assert extra_lora[lora_index][1] is None\n extra_lora[lora_index] = (first_lora[0].clone(), torch.zeros_like(first_lora[1]))\n if all(lora[0].equal(extra_lora[0][0]) for lora in extra_lora):\n # if all extra LoRAs have the same lora_down, use it\n extra_lora_down = extra_lora[0][0]\n extra_lora_up = torch.cat([lora[1] for lora in extra_lora], dim=0)\n else:\n extra_lora_down = torch.cat([lora[0] for lora in extra_lora], dim=0)\n extra_lora_up_c = sum(lora[1].shape[0] for lora in extra_lora)\n extra_lora_up_r = sum(lora[1].shape[1] for lora in extra_lora)\n assert extra_lora_up_r == extra_lora_down.shape[0]\n extra_lora_up = torch.zeros((extra_lora_up_c, extra_lora_up_r), dtype=extra_lora_down.dtype)\n c, r = 0, 0\n for lora in extra_lora:\n c_next, r_next = c + lora[1].shape[0], r + lora[1].shape[1]\n extra_lora_up[c:c_next, r:r_next] = lora[1]\n c, r = c_next, r_next\n else:\n extra_lora_down, extra_lora_up = extra_lora[0]\n extra_lora: tuple[torch.Tensor, torch.Tensor] = (extra_lora_down, extra_lora_up)\n print(f\" - Found {candidate_block_name} LoRA of {candidate_local_names} (rank: {extra_lora[0].shape[0]})\")\n # endregion\n # region merge LoRA\n if orig_lora is None:\n if extra_lora is None:\n lora = None\n else:\n print(\" - Using extra LoRA\")\n lora = (extra_lora[0].to(default_dtype), extra_lora[1].to(default_dtype))\n elif extra_lora is None:\n print(\" - Using original LoRA\")\n lora = orig_lora\n else:\n lora = (\n torch.cat([orig_lora[0], extra_lora[0].to(orig_lora[0].dtype)], dim=0),\n torch.cat([orig_lora[1], extra_lora[1].to(orig_lora[1].dtype)], dim=1),\n )\n print(f\" - Merging original and extra LoRA (rank: {lora[0].shape[0]})\")\n # endregion\n if lora is not None:\n if convert_map[converted_local_name] == \"adanorm_single\":\n update_state_dict(\n converted,\n {\n \"lora_down\": pad(lora[0], divisor=16, dim=0),\n \"lora_up\": pad(reorder_adanorm_lora_up(lora[1], splits=3), divisor=16, dim=1),\n },\n prefix=converted_local_name,\n )\n elif convert_map[converted_local_name] == \"adanorm_zero\":\n update_state_dict(\n converted,\n {\n \"lora_down\": pad(lora[0], divisor=16, dim=0),\n \"lora_up\": pad(reorder_adanorm_lora_up(lora[1], splits=6), divisor=16, dim=1),\n },\n prefix=converted_local_name,\n )\n elif convert_map[converted_local_name] == \"linear\":\n update_state_dict(\n converted,\n {\n \"lora_down\": packer.pack_lowrank_weight(lora[0], down=True),\n \"lora_up\": packer.pack_lowrank_weight(lora[1], down=False),\n },\n prefix=converted_local_name,\n )\n return converted\n\n\ndef convert_to_nunchaku_flux_single_transformer_block_lowrank_dict(\n orig_state_dict: dict[str, torch.Tensor],\n extra_lora_dict: dict[str, torch.Tensor],\n converted_block_name: str,\n candidate_block_name: str,\n default_dtype: torch.dtype = torch.bfloat16,\n) -> dict[str, torch.Tensor]:\n if f\"{candidate_block_name}.proj_out.lora_A.weight\" in extra_lora_dict:\n assert f\"{converted_block_name}.out_proj.qweight\" in orig_state_dict\n assert f\"{converted_block_name}.mlp_fc2.qweight\" in orig_state_dict\n n1 = orig_state_dict[f\"{converted_block_name}.out_proj.qweight\"].shape[1] * 2\n n2 = orig_state_dict[f\"{converted_block_name}.mlp_fc2.qweight\"].shape[1] * 2\n lora_down = extra_lora_dict[f\"{candidate_block_name}.proj_out.lora_A.weight\"]\n lora_up = extra_lora_dict[f\"{candidate_block_name}.proj_out.lora_B.weight\"]\n assert lora_down.shape[1] == n1 + n2\n extra_lora_dict[f\"{candidate_block_name}.proj_out.linears.0.lora_A.weight\"] = lora_down[:, :n1].clone()\n extra_lora_dict[f\"{candidate_block_name}.proj_out.linears.0.lora_B.weight\"] = lora_up.clone()\n extra_lora_dict[f\"{candidate_block_name}.proj_out.linears.1.lora_A.weight\"] = lora_down[:, n1:].clone()\n extra_lora_dict[f\"{candidate_block_name}.proj_out.linears.1.lora_B.weight\"] = lora_up.clone()\n extra_lora_dict.pop(f\"{candidate_block_name}.proj_out.lora_A.weight\")\n extra_lora_dict.pop(f\"{candidate_block_name}.proj_out.lora_B.weight\")\n\n return convert_to_nunchaku_transformer_block_lowrank_dict(\n orig_state_dict=orig_state_dict,\n extra_lora_dict=extra_lora_dict,\n converted_block_name=converted_block_name,\n candidate_block_name=candidate_block_name,\n local_name_map={\n \"norm.linear\": \"norm.linear\",\n \"qkv_proj\": [\"attn.to_q\", \"attn.to_k\", \"attn.to_v\"],\n \"norm_q\": \"attn.norm_q\",\n \"norm_k\": \"attn.norm_k\",\n \"out_proj\": \"proj_out.linears.0\",\n \"mlp_fc1\": \"proj_mlp\",\n \"mlp_fc2\": \"proj_out.linears.1\",\n },\n convert_map={\n \"norm.linear\": \"adanorm_single\",\n \"qkv_proj\": \"linear\",\n \"out_proj\": \"linear\",\n \"mlp_fc1\": \"linear\",\n \"mlp_fc2\": \"linear\",\n },\n default_dtype=default_dtype,\n )\n\n\ndef convert_to_nunchaku_flux_transformer_block_lowrank_dict(\n orig_state_dict: dict[str, torch.Tensor],\n extra_lora_dict: dict[str, torch.Tensor],\n converted_block_name: str,\n candidate_block_name: str,\n default_dtype: torch.dtype = torch.bfloat16,\n) -> dict[str, torch.Tensor]:\n return convert_to_nunchaku_transformer_block_lowrank_dict(\n orig_state_dict=orig_state_dict,\n extra_lora_dict=extra_lora_dict,\n converted_block_name=converted_block_name,\n candidate_block_name=candidate_block_name,\n local_name_map={\n \"norm1.linear\": \"norm1.linear\",\n \"norm1_context.linear\": \"norm1_context.linear\",\n \"qkv_proj\": [\"attn.to_q\", \"attn.to_k\", \"attn.to_v\"],\n \"qkv_proj_context\": [\"attn.add_q_proj\", \"attn.add_k_proj\", \"attn.add_v_proj\"],\n \"norm_q\": \"attn.norm_q\",\n \"norm_k\": \"attn.norm_k\",\n \"norm_added_q\": \"attn.norm_added_q\",\n \"norm_added_k\": \"attn.norm_added_k\",\n \"out_proj\": \"attn.to_out.0\",\n \"out_proj_context\": \"attn.to_add_out\",\n \"mlp_fc1\": \"ff.net.0.proj\",\n \"mlp_fc2\": \"ff.net.2\",\n \"mlp_context_fc1\": \"ff_context.net.0.proj\",\n \"mlp_context_fc2\": \"ff_context.net.2\",\n },\n convert_map={\n \"norm1.linear\": \"adanorm_zero\",\n \"norm1_context.linear\": \"adanorm_zero\",\n \"qkv_proj\": \"linear\",\n \"qkv_proj_context\": \"linear\",\n \"out_proj\": \"linear\",\n \"out_proj_context\": \"linear\",\n \"mlp_fc1\": \"linear\",\n \"mlp_fc2\": \"linear\",\n \"mlp_context_fc1\": \"linear\",\n \"mlp_context_fc2\": \"linear\",\n },\n default_dtype=default_dtype,\n )\n\n\ndef convert_to_nunchaku_flux_lowrank_dict(\n orig_state_dict: dict[str, torch.Tensor],\n extra_lora_dict: dict[str, torch.Tensor],\n default_dtype: torch.dtype = torch.bfloat16,\n) -> dict[str, torch.Tensor]:\n block_names: set[str] = set()\n for param_name in orig_state_dict.keys():\n if param_name.startswith((\"transformer_blocks.\", \"single_transformer_blocks.\")):\n block_names.add(\".\".join(param_name.split(\".\")[:2]))\n block_names = sorted(block_names, key=lambda x: (x.split(\".\")[0], int(x.split(\".\")[-1])))\n print(f\"Converting {len(block_names)} transformer blocks...\")\n converted: dict[str, torch.Tensor] = {}\n for block_name in block_names:\n if block_name.startswith(\"transformer_blocks\"):\n convert_fn = convert_to_nunchaku_flux_transformer_block_lowrank_dict\n else:\n convert_fn = convert_to_nunchaku_flux_single_transformer_block_lowrank_dict\n update_state_dict(\n converted,\n convert_fn(\n orig_state_dict=orig_state_dict,\n extra_lora_dict=extra_lora_dict,\n converted_block_name=block_name,\n candidate_block_name=block_name,\n default_dtype=default_dtype,\n ),\n prefix=block_name,\n )\n return converted\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--quant-path\", type=str, required=True, help=\"path to the quantized model safetensor file\")\n parser.add_argument(\"--lora-path\", type=str, required=True, help=\"path to LoRA weights safetensor file\")\n parser.add_argument(\"--output-root\", type=str, default=\"\", help=\"root to the output safetensor file\")\n parser.add_argument(\"--lora-name\", type=str, default=None, help=\"name of the LoRA weights\")\n parser.add_argument(\n \"--dtype\",\n type=str,\n default=\"bfloat16\",\n choices=[\"bfloat16\", \"float16\"],\n help=\"default data type of the converted LoRA weights\",\n )\n args = parser.parse_args()\n\n if not args.output_root:\n # output to the parent directory of the quantized model safetensor file\n args.output_root = os.path.dirname(args.quant_path)\n if args.lora_name is None:\n assert args.lora_path is not None, \"LoRA name or path must be provided\"\n lora_name = args.lora_path.rstrip(os.sep).split(os.sep)[-1].replace(\".safetensors\", \"\")\n print(f\"Lora name not provided, using {lora_name} as the LoRA name\")\n else:\n lora_name = args.lora_name\n assert lora_name, \"LoRA name must be provided.\"\n\n assert args.quant_path.endswith(\".safetensors\"), \"Quantized model must be a safetensor file\"\n assert args.lora_path.endswith(\".safetensors\"), \"LoRA weights must be a safetensor file\"\n orig_state_dict = load_state_dict_in_safetensors(args.quant_path)\n extra_lora_dict = load_state_dict_in_safetensors(args.lora_path, filter_prefix=\"transformer.\")\n converted = convert_to_nunchaku_flux_lowrank_dict(\n orig_state_dict=orig_state_dict,\n extra_lora_dict=extra_lora_dict,\n default_dtype=torch.bfloat16 if args.dtype == \"bfloat16\" else torch.float16,\n )\n os.makedirs(args.output_root, exist_ok=True)\n safetensors.torch.save_file(converted, os.path.join(args.output_root, f\"{lora_name}.safetensors\"))\n print(f\"Saved LoRA weights to {args.output_root}.\")\n"
},
{
"path": "deepcompressor/backend/nunchaku/utils.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Nunchaku backend utilities.\"\"\"\n\nimport torch\n\nfrom ..tinychat.utils import convert_to_tinychat_w4x16y16_linear_weight\nfrom ..utils import MmaWeightPackerBase, ceil_divide, fp_quantize, pad\n\n__all__ = [\n \"convert_to_nunchaku_w4x4y16_linear_weight\",\n \"convert_to_nunchaku_w8x8y16_linear_weight\",\n \"convert_to_nunchaku_w4x16_linear_weight\",\n]\n\n\nclass NunchakuWeightPacker(MmaWeightPackerBase):\n def __init__(self, bits: int, warp_n: int = 128):\n super().__init__(bits=bits, warp_n=warp_n)\n self.num_k_unrolls = 2\n\n def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:\n assert weight.dtype == torch.int32, f\"quantized weight should be torch.int32, but got {weight.dtype}.\"\n n, k = weight.shape\n assert n % self.mem_n == 0, f\"output channel size ({n}) should be divisible by mem_n ({self.mem_n}).\"\n # currently, Nunchaku did not check the boundry of unrolled `k` dimension\n assert k % (self.mem_k * self.num_k_unrolls) == 0, (\n f\"input channel size ({k}) should be divisible by \"\n f\"mem_k ({self.mem_k}) * num_k_unrolls ({self.num_k_unrolls}).\"\n )\n n_tiles, k_tiles = n // self.mem_n, k // self.mem_k\n weight = weight.reshape(\n n_tiles,\n self.num_n_packs, # 8 when warp_n = 128\n self.n_pack_size, # always 2 in nunchaku\n self.num_n_lanes, # constant 8\n self.reg_n, # constant 1\n k_tiles,\n self.num_k_packs, # 1\n self.k_pack_size, # always 2 in nunchaku\n self.num_k_lanes, # constant 4\n self.reg_k, # always 8 = 32 bits / 4 bits\n )\n # (n_tiles, num_n_packs, n_pack_size, num_n_lanes, reg_n, k_tiles, num_k_packs, k_pack_size, num_k_lanes, reg_k)\n # =>\n # (n_tiles, k_tiles, num_k_packs, num_n_packs, num_n_lanes, num_k_lanes, n_pack_size, k_pack_size, reg_n, reg_k)\n weight = weight.permute(0, 5, 6, 1, 3, 8, 2, 7, 4, 9).contiguous()\n assert weight.shape[4:-2] == (8, 4, 2, 2)\n if self.bits == 4:\n weight = weight.bitwise_and_(0xF)\n shift = torch.arange(0, 32, 4, dtype=torch.int32, device=weight.device)\n weight = weight.bitwise_left_shift_(shift)\n weight = weight.sum(dim=-1, dtype=torch.int32)\n elif self.bits == 8:\n weight = weight.bitwise_and_(0xFF)\n shift = torch.arange(0, 32, 8, dtype=torch.int32, device=weight.device)\n weight = weight.bitwise_left_shift_(shift)\n weight = weight.sum(dim=-1, dtype=torch.int32)\n else:\n raise NotImplementedError(f\"weight bits {self.bits} is not supported.\")\n return weight.view(dtype=torch.int8).view(n, -1) # assume little-endian\n\n def pack_scale(self, scale: torch.Tensor, group_size: int) -> torch.Tensor:\n if self.check_if_micro_scale(group_size=group_size):\n return self.pack_micro_scale(scale, group_size=group_size)\n # note: refer to https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#mma-16864-c\n assert scale.dtype in (torch.float16, torch.bfloat16), \"currently nunchaku only supports fp16 and bf16.\"\n n = scale.shape[0]\n # nunchaku load scales all in one access\n # for `[warp_n, warp_k]` weights, we load `[warp_n, warp_k / group_size]` scales\n # scale loading is parallelized in `n` dimension, that is,\n # `num_s_lanes` in a warp load `num_s_packs` of `s_pack_size` elements, in total `warp_s` elements\n # each element in `n` dimension is 16 bit as it contains 1 fp16\n # min `s_pack_size` set to 2 element, since each lane at least holds 2 accumulator results in `n` dimension\n # max `s_pack_size` set to 128b/16b = 8 elements\n # for `warp_n = 8`, we have\n # `s_pack_size = 2`, `num_s_lanes = 4`, `num_s_packs = 1`\n # for `warp_n = 128`, we have\n # `s_pack_size = 4`, `num_s_lanes = 32`, `num_s_packs = 1`\n # for `warp_n = 512`, we have\n # `s_pack_size = 8`, `num_s_lanes = 32`, `num_s_packs = 2`\n s_pack_size = min(max(self.warp_n // self.num_lanes, 2), 8)\n num_s_lanes = min(self.num_lanes, self.warp_n // s_pack_size)\n num_s_packs = self.warp_n // (s_pack_size * num_s_lanes)\n warp_s = num_s_packs * num_s_lanes * s_pack_size\n assert warp_s == self.warp_n, \"warp_n for scales should be equal to warp_n for weights.\"\n # `num_n_lanes = 8 (constant)` generates 8 elements consecutive in `n` dimension\n # however, they are held by 4 lanes, each lane holds 2 elements in `n` dimension\n # thus, we start from first 4 lanes, assign 2 elements to each lane, until all 8 elements are assigned\n # we then repeat the process for the same 4 lanes, until each lane holds `s_pack_size` elements\n # finally, we move to next 4 lanes, and repeat the process until all `num_s_lanes` lanes are assigned\n # the process is repeated for `num_s_packs` times\n # here is an example for `warp_n = 128, s_pack_size = 4, num_s_lanes = 32, num_s_packs = 1`\n # wscales store order:\n # 0 1 8 9 <-- load by lane 0, broadcast to lane {0, 4, 8, ..., 28} (8x)\n # 2 3 10 11 <-- load by lane 1, broadcast to lane {1, 5, 9, ..., 29} (8x)\n # 4 5 12 13 <-- load by lane 2, broadcast to lane {2, 6, 10, ..., 30} (8x)\n # 6 7 14 15 <-- load by lane 3, broadcast to lane {3, 7, 11, ..., 31} (8x)\n # 16 17 24 25 <-- load by lane 4, broadcast to lane {0, 4, 8, ..., 28} (8x)\n # ...\n # 22 23 30 31 <-- load by lane 7, broadcast to lane {3, 7, 11, ..., 31} (8x)\n # ... ...\n # 112 113 120 121 <-- load by lane 28, broadcast to lane {0, 4, 8, ..., 28} (8x)\n # ...\n # 118 119 126 127 <-- load by lane 31, broadcast to lane {3, 7, 11, ..., 31} (8x)\n scale = scale.reshape(n // warp_s, num_s_packs, num_s_lanes // 4, s_pack_size // 2, 4, 2, -1)\n scale = scale.permute(0, 6, 1, 2, 4, 3, 5).contiguous()\n return scale.view(-1) if group_size == -1 else scale.view(-1, n) # the shape is just used for validation\n\n def pack_micro_scale(self, scale: torch.Tensor, group_size: int) -> torch.Tensor:\n assert scale.dtype in (torch.float16, torch.bfloat16), \"currently nunchaku only supports fp16 and bf16.\"\n assert scale.max() <= 448, \"scale should be less than 448.\"\n assert scale.min() >= -448, \"scale should be greater than -448.\"\n assert group_size == 16, \"currently only support group size 16.\"\n assert self.insn_k == 64, \"insn_k should be 64.\"\n scale = scale.to(dtype=torch.float8_e4m3fn)\n n = scale.shape[0]\n assert self.warp_n >= 32, \"currently only support warp_n >= 32.\"\n # for `[warp_n, warp_k]` weights, we load `[warp_n, warp_k / group_size]` scales\n # scale loading is parallelized in `n` dimension, that is,\n # `num_s_lanes` in a warp load `num_s_packs` of `s_pack_size` elements, in total `warp_s` elements\n # each element in `n` dimension is 32 bit as it contains 4 fp8 in `k` dimension\n # min `s_pack_size` set to 1 element\n # max `s_pack_size` set to 128b/32b = 4 elements\n # for `warp_n = 128`, we have\n # `s_pack_size = 4`, `num_s_lanes = 32`, `num_s_packs = 1`\n # for `warp_n = 512`, we have\n # `s_pack_size = 8`, `num_s_lanes = 32`, `num_s_packs = 2`\n s_pack_size = min(max(self.warp_n // self.num_lanes, 1), 4)\n num_s_lanes = 4 * 8 # 32 lanes is divided into 4 pieces, each piece has 8 lanes at a stride of 4\n num_s_packs = ceil_divide(self.warp_n, s_pack_size * num_s_lanes)\n warp_s = num_s_packs * num_s_lanes * s_pack_size\n assert warp_s == self.warp_n, \"warp_n for scales should be equal to warp_n for weights.\"\n # note: refer to https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#mma-scaling-thread-id-b-selection\n # we start from first 8 lines at a stride of 4, assign 1 element to each lane, until all 8 elements are assigned\n # we then move to next 8 lines at a stride of 4, and repeat the process until all 32 lanes are assigned\n # here is an example for `warp_n = 128, s_pack_size = 4, num_s_lanes = 32, num_s_packs = 1`\n # wscales store order:\n # 0 32 64 96 <-- load by lane 0\n # 8 40 72 104 <-- load by lane 1\n # 16 48 80 112 <-- load by lane 2\n # 24 56 88 120 <-- load by lane 3\n # 1 33 65 97 <-- load by lane 4\n # ...\n # 25 57 81 113 <-- load by lane 7\n # ...\n # 7 39 71 103 <-- load by lane 28\n # ...\n # 31 63 95 127 <-- load by lane 31\n scale = scale.view(n // warp_s, num_s_packs, s_pack_size, 4, 8, -1, self.insn_k // group_size)\n scale = scale.permute(0, 5, 1, 4, 3, 2, 6).contiguous()\n return scale.view(-1, n) # the shape is just used for validation\n\n def pack_lowrank_weight(self, weight: torch.Tensor, down: bool) -> torch.Tensor:\n \"\"\"Pack Low-Rank Weight.\n\n Args:\n weight (`torch.Tensor`):\n low-rank weight tensor.\n down (`bool`):\n whether the weight is for down projection in low-rank branch.\n \"\"\"\n assert weight.dtype in (torch.float16, torch.bfloat16), f\"Unsupported weight dtype {weight.dtype}.\"\n reg_n, reg_k = 1, 2 # reg_n is always 1, reg_k is 32 bits // 16 bits = 2\n pack_n = self.n_pack_size * self.num_n_lanes * reg_n\n pack_k = self.k_pack_size * self.num_k_lanes * reg_k\n weight = pad(weight, divisor=(pack_n, pack_k), dim=(0, 1))\n if down:\n r, c = weight.shape\n r_packs, c_packs = r // pack_n, c // pack_k\n weight = weight.view(r_packs, pack_n, c_packs, pack_k).permute(2, 0, 1, 3)\n else:\n c, r = weight.shape\n c_packs, r_packs = c // pack_n, r // pack_k\n weight = weight.view(c_packs, pack_n, r_packs, pack_k).permute(0, 2, 1, 3)\n weight = weight.reshape(\n c_packs, r_packs, self.n_pack_size, self.num_n_lanes, reg_n, self.k_pack_size, self.num_k_lanes, reg_k\n )\n # (c_packs, r_packs, n_pack_size, num_n_lanes, reg_n, k_pack_size, num_k_lanes, reg_k)\n # =>\n # (c_packs, r_packs, num_n_lanes, num_k_lanes, n_pack_size, k_pack_size, reg_n, reg_k)\n weight = weight.permute(0, 1, 3, 6, 2, 5, 4, 7).contiguous()\n return weight.view(c, r)\n\n def unpack_lowrank_weight(self, weight: torch.Tensor, down: bool) -> torch.Tensor:\n \"\"\"Unpack Low-Rank Weight.\n\n Args:\n weight (`torch.Tensor`):\n low-rank weight tensor.\n down (`bool`):\n whether the weight is for down projection in low-rank branch.\n \"\"\"\n c, r = weight.shape\n assert weight.dtype in (torch.float16, torch.bfloat16), f\"Unsupported weight dtype {weight.dtype}.\"\n reg_n, reg_k = 1, 2 # reg_n is always 1, reg_k is 32 bits // 16 bits = 2\n pack_n = self.n_pack_size * self.num_n_lanes * reg_n\n pack_k = self.k_pack_size * self.num_k_lanes * reg_k\n if down:\n r_packs, c_packs = r // pack_n, c // pack_k\n else:\n c_packs, r_packs = c // pack_n, r // pack_k\n weight = weight.view(\n c_packs, r_packs, self.num_n_lanes, self.num_k_lanes, self.n_pack_size, self.k_pack_size, reg_n, reg_k\n )\n # (c_packs, r_packs, num_n_lanes, num_k_lanes, n_pack_size, k_pack_size, reg_n, reg_k)\n # =>\n # (c_packs, r_packs, n_pack_size, num_n_lanes, reg_n, k_pack_size, num_k_lanes, reg_k)\n weight = weight.permute(0, 1, 4, 2, 6, 5, 3, 7).contiguous()\n weight = weight.view(c_packs, r_packs, pack_n, pack_k)\n if down:\n weight = weight.permute(1, 2, 0, 3).contiguous().view(r, c)\n else:\n weight = weight.permute(0, 2, 1, 3).contiguous().view(c, r)\n return weight\n\n def check_if_micro_scale(self, group_size: int) -> bool:\n return self.insn_k == group_size * 4\n\n def pad_weight(self, weight: torch.Tensor) -> torch.Tensor:\n assert weight.ndim == 2, \"weight tensor should be 2D.\"\n return pad(weight, divisor=(self.mem_n, self.mem_k * self.num_k_unrolls), dim=(0, 1))\n\n def pad_scale(self, scale: torch.Tensor, group_size: int) -> torch.Tensor:\n if group_size > 0 and scale.numel() > scale.shape[0]:\n scale = scale.view(scale.shape[0], 1, -1, 1)\n if self.check_if_micro_scale(group_size=group_size):\n scale = pad(scale, divisor=(self.warp_n, self.insn_k // group_size), dim=(0, 2), fill_value=1)\n else:\n scale = pad(scale, divisor=(self.warp_n, self.num_k_unrolls), dim=(0, 2), fill_value=1)\n else:\n scale = pad(scale, divisor=self.warp_n, dim=0, fill_value=1)\n return scale\n\n def pad_lowrank_weight(self, weight: torch.Tensor, down: bool) -> torch.Tensor:\n assert weight.ndim == 2, \"weight tensor should be 2D.\"\n return pad(weight, divisor=self.warp_n, dim=1 if down else 0)\n\n\ndef convert_to_nunchaku_w4x4y16_linear_weight(\n weight: torch.Tensor,\n scale: torch.Tensor,\n bias: torch.Tensor | None = None,\n smooth: torch.Tensor | None = None,\n lora: tuple[torch.Tensor, torch.Tensor] | None = None,\n float_point: bool = False,\n subscale: torch.Tensor | None = None,\n) -> tuple[\n torch.Tensor,\n torch.Tensor,\n torch.Tensor,\n torch.Tensor,\n tuple[torch.Tensor, torch.Tensor] | None,\n torch.Tensor | None,\n]:\n assert weight.ndim == 2, \"weight tensor should be 2D.\"\n device, dtype = weight.device, weight.dtype\n assert dtype in (torch.float16, torch.bfloat16), \"currently nunchaku only supports fp16 and bf16.\"\n assert scale is not None, \"scale tensor is required for quantization.\"\n\n oc, ic = weight.shape\n if scale.numel() == 1:\n scale = scale.view(-1).expand(oc).reshape(oc, 1, 1, 1)\n per_tensor_scale = True\n else:\n per_tensor_scale = False\n assert scale.ndim == 4, \"scale tensor should be 4D.\"\n assert scale.shape[1] == scale.shape[3] == 1\n assert scale.shape[0] == oc\n ng, gs = scale.shape[2], ic // scale.shape[2]\n assert ic == gs * ng, \"input channel size should be equal to group size times number of groups.\"\n if subscale is not None:\n assert subscale.ndim == 4, \"subscale tensor should be 4D.\"\n assert subscale.shape[1] == subscale.shape[3] == 1\n assert subscale.shape[0] == oc\n nsg, sgs = subscale.shape[2], ic // subscale.shape[2]\n assert ic == sgs * nsg, \"input channel size should be equal to subgroup size times number of subgroups.\"\n assert gs > sgs and gs % sgs == 0, \"group size should be divisible by subgroup size.\"\n else:\n nsg, sgs = ng, gs\n # region quantize and pack weight tensor\n weight = weight.to(dtype=torch.float32).view(oc, 1, ng, gs).div_(scale.to(dtype=torch.float32, device=device))\n if subscale is not None:\n weight = weight.view(oc, 1, nsg, sgs).div_(subscale.to(dtype=torch.float32, device=device))\n weight = weight.view(oc, ic)\n if float_point:\n weight = fp_quantize(weight)\n assert weight.min() >= 0 and weight.max() <= 15, \"quantized weight should be in [0, 15].\"\n else:\n weight = weight.round_()\n assert weight.min() >= -8 and weight.max() <= 7, \"quantized weight should be in [-8, 7].\"\n # endregion\n bias = torch.zeros([oc, 1], dtype=dtype, device=device) if bias is None else bias.view(-1, 1)\n smooth = torch.ones([ic, 1], dtype=dtype, device=device) if smooth is None else smooth.view(-1, 1)\n\n packer = NunchakuWeightPacker(bits=4)\n weight = packer.pad_weight(weight.to(dtype=torch.int32))\n scale = packer.pad_scale(scale.to(dtype=dtype), group_size=gs)\n if subscale is not None:\n subscale = packer.pad_scale(subscale.to(dtype=dtype), group_size=sgs)\n bias = packer.pad_scale(bias.to(dtype=dtype), group_size=-1)\n smooth = packer.pad_scale(smooth.to(dtype=dtype), group_size=-1)\n\n weight = packer.pack_weight(weight)\n scale = packer.pack_scale(scale, group_size=gs if gs < ic else -1)\n if subscale is not None:\n subscale = packer.pack_scale(subscale, group_size=sgs if sgs < ic else -1)\n bias = packer.pack_scale(bias, group_size=-1)\n smooth = packer.pack_scale(smooth, group_size=-1)\n if lora is not None:\n lora_down = packer.pack_lowrank_weight(packer.pad_lowrank_weight(lora[0], down=True), down=True)\n lora_up = packer.pack_lowrank_weight(packer.pad_lowrank_weight(lora[1], down=False), down=False)\n lora = (lora_down, lora_up)\n if per_tensor_scale:\n scale = scale.view(-1)[0].view([1])\n return weight, scale, bias, smooth, lora, subscale\n\n\ndef convert_to_nunchaku_w8x8y16_linear_weight(\n weight: torch.Tensor, scale: torch.Tensor, bias: torch.Tensor | None = None\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n assert weight.ndim == 2, \"weight tensor should be 2D.\"\n device, dtype = weight.device, weight.dtype\n assert dtype in (torch.float16, torch.bfloat16), \"currently nunchaku only supports fp16 and bf16.\"\n assert scale is not None, \"scale tensor is required for quantization.\"\n oc, ic = weight.shape\n if scale.numel() == 1:\n scale = scale.view(-1).expand(oc)\n scale = scale.reshape(oc, 1)\n weight = weight.to(dtype=torch.float32)\n weight = weight.div_(scale.to(dtype=torch.float32, device=device)).round_().to(torch.int32).view(oc, ic)\n assert weight.min() >= -128 and weight.max() <= 127, \"quantized weight should be in [-128, 127].\"\n # endregion\n bias = torch.zeros([oc, 1], dtype=dtype, device=device) if bias is None else bias.view(-1, 1)\n packer = NunchakuWeightPacker(bits=8)\n weight = packer.pack_weight(packer.pad_weight(weight))\n scale = packer.pack_scale(packer.pad_scale(scale.to(dtype=dtype), group_size=-1), group_size=-1)\n bias = packer.pack_scale(packer.pad_scale(bias.to(dtype=dtype), group_size=-1), group_size=-1).view(-1)\n return weight, scale, bias\n\n\ndef convert_to_nunchaku_w4x16_linear_weight(\n weight: torch.Tensor,\n scale: torch.Tensor,\n zero: torch.Tensor | None = None,\n bias: torch.Tensor | None = None,\n adanorm_splits: int = 1,\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n oc, ic = weight.shape\n assert scale.ndim == 4, \"scale tensor should be 4D.\"\n assert scale.shape[0] == oc\n assert scale.shape[1] == scale.shape[3] == 1\n ng = scale.shape[2]\n if bias is None:\n bias = torch.zeros([oc], dtype=weight.dtype, device=weight.device)\n assert oc % adanorm_splits == 0, \"output channel size should be divisible by splits.\"\n if adanorm_splits > 1:\n weight = weight.view(adanorm_splits, oc // adanorm_splits, ic).transpose(0, 1).reshape(oc, ic)\n scale = scale.view(adanorm_splits, oc // adanorm_splits, ng).transpose(0, 1).reshape(oc, 1, ng, 1)\n bias = bias.reshape(adanorm_splits, oc // adanorm_splits).transpose(0, 1)\n delta = [0] * adanorm_splits\n delta[1] = delta[-2] = 1\n bias = bias.add_(torch.tensor(delta, dtype=bias.dtype, device=bias.device))\n bias = bias.reshape(oc)\n weight, scale, zero = convert_to_tinychat_w4x16y16_linear_weight(\n weight=weight, scale=scale, zero=torch.full_like(scale, 7) if zero is None else zero, zero_pre_scaled=True\n )\n weight = weight.view(torch.int32)\n return weight, scale, zero, bias\n"
},
{
"path": "deepcompressor/backend/qserve/__init__.py",
"content": ""
},
{
"path": "deepcompressor/backend/qserve/convert.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"QServe state dict converter module.\"\"\"\n\nimport argparse\nimport os\n\nimport torch\nimport tqdm\n\nfrom .utils import convert_to_qserve_w4x8y16_linear_weight, convert_to_qserve_w8x8y16_linear_weight\n\n__all__ = [\"convert_to_qserve_state_dict\"]\n\n\ndef convert_to_qserve_w4x8y16_linear_state_dict(\n param_name: str,\n weight: torch.Tensor,\n scale: torch.Tensor,\n zero: torch.Tensor,\n subscale: torch.Tensor | None = None,\n zero_pre_scaled: bool = False,\n) -> dict[str, torch.Tensor]:\n \"\"\"Convert a weight tensor to QServe W4-X8-Y16 linear state dictionary.\n\n Args:\n param_name (`str`):\n parameter name.\n weight (`torch.Tensor`):\n weight tensor to be converted.\n scale (`torch.Tensor`):\n scale tensor for the weight tensor.\n zero (`torch.Tensor`):\n zero point tensor for the weight tensor.\n subscale (`torch.Tensor` or `None`, *optional*, defaults to `None`):\n subscale tensor for the weight tensor.\n zero_pre_scaled (`bool`, *optional*, defaults to `False`):\n whether zero point tensor is pre-scaled.\n\n Returns:\n `dict[str, torch.Tensor]`:\n state dictionary for the quantized weight tensor.\n \"\"\"\n module_name = param_name[:-7]\n weight, scale, zero, subscale = convert_to_qserve_w4x8y16_linear_weight(\n weight, scale=scale, zero=zero, subscale=subscale, zero_pre_scaled=zero_pre_scaled\n )\n state_dict: dict[str, torch.Tensor] = {}\n state_dict[f\"{module_name}.qweight\"] = weight.cpu()\n state_dict[f\"{module_name}.s1_scales\"] = scale.cpu()\n if subscale is None:\n state_dict[f\"{module_name}.s1_szeros\"] = zero.cpu()\n else:\n state_dict[f\"{module_name}.s2_scales\"] = subscale.cpu()\n state_dict[f\"{module_name}.s2_zeros\"] = zero.cpu()\n return state_dict\n\n\ndef convert_to_qserve_w8x8y16_linear_state_dict(\n param_name: str, weight: torch.Tensor, scale: torch.Tensor\n) -> dict[str, torch.Tensor]:\n \"\"\"Convert a weight tensor to QServe W8-X8-Y16 linear state dictionary.\n\n Args:\n param_name (`str`):\n parameter name.\n weight (`torch.Tensor`):\n weight tensor to be converted.\n scale (`torch.Tensor`):\n scale tensor for the weight tensor.\n\n Returns:\n `dict[str, torch.Tensor]`:\n state dictionary for the quantized weight tensor.\n \"\"\"\n module_name = param_name[:-7]\n weight, scale = convert_to_qserve_w8x8y16_linear_weight(weight, scale=scale)\n state_dict: dict[str, torch.Tensor] = {}\n state_dict[f\"{module_name}.weight\"] = weight.cpu()\n state_dict[f\"{module_name}.dequant_scale\"] = scale.cpu()\n return state_dict\n\n\ndef convert_to_qserve_state_dict(\n state_dict: dict[str, torch.Tensor], scale_dict: dict[str, torch.Tensor], weight_bits: int\n) -> dict[str, torch.Tensor]:\n assert weight_bits in [4, 8], \"weight bits should be 4 or 8.\"\n scales: dict[str, dict[tuple[int, ...], torch.Tensor]] = {}\n zeros: dict[str, tuple[torch.Tensor | None, bool]] = {}\n print(\"Loading scale tensors...\")\n for name, tensor in tqdm.tqdm(scale_dict.items(), desc=\"Loading scale tensors\", leave=False, dynamic_ncols=True):\n print(f\" - Loading tensor {name} (dtype: {tensor.dtype}, shape: {tensor.shape}, device: {tensor.device})\")\n if name.endswith(\"zero\"):\n # this is a zero point tensor\n zero = None if tensor is None or all(t.item() == 0 for t in tensor.flatten()) else tensor\n if name.endswith(\".scaled_zero\"):\n zeros[name[:-12]] = (zero, False) # zero point tensor is post-scaled\n else:\n zeros[name[:-5]] = (zero, True) # zero point tensor is pre-scaled\n else:\n assert \".weight.scale\" in name\n # this is a scale tensor\n idx = name.index(\".weight.scale\")\n param_name = name[: idx + 7]\n scale_level = tuple(map(int, name[idx + 14 :].split(\".\")))\n scales.setdefault(param_name, {})[scale_level] = tensor\n for param_name in zeros.keys():\n assert param_name in state_dict, f\"zero point tensor {param_name} not found in state dict.\"\n assert param_name in scales, f\"scale tensor {param_name} not found in scale dict.\"\n converted: dict[str, torch.Tensor] = {}\n print(\"Converting state dict...\")\n for param_name, param in tqdm.tqdm(state_dict.items(), desc=\"Converting state dict\", dynamic_ncols=True):\n if param_name in scales:\n print(f\" - Converting {param_name} (dtype: {param.dtype}, shape: {param.shape}, device: {param.device})\")\n weight = param.data.clone()\n if param_name in zeros:\n zero, zero_pre_scaled = zeros[param_name]\n zero = zero.clone() if zero is not None else None\n else:\n zero, zero_pre_scaled = None, False\n level_scales = sorted(scales[param_name].items(), key=lambda x: x[0])\n assert len(level_scales) <= 2, \"more than two scale levels are not supported.\"\n scale = level_scales[0][1].clone()\n subscale = level_scales[1][1].clone() if len(level_scales) > 1 else None\n if weight_bits == 4:\n converted.update(\n convert_to_qserve_w4x8y16_linear_state_dict(\n param_name,\n weight,\n scale=scale,\n zero=zero,\n subscale=subscale,\n zero_pre_scaled=zero_pre_scaled,\n )\n )\n else:\n assert zero is None, \"zero point tensor is not supported for W8 quantization.\"\n assert subscale is None, \"subscale tensor is not supported for W8 quantization.\"\n converted.update(convert_to_qserve_w8x8y16_linear_state_dict(param_name, weight, scale=scale))\n else:\n if isinstance(param, torch.Tensor):\n print(f\" - Copying {param_name} (dtype: {param.dtype}, shape: {param.shape}, device: {param.device})\")\n converted[param_name] = param.clone().cpu()\n else:\n print(f\" - Copying {param_name} (type: {type(param)}, value: {param})\")\n converted[param_name] = param\n return converted\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--quant-path\", type=str, required=True, help=\"path to the quantization checkpoint directory.\")\n parser.add_argument(\"--weight-bits\", type=int, required=True, help=\"quantized weight bits.\")\n parser.add_argument(\"--output-root\", type=str, default=\"\", help=\"root to the output checkpoint directory.\")\n parser.add_argument(\"--model-name\", type=str, default=None, help=\"name of the model.\")\n parser.add_argument(\"--model-path\", type=str, default=None, help=\"path to the huggingface model directory.\")\n parser.add_argument(\n \"--copy-on-save\",\n action=\"store_true\",\n help=\"copy the original tokenizer and configuration files to the output directory.\",\n )\n args = parser.parse_args()\n if not args.output_root:\n args.output_root = args.quant_path\n if args.model_name is None:\n assert args.model_path is not None, \"model name or path is required.\"\n model_name = args.model_path.rstrip(os.sep).split(os.sep)[-1]\n print(f\"Model name not provided. Using model name {model_name}.\")\n else:\n model_name = args.model_name\n assert model_name, \"model name is required.\"\n model_name = f\"{model_name}-w{args.weight_bits}a8\"\n output_dirpath = os.path.join(args.output_root, model_name)\n output_path = os.path.join(output_dirpath, \"quant_model.pt\")\n state_dict = torch.load(\n os.path.join(args.quant_path, \"model.pt\"),\n map_location=\"cuda\" if torch.cuda.is_available() and torch.cuda.device_count() > 0 else \"cpu\",\n )\n scale_dict = torch.load(os.path.join(args.quant_path, \"scale.pt\"), map_location=\"cpu\")\n converted = convert_to_qserve_state_dict(state_dict, scale_dict, weight_bits=args.weight_bits)\n os.makedirs(output_dirpath, exist_ok=True)\n torch.save(converted, output_path)\n if args.model_path and os.path.exists(args.model_path):\n for filename in os.listdir(args.model_path):\n if filename == \"tokenizer.model\" or (\n filename.endswith(\".json\") and filename != \"pytorch_model.bin.index.json\"\n ):\n filepath = os.path.abspath(os.path.join(args.model_path, filename))\n if args.copy_on_save:\n os.system(f\"cp {filepath} {output_dirpath}/\")\n else:\n os.system(f\"ln -s {filepath} {output_dirpath}/{filename}\")\n print(f\"Quantized model checkpoint saved to {output_path}.\")\n print(f\"Quantized model saved to {output_dirpath}.\")\n print(f\"Quantized model checkpoint saved to {output_path}.\")\n print(f\"Quantized model saved to {output_dirpath}.\")\n"
},
{
"path": "deepcompressor/backend/qserve/utils.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"QServe backend utilities.\"\"\"\n\nimport torch\n\nfrom ..utils import MmaWeightPackerBase\n\n__all__ = [\"convert_to_qserve_w4x8y16_linear_weight\", \"convert_to_qserve_w8x8y16_linear_weight\"]\n\n\nclass QServePacker(MmaWeightPackerBase):\n def __init__(self):\n super().__init__(bits=8, warp_n=32)\n assert self.num_n_packs >= 2 and self.num_n_packs % 2 == 0, (\n f\"num_n_packs should be even, but got {self.num_n_packs}.\"\n )\n\n def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:\n assert weight.min() >= 0, \"quantized weight should be non-negative.\"\n assert weight.max() <= 15, \"quantized weight should be less than 16.\"\n assert weight.dtype == torch.uint8, f\"quantized weight should be torch.uint8, but got {weight.dtype}.\"\n n, k = weight.shape\n assert n % self.mem_n == 0, f\"output channel size ({n}) should be divisible by mem_n ({self.mem_n}).\"\n assert k % self.mem_k == 0, f\"input channel size ({k}) should be divisible by mem_k ({self.mem_k}).\"\n n_tiles, k_tiles = n // self.mem_n, k // self.mem_k\n weight = weight.reshape(\n n_tiles,\n self.num_n_packs, # num_n_packs = 2 when warp_n = 32\n self.n_pack_size, # always 2 in QServe\n self.num_n_lanes, # constant 8\n self.reg_n, # constant 1\n k_tiles,\n self.num_k_packs, # constant 1\n self.k_pack_size, # always 2\n self.num_k_lanes, # constant 4\n self.reg_k, # always 4 = 32 bits / 8 bits in QServe\n )\n # (n_tiles, num_n_packs, n_pack_size, num_n_lanes, reg_n, k_tiles, num_k_packs, k_pack_size, num_k_lanes, reg_k)\n # =>\n # (num_n_packs, n_tiles, k_tiles, num_k_packs, num_n_lanes, num_k_lanes, k_pack_size, n_pack_size, reg_n, reg_k)\n weight = weight.permute(1, 0, 5, 6, 3, 8, 7, 2, 4, 9).contiguous()\n assert weight.shape[4:-2] == (8, 4, 2, 2)\n weight = (weight[1] << 4) + weight[0]\n return weight.view(torch.int8).view(n, k // 2)\n\n def pack_scale(\n self, scale: torch.Tensor, zero: torch.Tensor | None = None, subscale: torch.Tensor | None = None\n ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:\n scale = scale.view(-1)\n n = scale.shape[0]\n if subscale is None:\n zero = zero.view(-1)\n else:\n assert subscale.dtype == torch.int8, f\"subscale should be torch.int8, but got {subscale.dtype}.\"\n view_shape = (n // self.mem_n, self.num_n_packs, self.n_pack_size, self.num_n_lanes, self.reg_n, -1)\n # (n_tiles, num_n_packs, n_pack_size, num_n_lanes, reg_n, -1)\n # =>\n # (-1, n_tiles, num_n_packs, num_n_lanes, n_pack_size, reg_n)\n subscale = subscale.view(view_shape).permute(5, 0, 1, 3, 2, 4).contiguous().view(-1, n)\n zero = zero.view(view_shape).permute(5, 0, 1, 3, 2, 4).contiguous().view(-1, n)\n return scale, zero, subscale\n\n\ndef convert_to_qserve_w4x8y16_linear_weight(\n weight: torch.Tensor,\n scale: torch.Tensor,\n zero: torch.Tensor,\n subscale: torch.Tensor | None = None,\n zero_pre_scaled: bool = False,\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:\n \"\"\"Convert a weight tensor to QServe W4-X8-Y16 linear weight format.\n\n Args:\n weight (`torch.Tensor`):\n weight tensor to be converted.\n scale (`torch.Tensor`):\n scale tensor for the weight tensor.\n zero (`torch.Tensor`):\n zero point tensor for the weight tensor.\n subscale (`torch.Tensor` or `None`, *optional*, defaults to `None`):\n subscale tensor for the weight tensor.\n zero_pre_scaled (`bool`, *optional*, defaults to `False`):\n whether zero point tensor is pre-scaled.\n\n Returns:\n `tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]`:\n packed quantized weight tensor, scale tensor, zero point tensor, and subscale tensor.\n \"\"\"\n dtype = weight.dtype\n assert dtype == torch.float16, \"currently qserve only supports fp16.\"\n assert scale is not None, \"scale tensor is required for quantization.\"\n assert zero is not None, \"zero point tensor is required for quantization.\"\n weight = weight.to(dtype=torch.float32)\n scale = scale.to(dtype=torch.float32, device=weight.device)\n zero = zero.to(dtype=torch.float32, device=weight.device)\n oc, ic = weight.shape\n if subscale is not None: # per-group quantization\n subscale = subscale.to(dtype=weight.dtype, device=weight.device)\n # region reshape scale and zero point\n if scale.numel() == 1:\n scale = scale.view(-1).expand(oc)\n scale = scale.reshape(oc).contiguous().view(oc, 1)\n assert subscale.numel() > 1, \"subscale tensor is required for per-group quantization.\"\n subscale = subscale.view(oc, -1, 1).round_()\n ng = subscale.shape[1]\n gs = ic // ng\n assert ic == ng * gs, \"input channel size should be divisible by group size.\"\n if zero.numel() == 1:\n zero = zero.view(1, 1).expand(oc, ng)\n zero = zero.reshape(oc, ng).contiguous().view(oc, ng, 1).round_()\n # endregion\n # region quantize weight tensor\n weight = weight.div_(scale).round_()\n assert weight.min() >= -128, \"first-level quantized weight should be greater than or equal to -128.\"\n assert weight.max() <= 127, \"first-level quantized weight should be less than or equal to 127.\"\n weight = weight.view(oc, ng, gs)\n if not zero_pre_scaled: # zero point is int8\n weight = weight.add_(zero)\n weight = weight.div_(subscale)\n if zero_pre_scaled: # zero point is int4\n if zero.min() < 0: # sint4 zero point\n zero = zero.add_(8) # convert to uint4 zero point\n assert zero.min() >= 0, \"quantized zero point should be non-negative.\"\n assert zero.max() <= 15, \"quantized zero point should be less than 16.\"\n weight = weight.add_(zero)\n zero = zero.mul_(subscale)\n else:\n if weight.min() < 0: # sint4 weight\n weight = weight.add_(8) # convert to uint4 weight\n zero = zero.add_(8 * subscale)\n _weight = weight.mul(subscale)\n assert _weight.min() >= 0, \"first-level dequantize weight should be non-negative.\"\n assert _weight.max() <= 255, \"first-level dequantize weight should be less than 256.\"\n del _weight\n assert subscale.min() >= 0, \"subscale should be non-negative.\"\n assert subscale.max() <= 127, \"subscale should be less than or equal to 127.\"\n assert zero.min() >= 0, \"quantized zero point should be non-negative.\"\n assert zero.max() <= 255, \"quantized zero point should be less than 256.\"\n assert weight.min() >= 0, \"quantized weight should be non-negative.\"\n assert weight.max() <= 15, \"quantized weight should be less than 16.\"\n # endregion\n zero = -zero # ! for group quant, qserve uses q*s+z=r instead of q*s-z=r\n subscale = subscale.to(torch.int8)\n zero = zero.to(torch.int8)\n else: # per-channel quantization\n assert subscale is None, \"subscale tensor is not required for per-channel quantization.\"\n # region reshape scale and zero point\n if scale.numel() == 1:\n scale = scale.view(-1).expand(oc)\n scale = scale.reshape(oc).contiguous().view(oc, 1)\n if zero.numel() == 1:\n zero = zero.view(-1).expand(oc)\n zero = zero.reshape(oc).contiguous().view(oc, 1)\n # endregion\n # region quantize weight tensor\n if not zero_pre_scaled: # zero point is fp16\n weight = weight.add_(zero)\n weight = weight.div_(scale).round_()\n if zero_pre_scaled: # zero point is int4\n zero = zero.round_()\n if zero.min() < 0: # sint4 zero point\n zero = zero.add_(8) # convert to uint4 zero point\n assert zero.min() >= 0, \"quantized zero point should be non-negative.\"\n assert zero.max() <= 15, \"quantized zero point should be less than 16.\"\n weight = weight.add_(zero)\n zero = zero.mul_(scale)\n else:\n if weight.min() < 0: # sint4 weight\n weight = weight.add_(8) # convert to uint4 weight\n zero = zero.add_(8 * scale)\n assert weight.min() >= 0, \"quantized weight should be non-negative.\"\n assert weight.max() <= 15, \"quantized weight should be less than 16.\"\n # endregion\n zero = zero.to(dtype=dtype)\n scale = scale.to(dtype=dtype)\n packer = QServePacker()\n weight = packer.pack_weight(weight.view(oc, ic).to(torch.uint8))\n scale, zero, subscale = packer.pack_scale(scale=scale, zero=zero, subscale=subscale)\n return weight, scale, zero, subscale\n\n\ndef convert_to_qserve_w8x8y16_linear_weight(\n weight: torch.Tensor, scale: torch.Tensor\n) -> tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Convert a weight tensor to QServe W8-X8-Y16 linear weight format.\n\n Args:\n weight (`torch.Tensor`):\n weight tensor to be converted.\n scale (`torch.Tensor`):\n scale tensor for the weight tensor.\n\n Returns:\n `tuple[torch.Tensor, torch.Tensor]`:\n packed quantized weight tensor and scale tensor.\n \"\"\"\n dtype = weight.dtype\n assert dtype == torch.float16, \"currently qserve only supports fp16.\"\n assert scale is not None, \"scale tensor is required for quantization.\"\n weight = weight.to(dtype=torch.float32)\n scale = scale.to(dtype=torch.float32, device=weight.device)\n oc = weight.shape[0]\n if scale.numel() == 1:\n scale = scale.view(-1).expand(oc)\n scale = scale.reshape(oc).contiguous().view(oc, 1)\n weight = weight.div_(scale).round_()\n assert weight.min() >= -128, \"quantized weight should be greater than or equal to -128.\"\n assert weight.max() <= 127, \"quantized weight should be less than or equal to 127.\"\n weight = weight.contiguous().to(torch.int8)\n scale = scale.view(oc).to(dtype=dtype)\n return weight, scale\n"
},
{
"path": "deepcompressor/backend/tinychat/__init__.py",
"content": ""
},
{
"path": "deepcompressor/backend/tinychat/convert.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"QServe state dict converter module.\"\"\"\n\nimport argparse\nimport os\n\nimport safetensors.torch\nimport torch\nimport tqdm\n\nfrom .utils import convert_to_tinychat_w4x16y16_linear_weight\n\n\ndef convert_to_tinychat_w4x16y16_linear_state_dict(\n param_name: str,\n weight: torch.Tensor,\n scale: torch.Tensor,\n zero: torch.Tensor,\n zero_pre_scaled: bool = False,\n) -> dict[str, torch.Tensor]:\n \"\"\"Convert a weight tensor to TinyChat W4-X16-Y16 linear state dictionary.\n\n Args:\n param_name (`str`):\n parameter name.\n weight (`torch.Tensor`):\n weight tensor to be converted.\n scale (`torch.Tensor`):\n scale tensor for the weight tensor.\n zero (`torch.Tensor`):\n zero point tensor for the weight tensor.\n zero_pre_scaled (`bool`, *optional*, defaults to `False`):\n whether zero point tensor is pre-scaled.\n\n Returns:\n `dict[str, torch.Tensor]`:\n state dictionary for the quantized weight tensor.\n \"\"\"\n module_name = param_name[:-7]\n weight, scale, zero = convert_to_tinychat_w4x16y16_linear_weight(\n weight, scale=scale, zero=zero, zero_pre_scaled=zero_pre_scaled\n )\n state_dict: dict[str, torch.Tensor] = {}\n state_dict[f\"{module_name}.qweight\"] = weight.cpu()\n state_dict[f\"{module_name}.scales\"] = scale.cpu()\n state_dict[f\"{module_name}.scaled_zeros\"] = zero.cpu()\n return state_dict\n\n\ndef convert_to_tinychat_state_dict(\n state_dict: dict[str, torch.Tensor], scale_dict: dict[str, torch.Tensor]\n) -> dict[str, torch.Tensor]:\n scales: dict[str, dict[tuple[int, ...], torch.Tensor]] = {}\n zeros: dict[str, tuple[torch.Tensor | None, bool]] = {}\n print(\"Loading scale tensors...\")\n for name, tensor in tqdm.tqdm(scale_dict.items(), desc=\"Loading scale tensors\", leave=False, dynamic_ncols=True):\n print(f\" - Loading tensor {name} (dtype: {tensor.dtype}, shape: {tensor.shape}, device: {tensor.device})\")\n if name.endswith(\"zero\"):\n # this is a zero point tensor\n zero = None if tensor is None or all(t.item() == 0 for t in tensor.flatten()) else tensor\n if name.endswith(\".scaled_zero\"):\n zeros[name[:-12]] = (zero, False) # zero point tensor is post-scaled\n else:\n zeros[name[:-5]] = (zero, True) # zero point tensor is pre-scaled\n else:\n assert \".weight.scale\" in name\n # this is a scale tensor\n idx = name.index(\".weight.scale\")\n param_name = name[: idx + 7]\n scale_level = tuple(map(int, name[idx + 14 :].split(\".\")))\n scales.setdefault(param_name, {})[scale_level] = tensor\n for param_name in zeros.keys():\n assert param_name in state_dict, f\"zero point tensor {param_name} not found in state dict.\"\n assert param_name in scales, f\"scale tensor {param_name} not found in scale dict.\"\n converted: dict[str, torch.Tensor] = {}\n print(\"Converting state dict...\")\n for param_name, param in tqdm.tqdm(state_dict.items(), desc=\"Converting state dict\", dynamic_ncols=True):\n if param_name in scales:\n print(f\" - Converting {param_name} (dtype: {param.dtype}, shape: {param.shape}, device: {param.device})\")\n weight = param.data.clone()\n if param_name in zeros:\n zero, zero_pre_scaled = zeros[param_name]\n zero = zero.clone() if zero is not None else None\n else:\n zero, zero_pre_scaled = None, False\n level_scales = sorted(scales[param_name].items(), key=lambda x: x[0])\n assert len(level_scales) == 1, \"more than one scale levels are not supported.\"\n scale = level_scales[0][1].clone()\n converted.update(\n convert_to_tinychat_w4x16y16_linear_state_dict(\n param_name, weight, scale=scale, zero=zero, zero_pre_scaled=zero_pre_scaled\n )\n )\n else:\n if isinstance(param, torch.Tensor):\n print(f\" - Copying {param_name} (dtype: {param.dtype}, shape: {param.shape}, device: {param.device})\")\n converted[param_name] = param.clone().cpu()\n else:\n print(f\" - Copying {param_name} (type: {type(param)}, value: {param})\")\n converted[param_name] = param\n return converted\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--quant-path\", type=str, required=True, help=\"path to the quantization checkpoint directory.\")\n parser.add_argument(\"--output-root\", type=str, default=\"\", help=\"root to the output checkpoint directory.\")\n parser.add_argument(\"--model-name\", type=str, default=None, help=\"model name.\")\n parser.add_argument(\"--model-path\", type=str, default=None, help=\"path to the huggingface model directory.\")\n parser.add_argument(\"--copy-on-save\", action=\"store_true\", help=\"copy files on save.\")\n args = parser.parse_args()\n if not args.output_root:\n args.output_root = args.quant_path\n if args.model_name is None:\n assert args.model_path is not None, \"model name or path is required.\"\n model_name = args.model_path.rstrip(os.sep).split(os.sep)[-1]\n print(f\"Model name not provided. Using model name {model_name}.\")\n else:\n model_name = args.model_name\n state_dict = torch.load(\n os.path.join(args.quant_path, \"model.pt\"),\n map_location=\"cuda\" if torch.cuda.is_available() and torch.cuda.device_count() > 0 else \"cpu\",\n )\n scale_dict = torch.load(os.path.join(args.quant_path, \"scale.pt\"), map_location=\"cpu\")\n converted = convert_to_tinychat_state_dict(state_dict, scale_dict)\n model_name = f\"{args.model_name}-w4a16\"\n output_dirpath = os.path.join(args.output_root, model_name)\n\n os.makedirs(output_dirpath, exist_ok=True)\n if args.model_path and os.path.exists(args.model_path):\n output_path = os.path.join(output_dirpath, \"model.safetensors\")\n safetensors.torch.save_file(converted, output_path)\n print(f\"Quantized model checkpoint saved to {output_path}.\")\n for filename in os.listdir(args.model_path):\n if filename == \"tokenizer.model\" or (\n filename.endswith(\".json\") and filename != \"pytorch_model.bin.index.json\"\n ):\n filepath = os.path.abspath(os.path.join(args.model_path, filename))\n if args.copy_on_save:\n os.system(f\"cp {filepath} {output_dirpath}/\")\n else:\n os.system(f\"ln -s {filepath} {output_dirpath}/{filename}\")\n else:\n output_path = os.path.join(output_dirpath, \"tinychat-v2.pt\")\n torch.save(converted, output_path)\n print(f\"Quantized model checkpoint saved to {output_path}.\")\n print(f\"Quantized model saved to {output_dirpath}.\")\n"
},
{
"path": "deepcompressor/backend/tinychat/csrc/load.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"TinyChat Extension.\"\"\"\n\nimport os\n\nfrom torch.utils.cpp_extension import load\n\n__all__ = [\"_C\"]\n\ndirpath = os.path.dirname(__file__)\n\n_C = load(\n name=\"deepcompressor_tinychat_C\",\n sources=[\n f\"{dirpath}/pybind.cpp\",\n f\"{dirpath}/quantization/gemv/gemv_cuda.cu\",\n f\"{dirpath}/quantization/gemm/gemm_cuda.cu\",\n ],\n extra_cflags=[\"-g\", \"-O3\", \"-fopenmp\", \"-lgomp\", \"-std=c++20\"],\n extra_cuda_cflags=[\n \"-O3\",\n \"-std=c++20\",\n \"-U__CUDA_NO_HALF_OPERATORS__\",\n \"-U__CUDA_NO_HALF_CONVERSIONS__\",\n \"-U__CUDA_NO_HALF2_OPERATORS__\",\n \"-U__CUDA_NO_HALF2_CONVERSIONS__\",\n \"-U__CUDA_NO_BFLOAT16_OPERATORS__\",\n \"-U__CUDA_NO_BFLOAT16_CONVERSIONS__\",\n \"-U__CUDA_NO_BFLOAT162_OPERATORS__\",\n \"-U__CUDA_NO_BFLOAT162_CONVERSIONS__\",\n \"--expt-relaxed-constexpr\",\n \"--expt-extended-lambda\",\n \"--use_fast_math\",\n \"--ptxas-options=--allow-expensive-optimizations=true\",\n \"--threads=8\",\n ],\n)\n"
},
{
"path": "deepcompressor/backend/tinychat/csrc/pybind.cpp",
"content": "#include \n#include \n#include \"quantization/gemm/gemm_cuda.h\"\n#include \"quantization/gemv/gemv_cuda.h\"\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m)\n{\n m.def(\"awq_gemm_forward_cuda\", &awq_gemm_forward_cuda, \"AWQ quantized GEMM kernel.\");\n m.def(\"awq_gemv_forward_cuda\", &awq_gemv_forward_cuda, \"AWQ quantized GEMV kernel.\");\n}\n"
},
{
"path": "deepcompressor/backend/tinychat/csrc/quantization/dequantize.cuh",
"content": "/*\nModified from NVIDIA FasterTransformer: https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h\n\n@article{lin2023awq,\n title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration},\n author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Dang, Xingyu and Han, Song},\n journal={arXiv},\n year={2023}\n}\n*/\n#pragma once\n\n#include \n#include \n\ntemplate \n__device__ __forceinline__ void dequantize_s4_to_f16x2(T const &source, uint4 *result);\n\ntemplate <>\n__device__ __forceinline__ void dequantize_s4_to_f16x2(half2 const &source, uint4 *result)\n{\n uint32_t *h = reinterpret_cast(result);\n uint32_t const i4s = reinterpret_cast(source);\n\n // First, we extract the i4s and construct an intermediate fp16 number.\n static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;\n static constexpr uint32_t BOTTOM_MASK = 0x000f000f;\n static constexpr uint32_t TOP_MASK = 0x00f000f0;\n static constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400;\n\n // Note that the entire sequence only requires 1 shift instruction. This is thanks to the register packing\n // format and the fact that we force our integers to be unsigned, and account for this in the fp16 subtractions.\n // In addition, I exploit the fact that sub and fma have the same throughput in order to convert elt_23 and\n // elt_67 to fp16 without having to shift them to the bottom bits before hand.\n\n // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide RAW dependency if we issue\n // immediately before required.\n const uint32_t top_i4s = i4s >> 8;\n // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400\n asm volatile(\"lop3.b32 %0, %1, %2, %3, %4;\\n\"\n : \"=r\"(h[0])\n : \"r\"(i4s), \"n\"(BOTTOM_MASK), \"n\"(I4s_TO_F16s_MAGIC_NUM), \"n\"(immLut));\n // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400\n asm volatile(\"lop3.b32 %0, %1, %2, %3, %4;\\n\"\n : \"=r\"(h[1])\n : \"r\"(i4s), \"n\"(TOP_MASK), \"n\"(I4s_TO_F16s_MAGIC_NUM), \"n\"(immLut));\n // Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400\n asm volatile(\"lop3.b32 %0, %1, %2, %3, %4;\\n\"\n : \"=r\"(h[2])\n : \"r\"(top_i4s), \"n\"(BOTTOM_MASK), \"n\"(I4s_TO_F16s_MAGIC_NUM), \"n\"(immLut));\n // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400\n asm volatile(\"lop3.b32 %0, %1, %2, %3, %4;\\n\"\n : \"=r\"(h[3])\n : \"r\"(top_i4s), \"n\"(TOP_MASK), \"n\"(I4s_TO_F16s_MAGIC_NUM), \"n\"(immLut));\n\n // I use inline PTX below because I am not sure if the compiler will emit float2half instructions if I use the\n // half2 ctor. In this case, I chose performance reliability over code readability.\n\n // This is the half2 {1032, 1032} represented as an integer.\n // static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64086408;\n // Haotian: subtract {1024, 1024} instead, we do not need to map to [-8, 7]\n static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64006400;\n // This is the half2 {1 / 16, 1 / 16} represented as an integer.\n static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00;\n // This is the half2 {-72, -72} represented as an integer.\n // static constexpr uint32_t NEG_72 = 0xd480d480;\n // Haotian: Let's use {-64, -64}.\n static constexpr uint32_t NEG_64 = 0xd400d400;\n\n // Finally, we construct the output numbers.\n // Convert elt_01\n asm volatile(\"sub.f16x2 %0, %1, %2;\\n\" : \"=r\"(h[0]) : \"r\"(h[0]), \"r\"(FP16_TOP_MAGIC_NUM));\n // Convert elt_23\n asm volatile(\"fma.rn.f16x2 %0, %1, %2, %3;\\n\" : \"=r\"(h[1]) : \"r\"(h[1]), \"r\"(ONE_SIXTEENTH), \"r\"(NEG_64));\n // Convert elt_45\n asm volatile(\"sub.f16x2 %0, %1, %2;\\n\" : \"=r\"(h[2]) : \"r\"(h[2]), \"r\"(FP16_TOP_MAGIC_NUM));\n // Convert elt_67\n asm volatile(\"fma.rn.f16x2 %0, %1, %2, %3;\\n\" : \"=r\"(h[3]) : \"r\"(h[3]), \"r\"(ONE_SIXTEENTH), \"r\"(NEG_64));\n}\n\ntemplate <>\n__device__ __forceinline__ void dequantize_s4_to_f16x2<__nv_bfloat162>(__nv_bfloat162 const &source, uint4 *result)\n{\n uint32_t *h = reinterpret_cast(result);\n uint32_t const source_i4s = reinterpret_cast(source);\n\n // First, we extract the i4s and construct an intermediate bf16 number.\n static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;\n static constexpr uint32_t MASK = 0x000f000f;\n static constexpr uint32_t I4s_TO_BF16s_MAGIC_NUM = 0x43004300;\n\n uint32_t i4s = source_i4s;\n // Extract elt_01 - (i4s & 0x000f000f) | 0x43004300\n asm volatile(\"lop3.b32 %0, %1, %2, %3, %4;\\n\"\n : \"=r\"(h[0])\n : \"r\"(i4s), \"n\"(MASK), \"n\"(I4s_TO_BF16s_MAGIC_NUM), \"n\"(immLut));\n // Extract elt_23 (i4s & 0x00f000f0) | 0x43004300\n i4s >>= 4;\n asm volatile(\"lop3.b32 %0, %1, %2, %3, %4;\\n\"\n : \"=r\"(h[1])\n : \"r\"(i4s), \"n\"(MASK), \"n\"(I4s_TO_BF16s_MAGIC_NUM), \"n\"(immLut));\n // Extract elt_45 (top_i4s & 0x000f000f) | 0x43004300\n i4s >>= 4;\n asm volatile(\"lop3.b32 %0, %1, %2, %3, %4;\\n\"\n : \"=r\"(h[2])\n : \"r\"(i4s), \"n\"(MASK), \"n\"(I4s_TO_BF16s_MAGIC_NUM), \"n\"(immLut));\n // Extract elt_67 (top_i4s & 0x00f000f0) | 0x43004300\n i4s >>= 4;\n asm volatile(\"lop3.b32 %0, %1, %2, %3, %4;\\n\"\n : \"=r\"(h[3])\n : \"r\"(i4s), \"n\"(MASK), \"n\"(I4s_TO_BF16s_MAGIC_NUM), \"n\"(immLut));\n\n // This is the BF16 {-136, -136} represented as an integer.\n // static constexpr uint32_t BF16_BIAS = 0xC308C308;\n // This is the BF16 {-128, -128} represented as an integer, we do not need to map to [-8, 7]\n static constexpr uint32_t NEG_128 = 0xC300C300;\n static constexpr uint32_t ONE = 0x3F803F80;\n\n // Finally, we construct the output numbers.\n // Convert elt_01\n asm volatile(\"fma.rn.bf16x2 %0, %1, %2, %3;\\n\" : \"=r\"(h[0]) : \"r\"(h[0]), \"r\"(ONE), \"r\"(NEG_128));\n // Convert elt_23\n asm volatile(\"fma.rn.bf16x2 %0, %1, %2, %3;\\n\" : \"=r\"(h[1]) : \"r\"(h[1]), \"r\"(ONE), \"r\"(NEG_128));\n // Convert elt_45\n asm volatile(\"fma.rn.bf16x2 %0, %1, %2, %3;\\n\" : \"=r\"(h[2]) : \"r\"(h[2]), \"r\"(ONE), \"r\"(NEG_128));\n // Convert elt_67\n asm volatile(\"fma.rn.bf16x2 %0, %1, %2, %3;\\n\" : \"=r\"(h[3]) : \"r\"(h[3]), \"r\"(ONE), \"r\"(NEG_128));\n}"
},
{
"path": "deepcompressor/backend/tinychat/csrc/quantization/gemm/gemm_cuda.cu",
"content": "#include \n#include \n#include \"semaphore.h\"\n#include \"gemm_cuda.h\"\n#include \"../dequantize.cuh\"\n#include \"../../utils.cuh\"\n#include \n#include \n\n#define kInterleave 4\n#define OP_M 16\n#define OP_N 8\n#define OP_K 16\n#define INTRIN_M 16\n#define INTRIN_N 16\n#define INTRIN_K 16\n#define WARP_SIZE 32\n#define SMEM_PAD_A 0\n#define SMEM_PAD_B 0\n#define PACK_SIZE 8\n#if (__CUDACC_VER_MAJOR__ >= 11) && (__CUDACC_VER_MINOR__ >= 4)\n#define L2_CACHEHINT(size) \".L2::\" #size \"B\"\n#else\n#define L2_CACHEHINT(size)\n#endif\n\n#define KERNEL_LAUNCH_CODE \\\n int num_mn_tiles = (num_in_feats + CTA_M - 1) / CTA_M * (num_out_channels + CTA_N - 1) / CTA_N; \\\n torch::Tensor _semaphores = torch::empty({num_mn_tiles}, options_int); \\\n auto semaphores = reinterpret_cast(_semaphores.data_ptr()); \\\n constexpr int NUM_WARPS = (CTA_M / WARP_M) * (CTA_N / WARP_N) * (CTA_K / WARP_K); \\\n constexpr int SCALES_SMEM_SIZE = (G >= CTA_K) ? (CTA_N / (G / CTA_K) * STAGES * 2) : (CTA_N * (CTA_K / G) * STAGES * 2); \\\n constexpr int kSmemByteSize = (CTA_M * (CTA_K + SMEM_PAD_A) + CTA_N * (CTA_K + SMEM_PAD_B) / kInterleave + SCALES_SMEM_SIZE) * STAGES * sizeof(f16_t); \\\n if (kSmemByteSize >= 99 * 1024) \\\n { \\\n printf(\"This kernel requires %d Bytes of shared memory, which exceeds device limit.\\n\", kSmemByteSize); \\\n return _out_feats; \\\n } \\\n int j_factors1 = num_out_channels / CTA_N / 1; \\\n dim3 num_blocks((num_out_feats + CTA_M - 1) / CTA_M * j_factors1 * SPLITK); \\\n dim3 threads_per_block(WARP_SIZE, NUM_WARPS); \\\n auto kernel_func = gemm_w4a16_T1; \\\n cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize); \\\n kernel_func<<>>( \\\n in_feats, kernel, scales, zeros, out_feats, semaphores, num_in_feats, num_out_channels, num_in_channels);\n\ntemplate \n__inline__ __host__ __device__ int get_log_tile(int n)\n{\n if (N >= 8 && n >= 6)\n return 3;\n else if (N >= 4 && n >= 3)\n return 2;\n else if (N >= 2 && n >= 2)\n return 1;\n else\n return 0;\n}\n\n__inline__ __device__ uint2 get_block_idx_mapping(int blockIdx_x, int blockIdx_y, int log_tile)\n{\n return make_uint2((blockIdx_x >> log_tile), (blockIdx_y << log_tile) + ((blockIdx_x) & ((1 << (log_tile)) - 1)));\n}\n\ntemplate \n__device__ void sync_slice(int slice_id)\n{\n if constexpr (SLICES == 1)\n {\n __syncthreads();\n }\n else\n {\n constexpr int SLICE_GROUP = (SLICES + 7) / 8;\n constexpr uint32_t num_threads = NUM_WARPS_MN * WARP_SIZE;\n const uint32_t barrier_id = slice_id / SLICE_GROUP + 1;\n asm volatile(\"bar.sync %0, %1;\" : : \"r\"(barrier_id), \"n\"(num_threads));\n }\n}\n\n__inline__ __device__ uint32_t cast_smem_ptr_to_uint(void const *const ptr)\n{\n uint32_t smem_int_ptr;\n\n asm(\"{.reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }\\n\"\n : \"=r\"(smem_int_ptr)\n : \"l\"(ptr));\n\n return smem_int_ptr;\n}\n\ntemplate \n__inline__ __device__ void ldmatrix_m8n8_x4_b16(f16_t *shared_warp, int ax0_0, uint32_t addr)\n{\n static_assert(std::is_same::value || std::is_same::value,\n \"ldmatrix_m8n8_x4_b16 supports only half or __nv_bfloat16 types.\");\n __asm__ __volatile__(\n \"ldmatrix.sync.aligned.m8n8.x4.shared.b16\"\n \"{%0, %1, %2, %3}, [%4];\"\n : \"=r\"(((unsigned *)(shared_warp + (ax0_0 * 8)))[0]), \"=r\"(((unsigned *)(shared_warp + (ax0_0 * 8)))[1]), \"=r\"(((unsigned *)(shared_warp + (ax0_0 * 8)))[2]), \"=r\"(((unsigned *)(shared_warp + (ax0_0 * 8)))[3])\n : \"r\"(addr));\n}\n\ntemplate \n__inline__ __device__ void ldmatrix_m8n8_x4_trans_b16(f16_t *shared_warp, int ax0_0, uint32_t addr)\n{\n static_assert(std::is_same::value || std::is_same::value,\n \"ldmatrix_m8n8_x4_trans_b16 supports only half or __nv_bfloat16 types.\");\n __asm__ __volatile__(\n \"ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16\"\n \"{%0, %1, %2, %3}, [%4];\"\n : \"=r\"(((unsigned *)(shared_warp + (ax0_0 * 8)))[0]), \"=r\"(((unsigned *)(shared_warp + (ax0_0 * 8)))[1]), \"=r\"(((unsigned *)(shared_warp + (ax0_0 * 8)))[2]), \"=r\"(((unsigned *)(shared_warp + (ax0_0 * 8)))[3])\n : \"r\"(addr));\n}\n\n__inline__ __device__ void cp_async_cg_A(uint32_t smem_int_ptr, const uint4 *__restrict__ src, bool mask)\n{\n const int cp_size = 16;\n asm volatile(\"{\"\n \" .reg .pred p;\"\n \" setp.ne.b32 p, %0, 0;\"\n \" @p cp.async.cg.shared.global\" L2_CACHEHINT(128) \" [%1], [%2], %3;\"\n \"}\" ::\"r\"((int)mask),\n \"r\"(smem_int_ptr),\n \"l\"(src),\n \"n\"(cp_size));\n}\n\ntemplate \n__device__ __inline__ void mma_m16n8k16(float *C_warp, f16_t *A_shared_warp, f16_t *B_shared_warp);\n\ntemplate <>\n__device__ __inline__ void mma_m16n8k16(float *C_warp, half *A_shared_warp, half *B_shared_warp)\n{\n __asm__ __volatile__(\n \"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32\"\n \"{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};\"\n : \"=f\"(((float *)C_warp)[0]), \"=f\"(((float *)C_warp)[1]), \"=f\"(((float *)C_warp)[2]), \"=f\"(((float *)C_warp)[3])\n : \"r\"(((unsigned *)A_shared_warp)[0]), \"r\"(((unsigned *)A_shared_warp)[1]), \"r\"(((unsigned *)A_shared_warp)[2]), \"r\"(((unsigned *)A_shared_warp)[3]), \"r\"(((unsigned *)B_shared_warp)[0]), \"r\"(((unsigned *)B_shared_warp)[1]), \"f\"(((float *)C_warp)[0]), \"f\"(((float *)C_warp)[1]), \"f\"(((float *)C_warp)[2]), \"f\"(((float *)C_warp)[3]));\n}\n\ntemplate <>\n__device__ __inline__ void mma_m16n8k16<__nv_bfloat16>(float *C_warp, __nv_bfloat16 *A_shared_warp, __nv_bfloat16 *B_shared_warp)\n{\n __asm__ __volatile__(\n \"mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32\"\n \"{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};\"\n : \"=f\"(((float *)C_warp)[0]), \"=f\"(((float *)C_warp)[1]), \"=f\"(((float *)C_warp)[2]), \"=f\"(((float *)C_warp)[3])\n : \"r\"(((unsigned *)A_shared_warp)[0]), \"r\"(((unsigned *)A_shared_warp)[1]), \"r\"(((unsigned *)A_shared_warp)[2]), \"r\"(((unsigned *)A_shared_warp)[3]), \"r\"(((unsigned *)B_shared_warp)[0]), \"r\"(((unsigned *)B_shared_warp)[1]), \"f\"(((float *)C_warp)[0]), \"f\"(((float *)C_warp)[1]), \"f\"(((float *)C_warp)[2]), \"f\"(((float *)C_warp)[3]));\n}\n\ntemplate \n__device__ __inline__ void global_to_share_one_stage_A(f16_t *src, f16_t *dst, int global_nrows, int global_ncols, int cta_offset_m, int cta_offset_n, int cta_offset_k, int global_iter_k, int shared_iter_k, bool mask)\n{\n constexpr int threads_needed = (CTA_M * CTA_K) / PACK_SIZE / SHARED_K_ITERS;\n constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE;\n constexpr int total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / threads_used;\n constexpr int partial_global_iters = (total_global_iters + SHARED_K_ITERS - 1) / SHARED_K_ITERS;\n constexpr int cta_step_m_or_n = (threads_used * PACK_SIZE) / CTA_K;\n constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K;\n constexpr int threads_per_row = CTA_K / PACK_SIZE;\n constexpr int kSmemCol = CTA_K + SMEM_PAD_A;\n bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used);\n int ld_col = (threadIdx.x % threads_per_row);\n#pragma unroll\n for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter)\n {\n int global_iter = shared_iter_k * partial_global_iters + _global_iter;\n int ld_row = global_iter * cta_step_m_or_n + threadIdx.y * warp_step_m_or_n + (threadIdx.x / threads_per_row);\n int ld_col_swizzled = (ld_col ^ (ld_row) & 7) * PACK_SIZE;\n void *dst_ptr = (void *)(dst + ld_row * kSmemCol + ld_col_swizzled);\n uint4 *src_ptr = (uint4 *)(src + (ld_row + cta_offset_m) * global_ncols + ld_col * PACK_SIZE + global_iter_k * CTA_K + cta_offset_k); // cta_offset_m * global_ncols + global_iter * cta_step_m_or_n * global_ncols + threadIdx.y * warp_step_m_or_n * global_ncols + (threadIdx.x / threads_per_row) * global_ncols + global_iter_k * CTA_K + (threadIdx.x % threads_per_row) * PACK_SIZE);\n if constexpr (STAGES > 1)\n {\n uint32_t addr = cast_smem_ptr_to_uint(dst_ptr);\n cp_async_cg_A(addr, src_ptr, local_mask & (ld_row + cta_offset_m < global_nrows));\n }\n else\n {\n if (local_mask & (ld_row + cta_offset_m < global_nrows))\n *(uint4 *)dst_ptr = *src_ptr;\n }\n }\n}\n\ntemplate \n__device__ __inline__ void global_to_share_one_stage_B(f16_t *src, f16_t *dst, int global_ncols, int cta_offset_m, int cta_offset_n, int cta_offset_k, int global_iter_k, int shared_iter_k, bool mask)\n{\n constexpr int threads_needed = (CTA_N / kInterleave * CTA_K) / PACK_SIZE / SHARED_K_ITERS;\n constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE;\n constexpr int total_global_iters = (CTA_N / kInterleave * CTA_K) / PACK_SIZE / threads_used;\n constexpr int partial_global_iters = (total_global_iters + SHARED_K_ITERS - 1) / SHARED_K_ITERS;\n constexpr int cta_step_m_or_n = (threads_used * PACK_SIZE) / CTA_K;\n constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K;\n constexpr int threads_per_row = CTA_K / PACK_SIZE;\n constexpr int kSmemCol = CTA_K + SMEM_PAD_B;\n bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used);\n#pragma unroll\n for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter)\n {\n int global_iter = shared_iter_k * partial_global_iters + _global_iter;\n\n int ld_row = global_iter * cta_step_m_or_n + threadIdx.y * warp_step_m_or_n + (threadIdx.x / threads_per_row);\n int ld_col = (threadIdx.x % threads_per_row);\n int ld_col_swizzled = ld_col ^ (ld_row % 2) & 7;\n void *dst_ptr = (void *)(dst + (ld_row * kSmemCol + ld_col_swizzled * PACK_SIZE));\n uint4 *src_ptr = (uint4 *)(src + global_iter_k * CTA_K + cta_offset_n / kInterleave * global_ncols + ld_row * global_ncols + ld_col * PACK_SIZE + cta_offset_k);\n if constexpr (STAGES > 1)\n {\n uint32_t addr = cast_smem_ptr_to_uint(dst_ptr);\n cp_async_cg_A(addr, src_ptr, local_mask);\n }\n else\n {\n if (local_mask)\n *(uint4 *)dst_ptr = *src_ptr;\n }\n }\n}\n\ntemplate \n__device__ __inline__ void global_to_share_one_stage_scales(f16_t *src, f16_t *dst, f16_t *src_z, f16_t *dst_z, int global_ncols, int cta_offset_m, int cta_offset_n, int cta_offset_k, int global_iter_k, int shared_iter_k, bool mask)\n{\n constexpr int LD_AMOUNT = (G >= CTA_K) ? CTA_N : CTA_N * CTA_K / G;\n constexpr int threads_needed = LD_AMOUNT / PACK_SIZE / 1;\n constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE;\n constexpr int total_global_iters = LD_AMOUNT / PACK_SIZE / threads_used;\n constexpr int threads_per_row = CTA_N / PACK_SIZE;\n constexpr int kSmemCol = CTA_N;\n bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used);\n int g_idx = (cta_offset_k + global_iter_k * CTA_K) / G;\n\n void *dst_ptr = (void *)(dst + (threadIdx.x / threads_per_row) * kSmemCol + (threadIdx.x % threads_per_row) * PACK_SIZE);\n uint4 *src_ptr = (uint4 *)(src + g_idx * global_ncols + cta_offset_n + (threadIdx.x / threads_per_row) * global_ncols + (threadIdx.x % threads_per_row) * PACK_SIZE);\n void *dst_ptr_z = (void *)(dst_z + (threadIdx.x / threads_per_row) * kSmemCol + (threadIdx.x % threads_per_row) * PACK_SIZE);\n uint4 *src_ptr_z = (uint4 *)(src_z + g_idx * global_ncols + cta_offset_n + (threadIdx.x / threads_per_row) * global_ncols + (threadIdx.x % threads_per_row) * PACK_SIZE);\n if (STAGES > 1)\n {\n uint32_t addr = cast_smem_ptr_to_uint(dst_ptr);\n cp_async_cg_A(addr, src_ptr, local_mask);\n uint32_t addr_z = cast_smem_ptr_to_uint(dst_ptr_z);\n cp_async_cg_A(addr_z, src_ptr_z, local_mask);\n }\n else\n {\n if (local_mask)\n {\n *(uint4 *)dst_ptr = *src_ptr;\n *(uint4 *)dst_ptr_z = *src_ptr_z;\n }\n }\n}\n\ntemplate \n__device__ __inline__ void share_to_reg_one_stage_A(f16_t *src, f16_t *dst, int warp_offset_m, int warp_offset_n, int warp_offset_k, int k_0_1)\n{\n constexpr int kSmemCol = CTA_K + SMEM_PAD_A;\n\n for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter)\n {\n\n int ld_row = warp_offset_m + shared_iter * OP_M + (threadIdx.x % 16);\n int ld_col = k_0_1 * 16 + (threadIdx.x / 16) * 8 + warp_offset_k;\n int ld_col_swizzled = ((ld_col / PACK_SIZE) ^ (ld_row) & 7) * PACK_SIZE;\n void *addr_ptr = (void *)(src + ld_row * kSmemCol + ld_col_swizzled);\n\n uint32_t addr = cast_smem_ptr_to_uint(addr_ptr);\n ldmatrix_m8n8_x4_b16(dst, shared_iter, addr);\n }\n}\n\ntemplate \n__device__ __inline__ void share_to_reg_one_stage_B(f16_t *src, f16_t *src_scales, f16_t *src_zeros, f16_t *dst, f16_t *dst_fp16, int warp_offset_m, int warp_offset_n, int warp_offset_k, int k_0_1)\n{\n using f162_t = typename packed_as::type;\n\n constexpr int kSmemCol = CTA_K + SMEM_PAD_B;\n int r0 = ((threadIdx.x / 8 / 2) * 8 + threadIdx.x % 8);\n int c0 = ((threadIdx.x / 8) % 2) * 8;\n int r = r0 / 4;\n int c = (r0 % 4) * 16 + c0;\n int c_swizzled = ((c / PACK_SIZE) ^ (r % 2) & 7) * PACK_SIZE;\n\n if constexpr (ldmatrix)\n {\n#pragma unroll\n for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter)\n {\n void *addr_ptr = (void *)(src + warp_offset_n / kInterleave * kSmemCol + shared_iter * 16 / kInterleave * kSmemCol + k_0_1 * 16 + r * kSmemCol + c_swizzled + warp_offset_k);\n uint32_t addr = cast_smem_ptr_to_uint(addr_ptr);\n ldmatrix_m8n8_x4_b16(dst, shared_iter, addr);\n }\n }\n\n#pragma unroll\n for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter)\n {\n f16_t scale = src_scales[(warp_offset_k / G) * CTA_N + warp_offset_n + 16 * shared_iter + 8 * (k_0_1 % 2) + threadIdx.x / 4];\n f16_t zero = src_zeros[(warp_offset_k / G) * CTA_N + warp_offset_n + 16 * shared_iter + 8 * (k_0_1 % 2) + threadIdx.x / 4];\n f162_t scale2 = f162f162(scale);\n f162_t zero2 = f162f162(zero);\n f162_t loaded[4];\n\n dequantize_s4_to_f16x2(*reinterpret_cast(dst + (k_0_1 % 2) * 4 + (k_0_1 / 2 * 2) + shared_iter * 8), reinterpret_cast(loaded));\n#pragma unroll\n for (int i = 0; i < 4; i++)\n {\n loaded[i] = __hfma2(loaded[i], scale2, zero2);\n }\n *reinterpret_cast(dst_fp16 + shared_iter * 16 + 8 * (k_0_1 % 2)) = *reinterpret_cast(loaded);\n }\n}\n\ntemplate \n__global__ void gemm_w4a16_T1(f16_t *__restrict__ A, f16_t *__restrict__ B, f16_t *__restrict__ scales, f16_t *__restrict__ zeros, f16_t *__restrict__ C, int *__restrict__ semaphores, int M, int N, int K)\n{\n using f162_t = typename packed_as::type;\n\n constexpr int NUM_WARPS_MN = CTA_M / WARP_M * CTA_N / WARP_N;\n constexpr int NUM_WARPS = NUM_WARPS_MN * CTA_K / WARP_K;\n constexpr int CTA_SIZE = NUM_WARPS * WARP_SIZE;\n constexpr int CTA_SIZE_MN = NUM_WARPS_MN * WARP_SIZE;\n constexpr int SLICES = CTA_K / WARP_K;\n int num_blocks_n = (N + CTA_N - 1) / CTA_N;\n int num_blocks_m = (M + CTA_M - 1) / CTA_M;\n int blockIdx_x = 0;\n int blockIdx_y = blockIdx.x % (num_blocks_m * num_blocks_n);\n int blockIdx_z = blockIdx.x / (num_blocks_m * num_blocks_n);\n const int log_tile = get_log_tile<1>((N + CTA_N - 1) / CTA_N);\n int blockIdx_m = blockIdx_y / (num_blocks_n >> log_tile);\n int blockIdx_n = blockIdx_y % (num_blocks_n >> log_tile);\n const uint2 block_idx_mapping = get_block_idx_mapping(blockIdx_m, blockIdx_n, log_tile);\n blockIdx_m = block_idx_mapping.x;\n blockIdx_n = block_idx_mapping.y;\n\n float C_warp[CTA_M * CTA_N / CTA_SIZE_MN];\n constexpr int kSmemPadKA = CTA_K + SMEM_PAD_A;\n constexpr int kSmemPadKB = CTA_K + SMEM_PAD_B;\n constexpr int kSmemSizeAPerStage = CTA_M * kSmemPadKA;\n constexpr int kSmemSizeBPerStage = CTA_N / kInterleave * kSmemPadKB;\n constexpr int kSmemSizeA = kSmemSizeAPerStage * STAGES;\n constexpr int kSmemSizeB = kSmemSizeBPerStage * STAGES;\n constexpr int scales_load_interval = G >= CTA_K ? G / CTA_K : 1;\n constexpr int scales_per_load = G < CTA_K ? CTA_K / G : 1;\n constexpr int kSmemSizeScales = CTA_N * STAGES / scales_load_interval * scales_per_load;\n constexpr int kSmemSizeZeros = CTA_N * STAGES / scales_load_interval * scales_per_load;\n extern __shared__ half mem_shared[];\n f16_t *A_shared = reinterpret_cast(mem_shared);\n f16_t *B_shared = reinterpret_cast(mem_shared + kSmemSizeA);\n f16_t *scales_shared = reinterpret_cast(mem_shared + kSmemSizeA + kSmemSizeB);\n f16_t *zeros_shared = reinterpret_cast(mem_shared + kSmemSizeA + kSmemSizeB + kSmemSizeScales);\n float *C_shared = reinterpret_cast(mem_shared);\n f16_t A_shared_warp_[2][WARP_M * INTRIN_K / WARP_SIZE];\n f16_t B_shared_warp_[2][WARP_N * 32 / WARP_SIZE];\n f16_t B_shared_warp_tmp_[2][WARP_N * 16 / WARP_SIZE];\n int cta_offset_m = blockIdx_m * CTA_M;\n int cta_offset_n = blockIdx_n * CTA_N;\n int cta_offset_k = blockIdx_z * (K / SPLITK);\n int warp_mn = threadIdx.y % NUM_WARPS_MN;\n int slice_id = threadIdx.y / NUM_WARPS_MN;\n int warp_offset_n = (warp_mn % (CTA_N / WARP_N)) * WARP_N;\n int warp_offset_m = (warp_mn / (CTA_N / WARP_N)) * WARP_M;\n int warp_offset_k = slice_id * WARP_K;\n\n for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE_MN; i++)\n C_warp[i] = 0.0;\n\n int gemm_iters = (K + CTA_K - 1) / CTA_K / SPLITK;\n int k_0_0_ld = 0;\n int k_0_0 = 0;\n constexpr int prologue_stages = STAGES == 1 ? 1 : STAGES - 1;\n#pragma unroll\n for (k_0_0_ld = 0; k_0_0_ld < prologue_stages; ++k_0_0_ld)\n {\n global_to_share_one_stage_A(A, A_shared + k_0_0_ld * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, 0, true);\n global_to_share_one_stage_B(B, B_shared + k_0_0_ld * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, 0, true);\n global_to_share_one_stage_scales(\n scales, scales_shared + (k_0_0_ld / scales_load_interval * scales_per_load) * CTA_N,\n zeros, zeros_shared + (k_0_0_ld / scales_load_interval * scales_per_load) * CTA_N,\n N, cta_offset_m, cta_offset_n, cta_offset_k,\n k_0_0_ld, 0, k_0_0_ld < gemm_iters && k_0_0_ld % scales_load_interval == 0);\n if constexpr (STAGES > 1)\n __pipeline_commit();\n }\n if constexpr (STAGES > 1)\n __pipeline_wait_prior(STAGES - 2);\n __syncthreads();\n\n share_to_reg_one_stage_A(A_shared, A_shared_warp_[0], warp_offset_m, warp_offset_n, warp_offset_k, 0);\n share_to_reg_one_stage_B(B_shared, scales_shared, zeros_shared, B_shared_warp_tmp_[0], B_shared_warp_[0], warp_offset_m, warp_offset_n, warp_offset_k, 0);\n constexpr int SHARED_K_ITERS = WARP_K / INTRIN_K;\n\n for (; k_0_0 < gemm_iters; ++k_0_0, ++k_0_0_ld)\n {\n int ld_stage = k_0_0_ld % STAGES;\n int compute_stage = k_0_0 % STAGES;\n f16_t *A_shared_this_compute_stage;\n f16_t *B_shared_this_compute_stage;\n f16_t *scales_shared_this_compute_stage;\n f16_t *zeros_shared_this_compute_stage;\n\n#pragma unroll\n for (int iter_k = 0; iter_k < SHARED_K_ITERS; ++iter_k)\n {\n A_shared_this_compute_stage = A_shared + compute_stage * kSmemSizeAPerStage;\n B_shared_this_compute_stage = B_shared + compute_stage * kSmemSizeBPerStage;\n scales_shared_this_compute_stage = scales_shared + (compute_stage / scales_load_interval * scales_per_load) * CTA_N;\n zeros_shared_this_compute_stage = zeros_shared + (compute_stage / scales_load_interval * scales_per_load) * CTA_N;\n share_to_reg_one_stage_A(A_shared_this_compute_stage, A_shared_warp_[(iter_k + 1) % 2], warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS);\n if ((iter_k + 1) % kInterleave == 0)\n {\n if (compute_stage % 2 == 1)\n {\n share_to_reg_one_stage_B(\n B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage,\n B_shared_warp_tmp_[1], B_shared_warp_[((iter_k + 1) / 2) % 2],\n warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS);\n }\n else\n {\n share_to_reg_one_stage_B(\n B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage,\n B_shared_warp_tmp_[0], B_shared_warp_[((iter_k + 1) / 2) % 2],\n warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS);\n }\n }\n else\n {\n if (compute_stage % 2 == 1)\n {\n share_to_reg_one_stage_B(\n B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage,\n B_shared_warp_tmp_[1], B_shared_warp_[((iter_k + 1) / 2) % 2],\n warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS);\n }\n else\n {\n share_to_reg_one_stage_B(\n B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage,\n B_shared_warp_tmp_[0], B_shared_warp_[((iter_k + 1) / 2) % 2],\n warp_offset_m, warp_offset_n, warp_offset_k, (iter_k + 1) % SHARED_K_ITERS);\n }\n }\n f16_t *A_shared_warp = A_shared_warp_[iter_k % 2];\n f16_t *B_shared_warp = B_shared_warp_[(iter_k / 2) % 2];\n\n for (int i_0_3 = 0; i_0_3 < WARP_M / INTRIN_M; ++i_0_3)\n {\n for (int j_0_4 = 0; j_0_4 < WARP_N / INTRIN_N; ++j_0_4)\n {\n mma_m16n8k16(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4);\n mma_m16n8k16(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4 + 8);\n }\n }\n\n if (iter_k < WARP_K / INTRIN_K - 1)\n {\n if constexpr (STAGES == 1)\n __syncthreads();\n global_to_share_one_stage_A(A, A_shared + ld_stage * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters);\n global_to_share_one_stage_B(B, B_shared + ld_stage * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters);\n }\n\n if (iter_k == WARP_K / INTRIN_K - 2)\n {\n if constexpr (STAGES == 1 && WARP_K / INTRIN_K > 2)\n {\n __syncthreads();\n }\n global_to_share_one_stage_A(A, A_shared + ld_stage * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, iter_k + 1, k_0_0_ld < gemm_iters);\n global_to_share_one_stage_B(B, B_shared + ld_stage * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, cta_offset_k, k_0_0_ld, iter_k + 1, k_0_0_ld < gemm_iters);\n global_to_share_one_stage_scales(\n scales, scales_shared + (ld_stage / scales_load_interval * scales_per_load) * CTA_N,\n zeros, zeros_shared + (ld_stage / scales_load_interval * scales_per_load) * CTA_N,\n N, cta_offset_m, cta_offset_n, cta_offset_k,\n k_0_0_ld, iter_k, k_0_0_ld < gemm_iters && k_0_0_ld % scales_load_interval == 0);\n if constexpr (STAGES > 1)\n {\n __pipeline_commit();\n __pipeline_wait_prior(STAGES - 2);\n }\n compute_stage = (k_0_0 + 1) % STAGES;\n __syncthreads();\n }\n }\n }\n __pipeline_commit();\n __pipeline_wait_prior(0);\n __syncthreads();\n if constexpr (SLICES > 1)\n {\n#pragma unroll\n for (int z = 0; z < SLICES; ++z)\n {\n if (slice_id == z)\n {\n#pragma unroll\n for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1)\n {\n#pragma unroll\n for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1)\n {\n#pragma unroll\n for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id)\n {\n if (z > 0)\n {\n C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] += C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2];\n }\n C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2] = C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id];\n };\n }\n }\n }\n __syncthreads();\n }\n if (slice_id == 0)\n {\n#pragma unroll\n for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1)\n {\n#pragma unroll\n for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1)\n {\n#pragma unroll\n for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id)\n {\n C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] = C_shared[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2];\n };\n }\n }\n }\n }\n\n if (slice_id == 0)\n {\n Semaphore semaphore(semaphores + blockIdx_y, threadIdx.x);\n\n if constexpr (SPLITK > 1)\n {\n semaphore.fetch();\n }\n\n if (blockIdx_z != 0)\n {\n semaphore.wait(blockIdx_z);\n for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1)\n {\n for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1)\n {\n for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2)\n {\n int write_row = cta_offset_m + warp_offset_m + ax0_0_1 * OP_M + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4));\n\n if (write_row < M)\n {\n f162_t *existing_psum_ptr = reinterpret_cast(\n C + write_row * N +\n cta_offset_n + warp_offset_n + ax1_0_1 * 16 +\n (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2);\n\n *existing_psum_ptr = __hadd2(\n *existing_psum_ptr,\n cuda_cast(*reinterpret_cast(\n C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id)));\n }\n };\n }\n }\n }\n else\n {\n for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1)\n {\n for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1)\n {\n for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2)\n {\n int write_row = cta_offset_m + warp_offset_m + ax0_0_1 * OP_M + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4));\n if (write_row < M)\n {\n *reinterpret_cast(\n C + write_row * N +\n cta_offset_n + warp_offset_n + ax1_0_1 * 16 +\n (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2) =\n cuda_cast(*reinterpret_cast(C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 +\n ax1_0_1 * 8 + local_id));\n }\n };\n }\n }\n }\n\n if constexpr (SPLITK > 1)\n {\n\n int lock = 0;\n if (SPLITK == blockIdx_z + 1)\n {\n\n lock = 0;\n }\n else\n {\n lock = blockIdx_z + 1;\n }\n semaphore.release(lock);\n }\n }\n}\n\ntemplate \n__device__ __inline__ void global_to_share_one_stage_A_T2(f16_t *src, f16_t *dst, int global_nrows, int global_ncols, int cta_offset_m, int cta_offset_n, int global_iter_k, int shared_iter_k, bool mask)\n{\n constexpr int threads_needed = (CTA_M * CTA_K) / PACK_SIZE / SHARED_K_ITERS;\n constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE;\n constexpr int total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / threads_used;\n constexpr int partial_global_iters = (total_global_iters + SHARED_K_ITERS - 1) / SHARED_K_ITERS;\n constexpr int cta_step_m_or_n = (threads_used * PACK_SIZE) / CTA_K;\n constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K;\n constexpr int threads_per_row = CTA_K / PACK_SIZE;\n constexpr int kSmemCol = CTA_K + SMEM_PAD_A;\n bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used);\n int ld_col = (threadIdx.x % threads_per_row);\n#pragma unroll\n for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter)\n {\n int global_iter = shared_iter_k * partial_global_iters + _global_iter;\n int ld_row = global_iter * cta_step_m_or_n + threadIdx.y * warp_step_m_or_n + (threadIdx.x / threads_per_row);\n int ld_col_swizzled = (ld_col ^ (ld_row) & 7) * PACK_SIZE;\n void *dst_ptr = (void *)(dst + ld_row * kSmemCol + ld_col_swizzled);\n uint4 *src_ptr = (uint4 *)(src + (ld_row + cta_offset_m) * global_ncols + ld_col * PACK_SIZE + global_iter_k * CTA_K); // cta_offset_m * global_ncols + global_iter * cta_step_m_or_n * global_ncols + threadIdx.y * warp_step_m_or_n * global_ncols + (threadIdx.x / threads_per_row) * global_ncols + global_iter_k * CTA_K + (threadIdx.x % threads_per_row) * PACK_SIZE);\n if constexpr (STAGES > 1)\n {\n uint32_t addr = cast_smem_ptr_to_uint(dst_ptr);\n cp_async_cg_A(addr, src_ptr, local_mask & (ld_row + cta_offset_m < global_nrows));\n }\n else\n {\n if (local_mask & (ld_row + cta_offset_m < global_nrows))\n *(uint4 *)dst_ptr = *src_ptr;\n }\n }\n}\n\ntemplate \n__device__ __inline__ void global_to_share_one_stage_B_T2(f16_t *src, f16_t *dst, int global_ncols, int cta_offset_m, int cta_offset_n, int global_iter_k, int shared_iter_k, bool mask)\n{\n constexpr int threads_needed = (CTA_N / kInterleave * CTA_K) / PACK_SIZE / SHARED_K_ITERS;\n constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE;\n constexpr int total_global_iters = (CTA_N / kInterleave * CTA_K) / PACK_SIZE / threads_used;\n constexpr int partial_global_iters = (total_global_iters + SHARED_K_ITERS - 1) / SHARED_K_ITERS;\n constexpr int cta_step_m_or_n = (threads_used * PACK_SIZE) / CTA_K;\n constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K;\n constexpr int threads_per_row = CTA_K / PACK_SIZE;\n constexpr int kSmemCol = CTA_K + SMEM_PAD_B;\n bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used);\n#pragma unroll\n for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter)\n {\n int global_iter = shared_iter_k * partial_global_iters + _global_iter;\n\n int ld_row = global_iter * cta_step_m_or_n + threadIdx.y * warp_step_m_or_n + (threadIdx.x / threads_per_row);\n int ld_col = (threadIdx.x % threads_per_row);\n int ld_col_swizzled = ld_col ^ (ld_row % 2) & 7;\n void *dst_ptr = (void *)(dst + (ld_row * kSmemCol + ld_col_swizzled * PACK_SIZE));\n uint4 *src_ptr = (uint4 *)(src + global_iter_k * CTA_K + cta_offset_n / kInterleave * global_ncols + ld_row * global_ncols + ld_col * PACK_SIZE);\n if constexpr (STAGES > 1)\n {\n uint32_t addr = cast_smem_ptr_to_uint(dst_ptr);\n cp_async_cg_A(addr, src_ptr, local_mask);\n }\n else\n {\n if (local_mask)\n *(uint4 *)dst_ptr = *src_ptr;\n }\n }\n}\n\ntemplate \n__device__ __inline__ void global_to_share_one_stage_scales_T2(f16_t *src, f16_t *dst, f16_t *src_z, f16_t *dst_z, int global_ncols, int cta_offset_m, int cta_offset_n, int global_iter_k, int shared_iter_k, bool mask)\n{\n constexpr int threads_needed = CTA_N / PACK_SIZE / 1;\n constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE;\n constexpr int total_global_iters = CTA_N / PACK_SIZE / threads_used;\n constexpr int threads_per_row = CTA_N / PACK_SIZE;\n constexpr int kSmemCol = CTA_N;\n bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used);\n int g_idx = global_iter_k * CTA_K / G;\n\n void *dst_ptr = (void *)(dst + (threadIdx.x % threads_per_row) * PACK_SIZE);\n uint4 *src_ptr = (uint4 *)(src + g_idx * global_ncols + cta_offset_n + (threadIdx.x % threads_per_row) * PACK_SIZE);\n void *dst_ptr_z = (void *)(dst_z + (threadIdx.x % threads_per_row) * PACK_SIZE);\n uint4 *src_ptr_z = (uint4 *)(src_z + g_idx * global_ncols + cta_offset_n + (threadIdx.x % threads_per_row) * PACK_SIZE);\n if (STAGES > 1)\n {\n uint32_t addr = cast_smem_ptr_to_uint(dst_ptr);\n cp_async_cg_A(addr, src_ptr, local_mask);\n uint32_t addr_z = cast_smem_ptr_to_uint(dst_ptr_z);\n cp_async_cg_A(addr_z, src_ptr_z, local_mask);\n }\n else\n {\n if (local_mask)\n {\n *(uint4 *)dst_ptr = *src_ptr;\n *(uint4 *)dst_ptr_z = *src_ptr_z;\n }\n }\n}\n\ntemplate \n__device__ __inline__ void share_to_reg_one_stage_A_T2(f16_t *src, f16_t *dst, int warp_offset_m, int warp_offset_n, int k_0_1)\n{\n constexpr int kSmemCol = CTA_K + SMEM_PAD_A;\n\n for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter)\n {\n\n int ld_row = warp_offset_m + shared_iter * OP_M + (threadIdx.x % 16);\n int ld_col = k_0_1 * 16 + (threadIdx.x / 16) * 8;\n int ld_col_swizzled = ((ld_col / PACK_SIZE) ^ (ld_row) & 7) * PACK_SIZE;\n void *addr_ptr = (void *)(src + ld_row * kSmemCol + ld_col_swizzled);\n\n uint32_t addr = cast_smem_ptr_to_uint(addr_ptr);\n ldmatrix_m8n8_x4_b16(dst, shared_iter, addr);\n }\n}\n\ntemplate \n__device__ __inline__ void share_to_reg_one_stage_B_T2(f16_t *src, f16_t *src_scales, f16_t *src_zeros, f16_t *dst, f16_t *dst_fp16, int warp_offset_m, int warp_offset_n, int k_0_1)\n{\n using f162_t = typename packed_as::type;\n constexpr int kSmemCol = CTA_K + SMEM_PAD_B;\n int r0 = ((threadIdx.x / 8 / 2) * 8 + threadIdx.x % 8);\n int c0 = ((threadIdx.x / 8) % 2) * 8;\n int r = r0 / 4;\n int c = (r0 % 4) * 16 + c0;\n int c_swizzled = ((c / PACK_SIZE) ^ (r % 2) & 7) * PACK_SIZE;\n\n if constexpr (ldmatrix)\n {\n#pragma unroll\n for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter)\n {\n void *addr_ptr = (void *)(src + warp_offset_n / kInterleave * kSmemCol + shared_iter * 16 / kInterleave * kSmemCol + k_0_1 * 16 + r * kSmemCol + c_swizzled);\n uint32_t addr = cast_smem_ptr_to_uint(addr_ptr);\n ldmatrix_m8n8_x4_b16(dst, shared_iter, addr);\n }\n }\n\n#pragma unroll\n for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter)\n {\n f16_t scale = src_scales[warp_offset_n + 16 * shared_iter + 8 * (k_0_1 % 2) + threadIdx.x / 4];\n f16_t zero = src_zeros[warp_offset_n + 16 * shared_iter + 8 * (k_0_1 % 2) + threadIdx.x / 4];\n f162_t scale2 = f162f162(scale);\n f162_t zero2 = f162f162(zero);\n f162_t loaded[4];\n dequantize_s4_to_f16x2(*reinterpret_cast(dst + (k_0_1 % 2) * 4 + (k_0_1 / 2 * 2) + shared_iter * 8), reinterpret_cast(loaded));\n#pragma unroll\n for (int i = 0; i < 4; i++)\n {\n loaded[i] = __hfma2(loaded[i], scale2, zero2);\n }\n *reinterpret_cast(dst_fp16 + shared_iter * 16 + 8 * (k_0_1 % 2)) = *reinterpret_cast(loaded);\n }\n}\n\ntemplate \n__global__ void gemm_w4a16_T2(f16_t *__restrict__ A, f16_t *__restrict__ B, f16_t *__restrict__ scales, f16_t *__restrict__ zeros, f16_t *__restrict__ C, int M, int N, int K)\n{\n using f162_t = typename packed_as::type;\n constexpr int NUM_WARPS = CTA_M / WARP_M * CTA_N / WARP_N;\n constexpr int CTA_SIZE = NUM_WARPS * WARP_SIZE;\n int num_blocks_n = (N + CTA_N - 1) / CTA_N;\n int num_blocks_m = (M + CTA_M - 1) / CTA_M;\n int blockIdx_x = 0;\n int blockIdx_y = blockIdx.x % (num_blocks_m * num_blocks_n);\n int blockIdx_z = blockIdx.x / (num_blocks_m * num_blocks_n);\n const int log_tile = get_log_tile<1>((N + CTA_N - 1) / CTA_N);\n int blockIdx_m = blockIdx_y / (num_blocks_n >> log_tile);\n int blockIdx_n = blockIdx_y % (num_blocks_n >> log_tile);\n const uint2 block_idx_mapping = get_block_idx_mapping(blockIdx_m, blockIdx_n, log_tile);\n blockIdx_m = block_idx_mapping.x;\n blockIdx_n = block_idx_mapping.y;\n\n float C_warp[CTA_M * CTA_N / CTA_SIZE];\n constexpr int kSmemPadKA = CTA_K + SMEM_PAD_A;\n constexpr int kSmemPadKB = CTA_K + SMEM_PAD_B;\n constexpr int kSmemSizeAPerStage = CTA_M * kSmemPadKA;\n constexpr int kSmemSizeBPerStage = CTA_N / kInterleave * kSmemPadKB;\n constexpr int kSmemSizeA = kSmemSizeAPerStage * STAGES;\n constexpr int kSmemSizeB = kSmemSizeBPerStage * STAGES;\n constexpr int kSmemSizeScales = CTA_N * STAGES / 2;\n constexpr int kSmemSizeZeros = CTA_N * STAGES / 2;\n constexpr int scales_load_interval = G / CTA_K;\n extern __shared__ half mem_shared[];\n f16_t *A_shared = reinterpret_cast(mem_shared);\n f16_t *B_shared = reinterpret_cast(mem_shared + kSmemSizeA);\n f16_t *scales_shared = reinterpret_cast(mem_shared + kSmemSizeA + kSmemSizeB);\n f16_t *zeros_shared = reinterpret_cast(mem_shared + kSmemSizeA + kSmemSizeB + kSmemSizeScales);\n f16_t A_shared_warp_[2][WARP_M * INTRIN_K / WARP_SIZE];\n f16_t B_shared_warp_[2][WARP_N * 32 / WARP_SIZE];\n f16_t B_shared_warp_tmp_[2][WARP_N * 16 / WARP_SIZE];\n int cta_offset_m = blockIdx_m * CTA_M;\n int cta_offset_n = blockIdx_n * CTA_N;\n int warp_offset_m = (threadIdx.y % (CTA_M / WARP_M)) * WARP_M;\n int warp_offset_n = (threadIdx.y / (CTA_M / WARP_M)) * WARP_N;\n\n for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE; i++)\n C_warp[i] = 0.0;\n\n int gemm_iters = (K + CTA_K - 1) / CTA_K;\n int k_0_0_ld = 0;\n int k_0_0 = 0;\n constexpr int prologue_stages = STAGES == 1 ? 1 : STAGES - 1;\n#pragma unroll\n for (k_0_0_ld = 0; k_0_0_ld < prologue_stages; ++k_0_0_ld)\n {\n global_to_share_one_stage_A_T2(A, A_shared + k_0_0_ld * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, k_0_0_ld, 0, true);\n global_to_share_one_stage_B_T2(B, B_shared + k_0_0_ld * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, k_0_0_ld, 0, true);\n global_to_share_one_stage_scales_T2(\n scales, scales_shared + (k_0_0_ld / scales_load_interval) * CTA_N,\n zeros, zeros_shared + (k_0_0_ld / scales_load_interval) * CTA_N,\n N, cta_offset_m, cta_offset_n, k_0_0_ld, 0, k_0_0_ld < gemm_iters && k_0_0_ld % scales_load_interval == 0);\n if constexpr (STAGES > 1)\n __pipeline_commit();\n }\n if constexpr (STAGES > 1)\n __pipeline_wait_prior(STAGES - 2);\n __syncthreads();\n\n share_to_reg_one_stage_A_T2(A_shared, A_shared_warp_[0], warp_offset_m, warp_offset_n, 0);\n share_to_reg_one_stage_B_T2(B_shared, scales_shared, zeros_shared, B_shared_warp_tmp_[0], B_shared_warp_[0], warp_offset_m, warp_offset_n, 0);\n constexpr int SHARED_K_ITERS = WARP_K / INTRIN_K;\n\n for (; k_0_0 < gemm_iters; ++k_0_0, ++k_0_0_ld)\n {\n int ld_stage = k_0_0_ld % STAGES;\n int compute_stage = k_0_0 % STAGES;\n f16_t *A_shared_this_compute_stage;\n f16_t *B_shared_this_compute_stage;\n f16_t *scales_shared_this_compute_stage;\n f16_t *zeros_shared_this_compute_stage;\n\n for (int iter_k = 0; iter_k < SHARED_K_ITERS; ++iter_k)\n {\n A_shared_this_compute_stage = A_shared + compute_stage * kSmemSizeAPerStage;\n B_shared_this_compute_stage = B_shared + compute_stage * kSmemSizeBPerStage;\n scales_shared_this_compute_stage = scales_shared + (compute_stage / scales_load_interval) * CTA_N;\n zeros_shared_this_compute_stage = zeros_shared + (compute_stage / scales_load_interval) * CTA_N;\n share_to_reg_one_stage_A_T2(A_shared_this_compute_stage, A_shared_warp_[(iter_k + 1) % 2], warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS);\n if ((iter_k + 1) % kInterleave == 0)\n {\n if (compute_stage % 2 == 1)\n {\n share_to_reg_one_stage_B_T2(\n B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage,\n B_shared_warp_tmp_[1], B_shared_warp_[((iter_k + 1) / 2) % 2],\n warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS);\n }\n else\n {\n share_to_reg_one_stage_B_T2(\n B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage,\n B_shared_warp_tmp_[0], B_shared_warp_[((iter_k + 1) / 2) % 2],\n warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS);\n }\n }\n else\n {\n if (compute_stage % 2 == 1)\n {\n share_to_reg_one_stage_B_T2(\n B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage,\n B_shared_warp_tmp_[1], B_shared_warp_[((iter_k + 1) / 2) % 2],\n warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS);\n }\n else\n {\n share_to_reg_one_stage_B_T2(\n B_shared_this_compute_stage, scales_shared_this_compute_stage, zeros_shared_this_compute_stage,\n B_shared_warp_tmp_[0], B_shared_warp_[((iter_k + 1) / 2) % 2],\n warp_offset_m, warp_offset_n, (iter_k + 1) % SHARED_K_ITERS);\n }\n }\n __syncthreads();\n f16_t *A_shared_warp = A_shared_warp_[iter_k % 2];\n f16_t *B_shared_warp = B_shared_warp_[(iter_k / 2) % 2];\n for (int i_0_3 = 0; i_0_3 < WARP_M / INTRIN_M; ++i_0_3)\n {\n for (int j_0_4 = 0; j_0_4 < WARP_N / INTRIN_N; ++j_0_4)\n {\n mma_m16n8k16(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4);\n mma_m16n8k16(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4, A_shared_warp + i_0_3 * 8, B_shared_warp + j_0_4 * 16 + (iter_k % 2) * 4 + 8);\n }\n }\n\n if (iter_k < WARP_K / INTRIN_K - 1)\n {\n if constexpr (STAGES == 1)\n __syncthreads();\n global_to_share_one_stage_A_T2(A, A_shared + ld_stage * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters);\n global_to_share_one_stage_B_T2(B, B_shared + ld_stage * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters);\n }\n\n if (iter_k == WARP_K / INTRIN_K - 2)\n {\n if constexpr (STAGES == 1 && WARP_K / INTRIN_K > 2)\n {\n __syncthreads();\n }\n global_to_share_one_stage_A_T2(A, A_shared + ld_stage * kSmemSizeAPerStage, M, K, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, k_0_0_ld < gemm_iters);\n global_to_share_one_stage_B_T2(B, B_shared + ld_stage * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k + 1, k_0_0_ld < gemm_iters);\n global_to_share_one_stage_scales_T2(\n scales, scales_shared + (ld_stage / scales_load_interval) * CTA_N,\n zeros, zeros_shared + (ld_stage / scales_load_interval) * CTA_N,\n N, cta_offset_m, cta_offset_n, k_0_0_ld, iter_k, k_0_0_ld < gemm_iters && k_0_0_ld % scales_load_interval == 0);\n if constexpr (STAGES > 1)\n {\n __pipeline_commit();\n __pipeline_wait_prior(STAGES - 2);\n }\n compute_stage = (k_0_0 + 1) % STAGES;\n __syncthreads();\n }\n }\n }\n for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1)\n {\n for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1)\n {\n for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2)\n {\n int write_row = cta_offset_m + warp_offset_m + ax0_0_1 * OP_M + ((local_id % 4) / 2 * 8 + (threadIdx.x / 4));\n if (write_row < M)\n {\n *reinterpret_cast(\n C + write_row * N +\n cta_offset_n + warp_offset_n + ax1_0_1 * 16 +\n (local_id / 4) * 8 + (local_id % 2) + (threadIdx.x % 4) * 2) =\n cuda_cast(*reinterpret_cast(C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 +\n ax1_0_1 * 8 + local_id));\n }\n };\n }\n }\n}\n\ntorch::Tensor awq_gemm_forward_cuda(\n torch::Tensor _in_feats,\n torch::Tensor _kernel,\n torch::Tensor _scales,\n torch::Tensor _zeros)\n{\n std::vector output_shape = _in_feats.sizes().vec();\n output_shape.back() = _kernel.size(0) * kInterleave;\n int num_in_feats = _in_feats.numel() / _in_feats.size(-1);\n int num_in_channels = _in_feats.size(-1);\n auto options =\n torch::TensorOptions().dtype(_in_feats.dtype()).device(_in_feats.device());\n auto options_int =\n torch::TensorOptions().dtype(torch::kInt32).device(_in_feats.device());\n at::Tensor _out_feats = torch::empty(output_shape, options);\n int num_out_feats = _out_feats.numel() / _out_feats.size(-1);\n int num_out_channels = _out_feats.size(-1);\n\n if (_in_feats.scalar_type() == at::ScalarType::Half)\n {\n using f16_t = half;\n\n auto in_feats = reinterpret_cast(_in_feats.data_ptr());\n auto kernel = reinterpret_cast(_kernel.data_ptr());\n auto scales = reinterpret_cast(_scales.data_ptr());\n auto zeros = reinterpret_cast(_zeros.data_ptr());\n auto out_feats = reinterpret_cast(_out_feats.data_ptr());\n\n if (num_out_feats <= 32)\n {\n constexpr int G = 128;\n constexpr int CTA_M = 16;\n constexpr int CTA_N = 128;\n constexpr int CTA_K = 128;\n constexpr int WARP_M = 16;\n constexpr int WARP_N = 32;\n constexpr int WARP_K = 64;\n constexpr int SPLITK = 2;\n constexpr int STAGES = 4;\n KERNEL_LAUNCH_CODE\n }\n else if (num_out_feats <= 64)\n {\n\n constexpr int G = 128;\n constexpr int CTA_M = 16;\n constexpr int CTA_N = 128;\n constexpr int CTA_K = 128;\n constexpr int WARP_M = 16;\n constexpr int WARP_N = 32;\n constexpr int WARP_K = 64;\n constexpr int SPLITK = 1;\n constexpr int STAGES = 3;\n KERNEL_LAUNCH_CODE\n }\n else if (num_out_feats <= 128)\n {\n constexpr int G = 128;\n constexpr int CTA_M = 32;\n constexpr int CTA_N = 128;\n constexpr int CTA_K = 128;\n constexpr int WARP_M = 32;\n constexpr int WARP_N = 32;\n constexpr int WARP_K = 64;\n constexpr int SPLITK = 1;\n constexpr int STAGES = 4;\n KERNEL_LAUNCH_CODE\n }\n else if (num_out_feats <= 192)\n {\n constexpr int G = 128;\n constexpr int CTA_M = 64;\n constexpr int CTA_N = 128;\n constexpr int CTA_K = 64;\n constexpr int WARP_M = 64;\n constexpr int WARP_N = 32;\n constexpr int WARP_K = 64;\n constexpr int SPLITK = 1;\n constexpr int STAGES = 4;\n KERNEL_LAUNCH_CODE\n }\n else\n {\n constexpr int G = 128;\n constexpr int CTA_M = 64;\n constexpr int CTA_N = 128;\n constexpr int CTA_K = 64;\n constexpr int WARP_M = 64;\n constexpr int WARP_N = 32;\n constexpr int WARP_K = 64;\n constexpr int STAGES = 4;\n\n constexpr int NUM_WARPS = (CTA_M / WARP_M) * (CTA_N / WARP_N);\n constexpr int kSmemByteSize = (CTA_M * (CTA_K + SMEM_PAD_A) + CTA_N * (CTA_K + SMEM_PAD_B) / kInterleave + CTA_N) * STAGES * sizeof(f16_t);\n if (kSmemByteSize >= 99 * 1024)\n {\n printf(\"This kernel requires %d Bytes of shared memory, which exceeds device limit.\\n\", kSmemByteSize);\n return _out_feats;\n }\n int j_factors1 = num_out_channels / CTA_N / 1;\n dim3 num_blocks((num_out_feats + CTA_M - 1) / CTA_M * j_factors1);\n dim3 threads_per_block(WARP_SIZE, NUM_WARPS);\n auto kernel_func = gemm_w4a16_T2;\n cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize);\n kernel_func<<>>(\n in_feats, kernel, scales, zeros, out_feats, num_in_feats, num_out_channels, num_in_channels);\n }\n }\n else if (_in_feats.scalar_type() == at::ScalarType::BFloat16)\n {\n using f16_t = __nv_bfloat16;\n\n auto in_feats = reinterpret_cast(_in_feats.data_ptr());\n auto kernel = reinterpret_cast(_kernel.data_ptr());\n auto scales = reinterpret_cast(_scales.data_ptr());\n auto zeros = reinterpret_cast(_zeros.data_ptr());\n auto out_feats = reinterpret_cast(_out_feats.data_ptr());\n\n if (num_out_feats <= 32)\n {\n constexpr int G = 128;\n constexpr int CTA_M = 16;\n constexpr int CTA_N = 128;\n constexpr int CTA_K = 128;\n constexpr int WARP_M = 16;\n constexpr int WARP_N = 32;\n constexpr int WARP_K = 64;\n constexpr int SPLITK = 2;\n constexpr int STAGES = 4;\n KERNEL_LAUNCH_CODE\n }\n else if (num_out_feats <= 64)\n {\n\n constexpr int G = 128;\n constexpr int CTA_M = 16;\n constexpr int CTA_N = 128;\n constexpr int CTA_K = 128;\n constexpr int WARP_M = 16;\n constexpr int WARP_N = 32;\n constexpr int WARP_K = 64;\n constexpr int SPLITK = 1;\n constexpr int STAGES = 3;\n KERNEL_LAUNCH_CODE\n }\n else if (num_out_feats <= 128)\n {\n constexpr int G = 128;\n constexpr int CTA_M = 32;\n constexpr int CTA_N = 128;\n constexpr int CTA_K = 128;\n constexpr int WARP_M = 32;\n constexpr int WARP_N = 32;\n constexpr int WARP_K = 64;\n constexpr int SPLITK = 1;\n constexpr int STAGES = 4;\n KERNEL_LAUNCH_CODE\n }\n else if (num_out_feats <= 192)\n {\n constexpr int G = 128;\n constexpr int CTA_M = 64;\n constexpr int CTA_N = 128;\n constexpr int CTA_K = 64;\n constexpr int WARP_M = 64;\n constexpr int WARP_N = 32;\n constexpr int WARP_K = 64;\n constexpr int SPLITK = 1;\n constexpr int STAGES = 4;\n KERNEL_LAUNCH_CODE\n }\n else\n {\n constexpr int G = 128;\n constexpr int CTA_M = 64;\n constexpr int CTA_N = 128;\n constexpr int CTA_K = 64;\n constexpr int WARP_M = 64;\n constexpr int WARP_N = 32;\n constexpr int WARP_K = 64;\n constexpr int STAGES = 4;\n\n constexpr int NUM_WARPS = (CTA_M / WARP_M) * (CTA_N / WARP_N);\n constexpr int kSmemByteSize = (CTA_M * (CTA_K + SMEM_PAD_A) + CTA_N * (CTA_K + SMEM_PAD_B) / kInterleave + CTA_N) * STAGES * sizeof(f16_t);\n if (kSmemByteSize >= 99 * 1024)\n {\n printf(\"This kernel requires %d Bytes of shared memory, which exceeds device limit.\\n\", kSmemByteSize);\n return _out_feats;\n }\n int j_factors1 = num_out_channels / CTA_N / 1;\n dim3 num_blocks((num_out_feats + CTA_M - 1) / CTA_M * j_factors1);\n dim3 threads_per_block(WARP_SIZE, NUM_WARPS);\n auto kernel_func = gemm_w4a16_T2;\n cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize);\n kernel_func<<>>(\n in_feats, kernel, scales, zeros, out_feats, num_in_feats, num_out_channels, num_in_channels);\n }\n }\n else\n {\n AT_ERROR(\"Unsupported input type\");\n }\n\n return _out_feats;\n}"
},
{
"path": "deepcompressor/backend/tinychat/csrc/quantization/gemm/gemm_cuda.h",
"content": "#include \n\ntorch::Tensor awq_gemm_forward_cuda(\n torch::Tensor _in_feats,\n torch::Tensor _kernel,\n torch::Tensor _scales,\n torch::Tensor _zeros);\n"
},
{
"path": "deepcompressor/backend/tinychat/csrc/quantization/gemm/semaphore.h",
"content": "/***************************************************************************************************\n * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n *\n * 1. Redistributions of source code must retain the above copyright notice, this\n * list of conditions and the following disclaimer.\n *\n * 2. Redistributions in binary form must reproduce the above copyright notice,\n * this list of conditions and the following disclaimer in the documentation\n * and/or other materials provided with the distribution.\n *\n * 3. Neither the name of the copyright holder nor the names of its\n * contributors may be used to endorse or promote products derived from\n * this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\n * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\n * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\n * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\n * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\n * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n **************************************************************************************************/\n/*! \\file\n \\brief Implementation of a CTA-wide semaphore for inter-CTA synchronization.\n*/\n\n#pragma once\n\n/////////////////////////////////////////////////////////////////////////////////////////////////\n\n// namespace cutlass {\n\n/////////////////////////////////////////////////////////////////////////////////////////////////\n\n/// CTA-wide semaphore for inter-CTA synchronization.\nclass Semaphore\n{\npublic:\n int *lock;\n bool wait_thread;\n int state;\n\npublic:\n /// Implements a semaphore to wait for a flag to reach a given value\n __host__ __device__ Semaphore(int *lock_, int thread_id) : lock(lock_),\n wait_thread(thread_id < 0 || thread_id == 0),\n state(-1)\n {\n }\n\n /// Permit fetching the synchronization mechanism early\n __device__ void fetch()\n {\n if (wait_thread)\n {\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700\n asm volatile(\"ld.global.acquire.gpu.b32 %0, [%1];\\n\" : \"=r\"(state) : \"l\"(lock));\n#else\n asm volatile(\"ld.global.cg.b32 %0, [%1];\\n\" : \"=r\"(state) : \"l\"(lock));\n#endif\n }\n }\n\n /// Gets the internal state\n __device__ int get_state() const\n {\n return state;\n }\n\n /// Waits until the semaphore is equal to the given value\n __device__ void wait(int status = 0)\n {\n while (__syncthreads_and(state != status))\n {\n fetch();\n }\n\n __syncthreads();\n }\n\n /// Updates the lock with the given result\n __device__ void release(int status = 0)\n {\n __syncthreads();\n\n if (wait_thread)\n {\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700\n asm volatile(\"st.global.release.gpu.b32 [%0], %1;\\n\" : : \"l\"(lock), \"r\"(status));\n#else\n asm volatile(\"st.global.cg.b32 [%0], %1;\\n\" : : \"l\"(lock), \"r\"(status));\n#endif\n }\n }\n};\n\n/////////////////////////////////////////////////////////////////////////////////////////////////\n\n// } // namespace cutlass\n\n/////////////////////////////////////////////////////////////////////////////////////////////////\n"
},
{
"path": "deepcompressor/backend/tinychat/csrc/quantization/gemv/gemv_cuda.cu",
"content": "/*\n * Modified from NVIDIA [TRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/tree/d37b507f41a87457fe9f10f7459d08f5db235745/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv)\n * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved.\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n/*\n@article{lin2023awq,\n title={AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration},\n author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Dang, Xingyu and Han, Song},\n journal={arXiv},\n year={2023}\n}\n*/\n\n#include \n#include \n#include \"gemv_cuda.h\"\n#include \"../dequantize.cuh\"\n#include \"../../utils.cuh\"\n\n#define PACK_FACTOR 8\n#define WARP_SIZE 32\n#define MEM_ACCESS_SIZE 128\n\n// Reduce sum within the warp using the tree reduction algorithm.\ntemplate \n__device__ __forceinline__ static void warp_reduce(fp_t *psum, float (*out_smem)[Num * 4])\n{\n // kInterleave = 4\n float fpsum[Num];\n#pragma unroll\n for (int i = 0; i < Num; ++i)\n {\n fpsum[i] = static_cast(psum[i]);\n }\n\n#pragma unroll\n for (int i = 0; i < Num; ++i)\n {\n // T0 + T1 + T8 + T9 + T16 + T17 + T24 + T25 (kInterleave = 4)\n fpsum[i] += __shfl_xor_sync(~0, fpsum[i], 16);\n fpsum[i] += __shfl_xor_sync(~0, fpsum[i], 8);\n fpsum[i] += __shfl_xor_sync(~0, fpsum[i], 1);\n }\n __syncthreads();\n int warp = threadIdx.x / WarpSize, lane = threadIdx.x % WarpSize;\n if (lane == 0 || lane == 2 || lane == 4 || lane == 6)\n {\n#pragma unroll\n for (int i = 0; i < Num; ++i)\n {\n out_smem[warp][i * 4 + lane / 2] = fpsum[i];\n }\n }\n __syncthreads();\n};\n\n__device__ __forceinline__ int make_divisible(int c, int divisor)\n{\n return (c + divisor - 1) / divisor;\n}\n\ntemplate \n__global__ void gemv_kernel(\n const f16_t *inputs, const uint32_t *weight, const f16_t *scales, const f16_t *zeros, f16_t *outputs,\n const int IC, const int OC)\n{\n\n using f162_t = typename packed_as::type;\n using accum_t = float;\n using accum2_t = typename packed_as::type;\n\n const int kStride = 64;\n const int kElemsPerThread = MEM_ACCESS_SIZE / 4;\n const int kThreadsNumPerTile = kStride / kElemsPerThread;\n\n static constexpr int kShuffleBasicTile = 2;\n static constexpr int kShuffleContinous = 4;\n static constexpr int kShuffleStrided = 4;\n\n constexpr int Num = NPerBlock * Batch;\n constexpr int kInterleave = 4;\n\n alignas(16) f16_t local_inputs[kElemsPerThread];\n alignas(16) uint32_t local_qweights[MEM_ACCESS_SIZE / 32];\n alignas(16) f16_t half_weight_buffer[kElemsPerThread];\n alignas(16) f16_t dequantized_weight[kElemsPerThread * NPerBlock];\n alignas(16) f16_t local_scale[NPerBlock];\n alignas(16) f16_t local_scaled_zeros[NPerBlock];\n\n accum_t psum[Num];\n for (int i = 0; i < Num; ++i)\n psum[i] = static_cast(0.f);\n\n extern __shared__ uint8_t shmem[];\n float(*out_smem)[Num * kInterleave] = reinterpret_cast(shmem);\n\n const int blk_row_offset = blockIdx.x * NPerBlock * kInterleave;\n const int thd_row_offset = (threadIdx.x / kThreadsNumPerTile) % kInterleave;\n const int act_k_offset = threadIdx.x / (kThreadsNumPerTile * kInterleave) * kStride + (threadIdx.x % kThreadsNumPerTile) * kElemsPerThread;\n const int group_offset = act_k_offset / GroupSize;\n // TODO: use make_divisible\n const uint32_t *blk_weight_ptr = weight + blk_row_offset * IC / PACK_FACTOR;\n const f16_t *scale_ptr = scales + blk_row_offset + thd_row_offset + group_offset * OC;\n const f16_t *zeros_ptr = zeros + blk_row_offset + thd_row_offset + group_offset * OC;\n const f16_t *inputs_ptr = inputs + act_k_offset;\n\n const int act_forward_step = BlockSize * kElemsPerThread / kInterleave;\n const int scale_forward_step = act_forward_step / GroupSize * OC;\n\n // Main loop iteration, each block completes the outputs for several OCs\n for (int kk = threadIdx.x * kElemsPerThread; kk < IC * kInterleave; kk += BlockSize * kElemsPerThread)\n {\n// Load qweight, scales and scaled_zeros\n#pragma unroll\n for (int idx = 0; idx < NPerBlock; ++idx)\n {\n // use float4 to load weights, each thread load 32 int4 numbers (1 x float4, 128 bit)\n *((float4 *)(local_qweights)) =\n *((float4 *)(blk_weight_ptr + (idx * kInterleave * IC + kk) / PACK_FACTOR));\n local_scale[idx] = *(scale_ptr + idx * kInterleave);\n local_scaled_zeros[idx] = *(zeros_ptr + idx * kInterleave);\n\n// Map int4 qweight to fp format\n#pragma unroll\n for (int i = 0; i < MEM_ACCESS_SIZE / 32; ++i)\n {\n // Converts 32 bits (8 x int4) to 8 fp16\n dequantize_s4_to_f16x2(*reinterpret_cast(local_qweights + i), reinterpret_cast(half_weight_buffer + i * PACK_FACTOR));\n }\n\n// Dequantize (apply s/z) and shuffle elements to match the weight packing format\n#pragma unroll\n for (int i = 0; i < kShuffleContinous; ++i)\n {\n#pragma unroll\n for (int j = 0; j < kShuffleStrided; ++j)\n {\n f162_t w =\n *reinterpret_cast(\n half_weight_buffer + (i + j * kShuffleContinous) * kShuffleBasicTile);\n w = __hfma2(w, f162f162(local_scale[idx]), f162f162(local_scaled_zeros[idx]));\n dequantized_weight[((i * kShuffleStrided + j) * kShuffleBasicTile + 0) * NPerBlock + idx] = w.x;\n dequantized_weight[((i * kShuffleStrided + j) * kShuffleBasicTile + 1) * NPerBlock + idx] = w.y;\n }\n }\n }\n#pragma unroll\n for (int batch_idx = 0; batch_idx < Batch; ++batch_idx)\n {\n const f16_t *local_inputs_ptr = inputs_ptr + batch_idx * IC;\n#pragma unroll\n for (int idx = 0; idx < kElemsPerThread / 8; ++idx)\n {\n // load activation, 8 halves (128 bits) / step.\n *((float4 *)(local_inputs + idx * 8)) = *((float4 *)(local_inputs_ptr + idx * 8));\n }\n// Perform the MACs\n#pragma unroll\n for (int x = 0; x < NPerBlock / 2; ++x)\n {\n#pragma unroll\n for (int y = 0; y < kElemsPerThread; ++y)\n {\n accum2_t prod = cuda_cast(__hmul2(\n *reinterpret_cast(dequantized_weight + y * NPerBlock + x * 2),\n f162f162(local_inputs[y])));\n *reinterpret_cast(psum + batch_idx * NPerBlock + x * 2) = prod + *reinterpret_cast(psum + batch_idx * NPerBlock + x * 2);\n }\n }\n }\n inputs_ptr += act_forward_step;\n scale_ptr += scale_forward_step;\n zeros_ptr += scale_forward_step;\n }\n\n warp_reduce(psum, out_smem);\n\n // Num * Interleave = batch * NPerBlock * Interleave -> 1 thread_block write back num\n for (int i = threadIdx.x; i < Num * kInterleave; i += BlockSize)\n {\n int batch_idx = i / (NPerBlock * kInterleave);\n int oc_idx = i % (NPerBlock * kInterleave);\n float acc = 0.f;\n for (int j = 0; j < BlockSize / WARP_SIZE; ++j)\n {\n acc += out_smem[j][i];\n }\n outputs[batch_idx * OC + blk_row_offset + oc_idx] = static_cast(acc);\n }\n}\n\n/*\nComputes GEMV (PyTorch interface).\n\nArgs:\n _in_feats: tensor of shape [B, IC];\n _kernel: int tensor of shape [OC, IC // 8];\n _zeros: int tensor of shape [OC, IC // G // 8];\n _scaling_factors: tensor of shape [OC, IC // G];\n blockDim_x: size of thread block, dimension x, where blockDim_x * workload_per_thread = IC;\n blockDim_y: size of thread block, dimension y, where blockDim_y * gridDim_y = OC;\n\nReturns:\n out_feats: tensor of shape [B, OC];\n*/\ntorch::Tensor awq_gemv_forward_cuda(\n torch::Tensor _in_feats,\n torch::Tensor _kernel,\n torch::Tensor _scaling_factors,\n torch::Tensor _zeros,\n int m,\n int n,\n int k,\n int group_size)\n{\n std::vector output_shape = _in_feats.sizes().vec();\n output_shape.back() = n;\n\n auto options = torch::TensorOptions().dtype(_in_feats.dtype()).device(_in_feats.device());\n at::Tensor _out_feats = torch::empty(output_shape, options);\n\n static constexpr int N_PER_BLOCK = 2;\n static constexpr int K_INTERLEAVE = 4;\n static constexpr int BLOCK_SIZE = 256;\n dim3 num_blocks(n / N_PER_BLOCK / K_INTERLEAVE);\n dim3 num_threads(BLOCK_SIZE);\n\n AT_DISPATCH_REDUCED_FLOATING_TYPES(\n _in_feats.scalar_type(),\n \"awq_gemv_forward_cuda\",\n [&]\n {\n using f16_t = typename to_cpp_t::type;\n auto in_feats = reinterpret_cast(_in_feats.data_ptr());\n auto kernel = reinterpret_cast(_kernel.data_ptr());\n auto zeros = reinterpret_cast(_zeros.data_ptr());\n auto scaling_factors = reinterpret_cast(_scaling_factors.data_ptr());\n auto out_feats = reinterpret_cast(_out_feats.data_ptr());\n\n if (group_size == 128)\n {\n switch (m)\n {\n case 1:\n gemv_kernel<<>>(\n in_feats, kernel, scaling_factors, zeros, out_feats, k, n);\n break;\n case 2:\n gemv_kernel<<>>(\n in_feats, kernel, scaling_factors, zeros, out_feats, k, n);\n break;\n case 3:\n gemv_kernel<<>>(\n in_feats, kernel, scaling_factors, zeros, out_feats, k, n);\n break;\n case 4:\n gemv_kernel<<>>(\n in_feats, kernel, scaling_factors, zeros, out_feats, k, n);\n break;\n case 5:\n gemv_kernel<<>>(\n in_feats, kernel, scaling_factors, zeros, out_feats, k, n);\n break;\n case 6:\n gemv_kernel<<>>(\n in_feats, kernel, scaling_factors, zeros, out_feats, k, n);\n break;\n case 7:\n gemv_kernel<<>>(\n in_feats, kernel, scaling_factors, zeros, out_feats, k, n);\n break;\n default:\n throw std::runtime_error(\"Unsupported batch size for gemv kernel.\\n\");\n }\n }\n else\n {\n throw std::runtime_error(\"Unsupported group size for gemv kernel.\\n\");\n }\n });\n\n return _out_feats;\n}\n"
},
{
"path": "deepcompressor/backend/tinychat/csrc/quantization/gemv/gemv_cuda.h",
"content": "#pragma once\n#include \n\ntorch::Tensor awq_gemv_forward_cuda(\n torch::Tensor _in_feats,\n torch::Tensor _kernel,\n torch::Tensor _scaling_factors,\n torch::Tensor _zeros,\n int m,\n int n,\n int k,\n int group_size);\n"
},
{
"path": "deepcompressor/backend/tinychat/csrc/utils.cuh",
"content": "// Adated from FasterTransformer, https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp\n\n#pragma once\n\n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#define ENABLE_BF16 1\n\ntemplate \nstruct to_cpp_t;\ntemplate <>\nstruct to_cpp_t\n{\n using type = half;\n};\n\ntemplate <>\nstruct to_cpp_t\n{\n using type = __nv_bfloat16;\n};\n\ntemplate \nstruct num_elems;\ntemplate <>\nstruct num_elems\n{\n static constexpr int value = 1;\n};\ntemplate <>\nstruct num_elems\n{\n static constexpr int value = 2;\n};\ntemplate <>\nstruct num_elems\n{\n static constexpr int value = 4;\n};\ntemplate <>\nstruct num_elems\n{\n static constexpr int value = 1;\n};\ntemplate <>\nstruct num_elems\n{\n static constexpr int value = 2;\n};\n#ifdef ENABLE_BF16\ntemplate <>\nstruct num_elems<__nv_bfloat16>\n{\n static constexpr int value = 1;\n};\ntemplate <>\nstruct num_elems<__nv_bfloat162>\n{\n static constexpr int value = 2;\n};\n#endif\n\ntemplate \nstruct packed_as;\ntemplate \nstruct packed_as\n{\n using type = T;\n};\ntemplate <>\nstruct packed_as\n{\n using type = half2;\n};\ntemplate <>\nstruct packed_as\n{\n using type = float2;\n};\ntemplate <>\nstruct packed_as\n{\n using type = int16_t;\n};\ntemplate <>\nstruct packed_as\n{\n using type = int2;\n};\ntemplate <>\nstruct packed_as\n{\n using type = half;\n};\ntemplate <>\nstruct packed_as\n{\n using type = float;\n};\n#ifdef ENABLE_BF16\ntemplate <>\nstruct packed_as<__nv_bfloat16, 2>\n{\n using type = __nv_bfloat162;\n};\ntemplate <>\nstruct packed_as<__nv_bfloat162, 1>\n{\n using type = __nv_bfloat16;\n};\n#endif\n#ifdef ENABLE_FP8\ntemplate <>\nstruct packed_as<__nv_fp8_e4m3, 2>\n{\n using type = __nv_fp8x2_e4m3;\n};\ntemplate <>\nstruct packed_as<__nv_fp8x2_e4m3, 1>\n{\n using type = __nv_fp8_e4m3;\n};\ntemplate <>\nstruct packed_as<__nv_fp8_e5m2, 2>\n{\n using type = __nv_fp8x2_e5m2;\n};\ntemplate <>\nstruct packed_as<__nv_fp8x2_e5m2, 1>\n{\n using type = __nv_fp8_e5m2;\n};\n#endif\n\ntemplate \n__device__ __forceinline__\n packed_as::type\n f162f162(f16_t x);\n\ntemplate <>\n__device__ __forceinline__\n packed_as::type\n f162f162(half x)\n{\n return __half2half2(x);\n}\n\n#ifdef ENABLE_BF16\ntemplate <>\n__device__ __forceinline__\n packed_as<__nv_bfloat16, 2>::type\n f162f162<__nv_bfloat16>(__nv_bfloat16 x)\n{\n return __bfloat162bfloat162(x);\n}\n# endif\n\ntemplate \n__device__ __forceinline__\n float2\n f1622float2(T val);\n\ntemplate <>\n__device__ __forceinline__\n float2\n f1622float2(half2 val)\n{\n return __half22float2(val);\n}\n\n#ifdef ENABLE_BF16\ntemplate <>\n__device__ __forceinline__\n float2\n f1622float2<__nv_bfloat162>(__nv_bfloat162 val)\n{\n return __bfloat1622float2(val);\n}\n# endif\n\ninline __device__ float2 operator*(float2 a, float2 b) { return make_float2(a.x * b.x, a.y * b.y); }\ninline __device__ float2 operator+(float2 a, float2 b) { return make_float2(a.x + b.x, a.y + b.y); }\ninline __device__ float2 operator-(float2 a, float2 b) { return make_float2(a.x - b.x, a.y - b.y); }\n\ninline __device__ float2 operator*(float2 a, float b) { return make_float2(a.x * b, a.y * b); }\ninline __device__ float2 operator+(float2 a, float b) { return make_float2(a.x + b, a.y + b); }\ninline __device__ float2 operator-(float2 a, float b) { return make_float2(a.x - b, a.y - b); }\n\nstatic inline __device__ int8_t float_to_int8_rn(float x)\n{\n uint32_t dst;\n asm volatile(\"cvt.rni.sat.s8.f32 %0, %1;\" : \"=r\"(dst) : \"f\"(x));\n return reinterpret_cast(dst);\n}\n\ntemplate \ninline __device__ T ldg(const T *val)\n{\n return __ldg(val);\n}\n\n#if ENABLE_BF16\n\n#define float22bf162 __float22bfloat162_rn\n\ninline __device__ int16_t bf1622int16(__nv_bfloat162 val)\n{\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800\n float2 f_val;\n f_val.x = max(min(__low2float(val), 127.f), -128.f);\n f_val.y = max(min(__high2float(val), 127.f), -128.f);\n\n union\n {\n int8_t int8[2];\n int16_t int16;\n };\n\n int8[0] = static_cast(static_cast(f_val.x));\n int8[1] = static_cast(static_cast(f_val.y));\n return int16;\n#else\n val = __hmin2(val, make_bfloat162(127., 127.));\n val = __hmax2(val, make_bfloat162(-128., -128.));\n\n union\n {\n int8_t int8[2];\n int16_t int16;\n };\n\n int8[0] = static_cast(static_cast(val.x));\n int8[1] = static_cast(static_cast(val.y));\n return int16;\n#endif\n}\n#endif\n\n#if ENABLE_BF16\ntemplate <>\ninline __device__ __nv_bfloat162 ldg(const __nv_bfloat162 *val)\n{\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800\n return val[0];\n#else\n return __ldg(val);\n#endif\n}\n\ntemplate <>\ninline __device__ __nv_bfloat16 ldg(const __nv_bfloat16 *val)\n{\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800\n return val[0];\n#else\n return __ldg(val);\n#endif\n}\n#endif // ENABLE_BF16\n\ntemplate \n__device__ inline T_OUT cuda_cast(T_IN val)\n{\n return val;\n}\n\ntemplate <>\n__device__ inline float2 cuda_cast(int2 val)\n{\n return make_float2(val.x, val.y);\n}\n\ntemplate <>\n__device__ inline float2 cuda_cast(float val)\n{\n return make_float2(val, val);\n}\n\ntemplate <>\n__device__ inline float2 cuda_cast(half2 val)\n{\n return __half22float2(val);\n}\n\ntemplate <>\n__device__ inline half2 cuda_cast(float2 val)\n{\n return __float22half2_rn(val);\n}\n\ntemplate <>\n__device__ inline half2 cuda_cast(float val)\n{\n return __float2half2_rn(val);\n}\n\ntemplate <>\n__device__ inline half2 cuda_cast(half val)\n{\n return __half2half2(val);\n}\n\ntemplate <>\n__device__ inline int8_t cuda_cast(half val)\n{\n union\n {\n int8_t int8[2];\n int16_t int16;\n };\n\n union\n {\n half fp16;\n int16_t int16_in;\n };\n\n fp16 = val;\n asm volatile(\"cvt.rni.sat.s8.f16 %0, %1;\" : \"=h\"(int16) : \"h\"(int16_in));\n return int8[0];\n}\n\ntemplate <>\n__device__ inline int16_t cuda_cast(half2 val)\n{\n union\n {\n int8_t int8[2];\n int16_t int16;\n };\n\n int8[0] = cuda_cast(val.x);\n int8[1] = cuda_cast(val.y);\n return int16;\n}\n\ntemplate <>\n__device__ inline int8_t cuda_cast(float val)\n{\n union\n {\n int8_t int8[2];\n int16_t int16;\n };\n\n asm volatile(\"cvt.rni.sat.s8.f32 %0, %1;\" : \"=h\"(int16) : \"f\"(val));\n return int8[0];\n}\n\ntemplate <>\n__device__ inline int16_t cuda_cast(float2 val)\n{\n union\n {\n int8_t int8[2];\n int16_t int16;\n };\n\n int8[0] = cuda_cast(val.x);\n int8[1] = cuda_cast(val.y);\n return int16;\n}\n\ntemplate <>\n__device__ inline half2 cuda_cast(int16_t val)\n{\n union\n {\n int8_t int8[2];\n int16_t int16;\n };\n\n int16 = val;\n return make_half2(int8[0], int8[1]);\n}\n\ntemplate <>\n__device__ inline float2 cuda_cast(int16_t val)\n{\n union\n {\n int8_t int8[2];\n int16_t int16;\n };\n\n int16 = val;\n return make_float2(int8[0], int8[1]);\n}\n\n#ifdef ENABLE_BF16\ntemplate <>\n__device__ inline __nv_bfloat16 cuda_cast(int32_t val)\n{\n return static_cast(val);\n}\n\ntemplate <>\n__device__ inline __nv_bfloat16 cuda_cast(int8_t val)\n{\n return static_cast(val);\n}\n\ntemplate <>\n__device__ inline int8_t cuda_cast(__nv_bfloat16 val)\n{\n return static_cast(val);\n}\n\ntemplate <>\n__device__ inline float cuda_cast(__nv_bfloat16 val)\n{\n return __bfloat162float(val);\n}\n\ntemplate <>\n__device__ inline float2 cuda_cast(__nv_bfloat162 val)\n{\n return __bfloat1622float2(val);\n}\n\ntemplate <>\n__device__ inline half cuda_cast(__nv_bfloat16 val)\n{\n return __float2half(__bfloat162float(val));\n}\n\ntemplate <>\n__device__ inline int16_t cuda_cast(__nv_bfloat162 val)\n{\n return bf1622int16(val);\n}\n\ntemplate <>\n__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, float>(float val)\n{\n return __float2bfloat16(val);\n}\n\ntemplate <>\n__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, half>(half val)\n{\n return __float2bfloat16(__half2float(val));\n}\n\ntemplate <>\n__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, __nv_bfloat16>(__nv_bfloat16 val)\n{\n return __bfloat162bfloat162(val);\n}\n\ntemplate <>\n__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float>(float val)\n{\n return __float2bfloat162_rn(val);\n}\n\ntemplate <>\n__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float2>(float2 val)\n{\n return float22bf162(val);\n}\n\ntemplate <>\n__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, int16_t>(int16_t val)\n{\n union\n {\n int8_t int8[2];\n int16_t int16;\n };\n\n int16 = val;\n __nv_bfloat162 res;\n res.x = cuda_cast<__nv_bfloat16>(int8[0]);\n res.y = cuda_cast<__nv_bfloat16>(int8[1]);\n return res;\n}\n\ntemplate <>\n__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, half2>(half2 val)\n{\n return float22bf162(__half22float2(val));\n}\n\n#endif // ENABLE BF16\n\ntemplate \n__device__ inline To cuda_sum(Ti val)\n{\n return cuda_cast(val);\n};\n\ntemplate \n__device__ inline To cuda_sum(float2 val)\n{\n return cuda_cast(val.x + val.y);\n};\n\n// Unary maximum: compute the max of a vector type\ntemplate \n__device__ inline To cuda_max(Ti val)\n{\n return cuda_cast(val);\n};\n\ntemplate <>\n__device__ inline float cuda_max(float2 val)\n{\n return fmaxf(val.x, val.y);\n}\n\ntemplate <>\n__device__ inline half cuda_max(half2 val)\n{\n return __hmax(val.x, val.y);\n}\n\n#ifdef ENABLE_BF16\ntemplate <>\n__device__ inline __nv_bfloat16 cuda_max(__nv_bfloat162 val)\n{\n#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))\n return __hmax(val.x, val.y);\n#endif\n}\n#endif\n\n// Binary maximum: compute the max of two scalar types\ntemplate \n__device__ inline T cuda_max(T val1, T val2)\n{\n return (val1 > val2) ? val1 : val2;\n}\n\ntemplate \n__device__ inline T cuda_abs(T val)\n{\n assert(false);\n return {};\n}\n\ntemplate <>\n__device__ inline float cuda_abs(float val)\n{\n return fabs(val);\n}\n\ntemplate <>\n__device__ inline float2 cuda_abs(float2 val)\n{\n return make_float2(fabs(val.x), fabs(val.y));\n}\n\ntemplate <>\n__device__ inline half cuda_abs(half val)\n{\n return __habs(val);\n}\n\ntemplate <>\n__device__ inline half2 cuda_abs(half2 val)\n{\n return __habs2(val);\n}\n\n#ifdef ENABLE_BF16\n\n#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)\ntemplate <>\n__device__ inline __nv_bfloat16 cuda_abs(__nv_bfloat16 val)\n{\n return __habs(val);\n}\n\ntemplate <>\n__device__ inline __nv_bfloat162 cuda_abs(__nv_bfloat162 val)\n{\n return __habs2(val);\n}\n#endif\n\n#endif\n"
},
{
"path": "deepcompressor/backend/tinychat/linear.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"TinyChat Quantized Linear Module\"\"\"\n\nimport warnings\n\nimport torch\nimport torch.nn as nn\n\nfrom .csrc.load import _C\nfrom .utils import ceil_num_groups, convert_to_tinychat_w4x16y16_linear_weight\n\n__all__ = [\"W4Linear\"]\n\nwarnings.warn(\n \"Module `tinychat.linear` will be moved to `Nunchaku` and deprecated in the future release.\",\n DeprecationWarning,\n stacklevel=2,\n)\n\n\nclass W4Linear(nn.Module):\n def __init__(\n self,\n in_features: int,\n out_features: int,\n bias: bool = False,\n group_size: int = 128,\n dtype: torch.dtype = torch.float16,\n device: str | torch.device = \"cuda\",\n ):\n super().__init__()\n assert dtype in (torch.float16, torch.bfloat16), f\"Unsupported dtype: {dtype}\"\n\n self.in_features = in_features\n self.out_features = out_features\n self.group_size = group_size if group_size != -1 else in_features\n assert self.in_features % self.group_size == 0\n assert out_features % (32 // self.weight_bits) == 0\n self.ceil_num_groups = ceil_num_groups(\n in_features=self.in_features,\n group_size=self.group_size,\n weight_bits=self.weight_bits,\n )\n\n assert out_features % (self.interleave) == 0\n self.register_buffer(\n \"qweight\",\n torch.zeros(\n (\n self.out_features // self.interleave,\n self.in_features // (16 // self.weight_bits) * self.interleave,\n ),\n dtype=torch.int16,\n device=device,\n ),\n )\n self.register_buffer(\n \"scales\",\n torch.zeros((self.ceil_num_groups, self.out_features), dtype=dtype, device=device),\n )\n self.register_buffer(\n \"scaled_zeros\",\n torch.zeros((self.ceil_num_groups, self.out_features), dtype=dtype, device=device),\n )\n if bias:\n self.register_buffer(\"bias\", torch.zeros((out_features), dtype=dtype, device=device))\n else:\n self.bias = None\n\n @property\n def weight_bits(self) -> int:\n return 4\n\n @property\n def interleave(self) -> int:\n return 4\n\n @torch.no_grad()\n def forward(self, x):\n if x.numel() / x.shape[-1] < 8:\n out = _C.awq_gemv_forward_cuda(\n x,\n self.qweight,\n self.scales,\n self.scaled_zeros,\n x.numel() // x.shape[-1],\n self.out_features,\n self.in_features,\n self.group_size,\n )\n else:\n out = _C.awq_gemm_forward_cuda(x, self.qweight, self.scales, self.scaled_zeros)\n out = out + self.bias if self.bias is not None else out\n return out\n\n @staticmethod\n def from_linear(\n linear: nn.Linear,\n group_size: int,\n init_only: bool = False,\n weight: torch.Tensor | None = None,\n scale: torch.Tensor | None = None,\n zero: torch.Tensor | None = None,\n zero_pre_scaled: bool = False,\n ) -> \"W4Linear\":\n \"\"\"Convert a linear layer to a TinyChat 4-bit weight-only quantized linear layer.\n\n Args:\n linear (`nn.Linear`):\n linear layer to be converted.\n group_size (`int`):\n quantization group size.\n init_only (`bool`, *optional*, defaults to `False`):\n whether to only initialize the quantized linear layer.\n weight (`torch.Tensor`, *optional*, defaults to `None`):\n weight tensor for the quantized linear layer.\n scale (`torch.Tensor`, *optional*, defaults to `None`):\n scale tensor for the quantized linear layer.\n zero (`torch.Tensor`, *optional*, defaults to `None`):\n zero point tensor for the quantized linear layer.\n zero_pre_scaled (`bool`, *optional*, defaults to `False`):\n whether zero point tensor is pre-scaled.\n\n Returns:\n `W4Linear`:\n quantized linear layer.\n \"\"\"\n assert isinstance(linear, nn.Linear)\n weight = linear.weight.data if weight is None else weight.data\n dtype, device = weight.dtype, weight.device\n oc, ic = linear.out_features, linear.in_features\n _linear = W4Linear(\n in_features=ic,\n out_features=oc,\n bias=linear.bias is not None,\n group_size=group_size,\n dtype=dtype,\n device=device,\n )\n if init_only:\n return _linear\n if linear.bias is not None:\n _linear.bias.data.copy_(linear.bias.data)\n if scale is None:\n assert zero is None, \"scale and zero point tensors should be provided together.\"\n group_size = ic if group_size <= 0 else group_size\n assert group_size <= ic, \"group size should be less than or equal to input channel size.\"\n assert ic % group_size == 0, \"input channel size should be divisible by group size.\"\n ng, gs = ic // group_size, group_size\n weight = weight.to(dtype=torch.float32).view(oc, 1, ng, gs)\n vmin, vmax = weight.amin(dim=-1, keepdim=True), weight.amax(dim=-1, keepdim=True)\n scale = (vmax - vmin).div_(15)\n scale[scale == 0] = 1.0\n if zero_pre_scaled:\n zero = vmin.neg_().div_(scale).round_().clamp_(0, 15)\n weight = weight.div_(scale).add_(zero).round_().clamp_(0, 15).sub_(zero).mul_(scale)\n else:\n zero = vmin.neg_().clamp_min(0)\n weight = weight.add_(zero).div_(scale).round_().clamp_(0, 15).mul_(scale).sub_(zero)\n weight = weight.to(dtype=dtype).view(oc, ic)\n scale = scale.to(dtype=dtype)\n zero = zero.to(dtype=dtype)\n weight, scale, zero = convert_to_tinychat_w4x16y16_linear_weight(\n weight=weight,\n scale=scale,\n zero=zero,\n zero_pre_scaled=zero_pre_scaled,\n )\n _linear.qweight.data.copy_(weight)\n _linear.scales.data.copy_(scale)\n _linear.scaled_zeros.data.copy_(zero)\n return _linear\n\n def extra_repr(self) -> str:\n return \"in_features={}, out_features={}, bias={}, weight_bits={}, group_size={}\".format(\n self.in_features,\n self.out_features,\n self.bias is not None,\n self.weight_bits,\n self.group_size,\n )\n"
},
{
"path": "deepcompressor/backend/tinychat/utils.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"TinyChat backend utilities.\"\"\"\n\nimport torch\n\nfrom ..utils import ceil_divide\n\n__all__ = [\"ceil_num_groups\", \"convert_to_tinychat_w4x16y16_linear_weight\"]\n\n\ndef ceil_num_groups(in_features: int, group_size: int, weight_bits: int = 4) -> int:\n \"\"\"Calculate the ceiling number of quantization groups.\n\n Args:\n in_features (`int`):\n input channel size.\n group_size (`int`):\n quantization group size.\n weight_bits (`int`, *optional*, defaults to `4`):\n quantized weight bits.\n\n Returns:\n `int`:\n ceiling number of quantization groups.\n \"\"\"\n assert in_features % group_size == 0, \"input channel size should be divisible by group size.\"\n num_groups = in_features // group_size\n assert weight_bits in (4, 2, 1), \"weight bits should be 4, 2, or 1.\"\n pack_size = 32 // weight_bits # one INT32 contains `pack_size` elements of weights\n num_packs = ceil_divide(num_groups, pack_size)\n if group_size >= 128:\n num_packs_factor = 1\n elif group_size == 64:\n num_packs_factor = 2\n elif group_size == 32:\n num_packs_factor = 4\n else:\n raise NotImplementedError\n # make sure num_packs is a multiple of num_packs_factor\n num_packs = ceil_divide(num_packs, num_packs_factor) * num_packs_factor\n num_groups = num_packs * pack_size\n return num_groups\n\n\ndef pack_w4(weight: torch.Tensor) -> torch.Tensor:\n assert weight.dtype == torch.int32, f\"quantized weight should be torch.int32, but got {weight.dtype}.\"\n oc, ic = weight.shape\n assert ic % 32 == 0, \"input channel size should be divisible by 32.\"\n # [0, 1, ..., 31] -> [0, 8, 16, 24, 1, 9, 17, 25, ..., 7, 15, 23, 31]\n weight = weight.view(-1, 4, 8)\n weight = weight[:, 0] | (weight[:, 1] << 4) | (weight[:, 2] << 8) | (weight[:, 3] << 12)\n weight = weight.view(oc // 4, 4, ic // 64, 16).permute(0, 2, 1, 3).reshape(oc // 4, ic)\n return weight.to(torch.int16)\n\n\ndef convert_to_tinychat_w4x16y16_linear_weight(\n weight: torch.Tensor,\n scale: torch.Tensor,\n zero: torch.Tensor,\n zero_pre_scaled: bool = False,\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"Convert a weight tensor to TinyChat W4-X16-Y16 linear weight format.\n\n Args:\n weight (`torch.Tensor`):\n weight tensor to be converted.\n scale (`torch.Tensor`):\n scale tensor for the weight tensor.\n zero (`torch.Tensor`):\n zero point tensor for the weight tensor.\n zero_pre_scaled (`bool`, *optional*, defaults to `False`):\n whether zero point tensor is pre-scaled.\n\n Returns:\n `tuple[torch.Tensor, torch.Tensor, torch.Tensor]`:\n packed quantized weight tensor, scale tensor, and zero point tensor.\n \"\"\"\n dtype, device = weight.dtype, weight.device\n assert dtype in (torch.float16, torch.bfloat16), \"currently tinychat only supports fp16 and bf16.\"\n assert scale is not None, \"scale tensor is required for quantization.\"\n assert zero is not None, \"zero point tensor is required for quantization.\"\n weight = weight.to(dtype=torch.float32)\n scale = scale.to(dtype=torch.float32, device=device)\n zero = zero.to(dtype=torch.float32, device=device)\n if zero_pre_scaled:\n zero = zero * scale\n oc, ic = weight.shape\n if scale.numel() == 1:\n scale = scale.view(1, 1).expand(oc, 1)\n ng, gs = 1, ic\n else:\n ng = scale.numel() // oc\n gs = ic // ng\n scale = scale.reshape(oc, ng).contiguous().view(oc, ng, 1)\n assert ic == gs * ng, \"input channel size should be equal to group size times number of groups.\"\n if zero.numel() == 1:\n zero = zero.view(1, 1).expand(oc, ng)\n zero = zero.reshape(oc, ng).contiguous().view(oc, ng, 1)\n weight = weight.view(oc, ng, -1).add_(zero).div_(scale).round_().view(oc, ic)\n assert weight.min() >= 0 and weight.max() <= 15, \"quantized weight should be in [0, 15].\"\n _weight = pack_w4(weight.to(torch.int32))\n _ng = ceil_num_groups(ic, gs, weight_bits=4)\n _scale = torch.zeros((_ng, oc), dtype=dtype, device=device)\n _zero = torch.zeros((_ng, oc), dtype=dtype, device=device)\n _scale[:ng] = scale.view(oc, ng).t().to(dtype=dtype)\n _zero[:ng] = zero.view(oc, ng).t().to(dtype=dtype).neg_()\n return _weight, _scale, _zero\n"
},
{
"path": "deepcompressor/backend/utils.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Backend utilities.\"\"\"\n\nimport typing as tp\n\nimport safetensors\nimport torch\n\n__all__ = [\"ceil_divide\", \"pad\", \"fp_quantize\", \"MmaWeightPackerBase\"]\n\n\ndef ceil_divide(x: int, divisor: int) -> int:\n \"\"\"Ceiling division.\n\n Args:\n x (`int`):\n dividend.\n divisor (`int`):\n divisor.\n\n Returns:\n `int`:\n ceiling division result.\n \"\"\"\n return (x + divisor - 1) // divisor\n\n\ndef pad(\n tensor: tp.Optional[torch.Tensor],\n divisor: int | tp.Sequence[int],\n dim: int | tp.Sequence[int],\n fill_value: float | int = 0,\n) -> torch.Tensor:\n if isinstance(divisor, int):\n if divisor <= 1:\n return tensor\n elif all(d <= 1 for d in divisor):\n return tensor\n if tensor is None:\n return None\n shape = list(tensor.shape)\n if isinstance(dim, int):\n assert isinstance(divisor, int)\n shape[dim] = ceil_divide(shape[dim], divisor) * divisor\n else:\n if isinstance(divisor, int):\n divisor = [divisor] * len(dim)\n for d, div in zip(dim, divisor, strict=True):\n shape[d] = ceil_divide(shape[d], div) * div\n result = torch.full(shape, fill_value, dtype=tensor.dtype, device=tensor.device)\n result[[slice(0, extent) for extent in tensor.shape]] = tensor\n return result\n\n\ndef load_state_dict_in_safetensors(\n path: str, device: str | torch.device = \"cpu\", filter_prefix: str = \"\"\n) -> dict[str, torch.Tensor]:\n \"\"\"Load state dict in SafeTensors.\n\n Args:\n path (`str`):\n file path.\n device (`str` | `torch.device`, optional, defaults to `\"cpu\"`):\n device.\n filter_prefix (`str`, optional, defaults to `\"\"`):\n filter prefix.\n\n Returns:\n `dict`:\n loaded SafeTensors.\n \"\"\"\n state_dict = {}\n with safetensors.safe_open(path, framework=\"pt\", device=device) as f:\n for k in f.keys():\n if filter_prefix and not k.startswith(filter_prefix):\n continue\n state_dict[k.removeprefix(filter_prefix)] = f.get_tensor(k)\n return state_dict\n\n\ndef fp_quantize(x: torch.Tensor, codebook: torch.Tensor | None = None) -> torch.Tensor:\n if codebook is None:\n codebook = torch.tensor(\n [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0],\n dtype=x.dtype,\n device=x.device,\n )\n return (x.unsqueeze(-1) - codebook.unsqueeze(0)).abs().argmin(dim=-1)\n\n\nclass MmaWeightPackerBase:\n def __init__(self, bits: int, warp_n: int, comp_n: int = None, comp_k: int = None):\n self.bits = bits\n assert self.bits in (1, 4, 8, 16, 32), \"weight bits should be 1, 4, 8, 16, or 32.\"\n\n # region compute tile size\n self.comp_n = comp_n if comp_n is not None else 16\n \"\"\"smallest tile size in `n` dimension for MMA computation.\"\"\"\n self.comp_k = comp_k if comp_k is not None else 256 // self.bits\n \"\"\"smallest tile size in `k` dimension for MMA computation.\"\"\"\n # the smallest MMA computation may contain several MMA instructions\n self.insn_n = 8 # mma instruction tile size in `n` dimension\n \"\"\"tile size in `n` dimension for MMA instruction.\"\"\"\n self.insn_k = self.comp_k\n \"\"\"tile size in `k` dimension for MMA instruction.\"\"\"\n assert self.insn_k * self.bits in (128, 256), (\n f\"insn_k ({self.insn_k}) * bits ({self.bits}) should be 128 or 256.\"\n )\n assert self.comp_n % self.insn_n == 0, f\"comp_n ({self.comp_n}) should be divisible by insn_n ({self.insn_n}).\"\n self.num_lanes = 32\n \"\"\"there are 32 lanes (or threds) in a warp.\"\"\"\n self.num_k_lanes = 4\n self.num_n_lanes = 8\n assert warp_n >= self.comp_n and warp_n % self.comp_n == 0, (\n f\"warp_n ({warp_n}) should be divisible by comp_n({self.comp_n}).\"\n )\n self.warp_n = warp_n\n # endregion\n # region memory\n self.reg_k = 32 // self.bits\n \"\"\"number of elements in a register in `k` dimension.\"\"\"\n self.reg_n = 1\n \"\"\"number of elements in a register in `n` dimension (always 1).\"\"\"\n self.k_pack_size = self.comp_k // (self.num_k_lanes * self.reg_k)\n \"\"\"number of elements in a pack in `k` dimension.\"\"\"\n self.n_pack_size = self.comp_n // (self.num_n_lanes * self.reg_n)\n \"\"\"number of elements in a pack in `n` dimension.\"\"\"\n self.pack_size = self.k_pack_size * self.n_pack_size\n \"\"\"number of elements in a pack accessed by a lane at a time.\"\"\"\n assert 1 <= self.pack_size <= 4, \"pack size should be less than or equal to 4.\"\n assert self.k_pack_size * self.num_k_lanes * self.reg_k == self.comp_k\n assert self.n_pack_size * self.num_n_lanes * self.reg_n == self.comp_n\n self.mem_k = self.comp_k\n \"\"\"the tile size in `k` dimension for one tensor memory access.\"\"\"\n self.mem_n = warp_n\n \"\"\"the tile size in `n` dimension for one tensor memory access.\"\"\"\n self.num_k_packs = self.mem_k // (self.k_pack_size * self.num_k_lanes * self.reg_k)\n \"\"\"number of packs in `k` dimension for one tensor memory access.\"\"\"\n self.num_n_packs = self.mem_n // (self.n_pack_size * self.num_n_lanes * self.reg_n)\n \"\"\"number of packs in `n` dimension for one tensor memory access.\"\"\"\n # endregion\n\n def get_view_shape(self, n: int, k: int) -> tuple[int, int, int, int, int, int, int, int, int, int]:\n assert n % self.mem_n == 0, \"output channel size should be divisible by mem_n.\"\n assert k % self.mem_k == 0, \"input channel size should be divisible by mem_k.\"\n return (\n n // self.mem_n,\n self.num_n_packs,\n self.n_pack_size,\n self.num_n_lanes,\n self.reg_n,\n k // self.mem_k,\n self.num_k_packs,\n self.k_pack_size,\n self.num_k_lanes,\n self.reg_k,\n )\n"
},
{
"path": "deepcompressor/calib/__init__.py",
"content": "# -*- coding: utf-8 -*-\n"
},
{
"path": "deepcompressor/calib/config/__init__.py",
"content": "# -*- coding: utf-8 -*-\n\nfrom .lowrank import QuantLowRankCalibConfig, SkipBasedQuantLowRankCalibConfig\nfrom .range import DynamicRangeCalibConfig, SkipBasedDynamicRangeCalibConfig\nfrom .reorder import ChannelOrderCalibConfig, SkipBasedChannelOrderConfig\nfrom .rotation import QuantRotationConfig\nfrom .search import (\n SearchBasedCalibConfig,\n SearchBasedCalibGranularity,\n SearchBasedCalibObjective,\n SearchBasedCalibStrategy,\n)\nfrom .smooth import SkipBasedSmoothCalibConfig, SmoothCalibConfig, SmoothSpanMode, SmoothTransfomerConfig\n"
},
{
"path": "deepcompressor/calib/config/lowrank.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Quantization SVD calibration configuration.\"\"\"\n\nfrom dataclasses import dataclass, field\n\nfrom omniconfig import configclass\n\nfrom ...quantizer.config import QuantLowRankConfig\nfrom ...utils.common import num2str\nfrom ...utils.config import SkipBasedConfig\nfrom .search import SearchBasedCalibConfig, SearchBasedCalibGranularity, SearchBasedCalibStrategy\n\n__all__ = [\"QuantLowRankCalibConfig\", \"SkipBasedQuantLowRankCalibConfig\"]\n\n\n@configclass\n@dataclass\nclass QuantLowRankCalibConfig(SearchBasedCalibConfig, QuantLowRankConfig):\n \"\"\"Configuration for quantization low-rank branch calibration.\n\n Args:\n rank (`int`, *optional*, defaults to `32`):\n The rank of the low-rank branch.\n exclusive (`bool`, *optional*, defaults to `False`):\n Whether to use exclusive low-rank branch for each weight sharing the inputs.\n compensate (`bool`, *optional*, defaults to `False`):\n Whether the low-rank branch compensates the quantization error.\n degree (`int`, *optional*, default=`2`):\n The power degree for the quantization error. Defaults to `2`.\n objective (`SearchBasedCalibObjective`, *optional*, default=`SearchBasedCalibObjective.OutputsError`):\n The objective for quantization calibration.\n sample_batch_size (`int`, *optional*, default=`-1`):\n The samples batch size for calibration.\n sample_size (`int`, *optional*, default=`-1`):\n The calibration sample size.\n outputs_device (`str`, *optional*, default=`\"cpu\"`):\n The device to store the precomputed outputs of the module.\n num_iters (`int`, *optional*, default=`1`):\n The number of iterations.\n early_stop (`bool`, *optional*, default=`False`):\n Whether to stop the calibration early.\n \"\"\"\n\n granularity: SearchBasedCalibGranularity = field(init=False, default=SearchBasedCalibGranularity.Layer)\n element_batch_size: int = field(init=False, default=-1)\n element_size: int = field(init=False, default=-1)\n pre_reshape: bool = field(init=False, default=True)\n num_iters: int = 1\n early_stop: bool = False\n\n def __post_init__(self):\n if self.strategy != SearchBasedCalibStrategy.Manual:\n self.strategy = SearchBasedCalibStrategy.GridSearch\n if self.compensate and self.num_iters <= 1:\n self.exclusive = True\n super().__post_init__()\n\n def generate_dirnames(self, *, prefix: str = \"\", **kwargs) -> list[str]:\n \"\"\"Generate the directory names of the configuration.\n\n Returns:\n list[str]: The directory names.\n \"\"\"\n names = super().generate_dirnames(**kwargs)\n name = f\"i{num2str(self.num_iters)}.r{num2str(self.rank)}\"\n if self.exclusive:\n name += \".exclusive\"\n if self.compensate:\n name += \".compensate\"\n if self.early_stop and self.num_iters > 1:\n name += \".earlystop\"\n names.append(name)\n if prefix:\n names = [f\"{prefix}.{name}\" for name in names]\n return names\n\n\n@configclass\n@dataclass\nclass SkipBasedQuantLowRankCalibConfig(SkipBasedConfig, QuantLowRankCalibConfig):\n \"\"\"Configuration for Quantization Low-Rank Branch calibration.\n\n Args:\n rank (`int`, *optional*, defaults to `32`):\n The rank of the low-rank branch.\n exclusive (`bool`, *optional*, defaults to `False`):\n Whether to use exclusive low-rank branch for each weight sharing the inputs.\n compensate (`bool`, *optional*, defaults to `False`):\n Whether the low-rank branch compensates the quantization error.\n degree (`int`, *optional*, default=`2`):\n The power degree for the quantization error. Defaults to `2`.\n objective (`SearchBasedCalibObjective`, *optional*, default=`SearchBasedCalibObjective.OutputsError`):\n The objective for quantization calibration.\n sample_batch_size (`int`, *optional*, default=`-1`):\n The samples batch size for calibration.\n sample_size (`int`, *optional*, default=`-1`):\n The calibration sample size.\n outputs_device (`str`, *optional*, default=`\"cpu\"`):\n The device to store the precomputed outputs of the module.\n num_iters (`int`, *optional*, default=`1`):\n The number of iterations.\n early_stop (`bool`, *optional*, default=`False`):\n Whether to stop the calibration early.\n skips (`list[str]`, *optional*, default=`[]`):\n The keys of the modules to skip.\n \"\"\"\n\n pass\n"
},
{
"path": "deepcompressor/calib/config/range.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Quantization dynamic range calibration configuration.\"\"\"\n\nfrom dataclasses import dataclass\n\nfrom omniconfig import configclass\n\nfrom ...utils.common import num2str\nfrom ...utils.config import SkipBasedConfig\nfrom .search import SearchBasedCalibConfig, SearchBasedCalibStrategy\n\n__all__ = [\"DynamicRangeCalibConfig\", \"SkipBasedDynamicRangeCalibConfig\"]\n\n\n@configclass\n@dataclass\nclass DynamicRangeCalibConfig(SearchBasedCalibConfig):\n \"\"\"Configuration for quantization dynamic range calibration.\n\n Args:\n degree (`int`, *optional*, default=`2`):\n The power degree for the quantization error. Defaults to `2`.\n objective (`SearchBasedCalibObjective`, *optional*, default=`SearchBasedCalibObjective.OutputsError`):\n The objective for quantization calibration.\n strategy (`SearchBasedCalibStrategy`, *optional*, default=`SearchBasedCalibStrategy.Manual`):\n The strategy for quantization calibration.\n granularity (`SearchBasedCalibGranularity`, *optional*, default=`SearchBasedCalibGranularity.Layer`):\n The granularity for quantization calibration.\n element_batch_size (`int`, *optional*, default=`-1`):\n The element batch size for calibration.\n sample_batch_size (`int`, *optional*, default=`-1`):\n The samples batch size for calibration.\n element_size (`int`, *optional*, default=`-1`):\n The calibration element size.\n sample_size (`int`, *optional*, default=`-1`):\n The calibration sample size.\n pre_reshape (`bool`, *optional*, default=`True`):\n Whether to enable reshaping the tensor before calibration.\n outputs_device (`str`, *optional*, default=`\"cpu\"`):\n The device to store the precomputed outputs of the module.\n ratio (`float`, *optional*, default=`1.0`):\n The dynamic range ratio.\n max_shrink (`float`, *optional*, default=`0.2`):\n Maximum shrinkage ratio.\n max_expand (`float`, *optional*, default=`1.0`):\n Maximum expansion ratio.\n num_grids (`int`, *optional*, default=`80`):\n Number of grids for linear range search.\n allow_scale (`bool`, *optional*, default=`False`):\n Whether to allow range dynamic scaling.\n \"\"\"\n\n ratio: float = 1.0\n max_shrink: float = 0.2\n max_expand: float = 1.0\n num_grids: int = 80\n allow_scale: bool = False\n\n def get_linear_ratios(self) -> list[float]:\n \"\"\"Get the ratios for linear range search.\n\n Returns:\n `list[float]`:\n The dynamic range ratio candidates for linear range search.\n \"\"\"\n num_grids, max_shrink, max_expand = self.num_grids, self.max_shrink, self.max_expand\n assert max_shrink < 1, \"maximal shrinkage ratio must be less than 1\"\n ratios = [1 - grid / num_grids * (1 - max_shrink) for grid in range(1, num_grids + 1)]\n if max_expand > 1:\n ratios += [1 + grid / num_grids * (max_expand - 1) for grid in range(1, num_grids + 1)]\n return ratios\n\n def get_ratios(self) -> list[list[float]]:\n \"\"\"Get the ratios for linear range search.\n\n Returns:\n `list[list[float]]`:\n The dynamic range ratio candidates for linear range search.\n \"\"\"\n if self.strategy == SearchBasedCalibStrategy.Manual:\n return [[self.ratio]]\n elif self.strategy == SearchBasedCalibStrategy.GridSearch:\n return [[1.0], self.get_linear_ratios()]\n else:\n raise ValueError(f\"Invalid strategy: {self.strategy}\")\n\n def generate_dirnames(self, *, prefix: str = \"\", **kwargs) -> list[str]:\n \"\"\"Generate the directory names of the configuration.\n\n Args:\n prefix (`str`, *optional*, default=`\"\"`):\n The prefix of the directory.\n\n Returns:\n `list[str]`:\n The directory names.\n \"\"\"\n names = super().generate_dirnames(**kwargs)\n if self.strategy == SearchBasedCalibStrategy.Manual:\n name = f\"r.[{num2str(self.ratio)}]\"\n elif self.strategy == SearchBasedCalibStrategy.GridSearch:\n name = f\"r.[{num2str(self.max_shrink)}.{num2str(self.max_expand)}].g{self.num_grids}\"\n else:\n raise ValueError(f\"Invalid strategy: {self.strategy}\")\n if self.allow_scale:\n name += \".scale\"\n names.append(name)\n if prefix:\n names = [f\"{prefix}.{name}\" for name in names]\n return names\n\n\n@configclass\n@dataclass\nclass SkipBasedDynamicRangeCalibConfig(SkipBasedConfig, DynamicRangeCalibConfig):\n \"\"\"Configuration for quantization dynamic range calibration.\n\n Args:\n degree (`int`, *optional*, default=`2`):\n The power degree for the quantization error. Defaults to `2`.\n objective (`SearchBasedCalibObjective`, *optional*, default=`SearchBasedCalibObjective.OutputsError`):\n The objective for quantization calibration.\n strategy (`SearchBasedCalibStrategy`, *optional*, default=`SearchBasedCalibStrategy.Manual`):\n The strategy for quantization calibration.\n granularity (`SearchBasedCalibGranularity`, *optional*, default=`SearchBasedCalibGranularity.Layer`):\n The granularity for quantization calibration.\n element_batch_size (`int`, *optional*, default=`-1`):\n The element batch size for calibration.\n sample_batch_size (`int`, *optional*, default=`-1`):\n The samples batch size for calibration.\n element_size (`int`, *optional*, default=`-1`):\n The calibration element size.\n sample_size (`int`, *optional*, default=`-1`):\n The calibration sample size.\n pre_reshape (`bool`, *optional*, default=`True`):\n Whether to enable reshaping the tensor before calibration.\n outputs_device (`str`, *optional*, default=`\"cpu\"`):\n The device to store the precomputed outputs of the module.\n ratio (`float`, *optional*, default=`1.0`):\n The dynamic range ratio.\n max_shrink (`float`, *optional*, default=`0.2`):\n Maximum shrinkage ratio.\n max_expand (`float`, *optional*, default=`1.0`):\n Maximum expansion ratio.\n num_grids (`int`, *optional*, default=`80`):\n Number of grids for linear range search.\n allow_scale (`bool`, *optional*, default=`False`):\n Whether to allow range dynamic scaling.\n skips (`list[str]`, *optional*, default=`[]`):\n The keys of the modules to skip.\n \"\"\"\n\n pass\n"
},
{
"path": "deepcompressor/calib/config/reorder.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Channel reorder configuration.\"\"\"\n\nimport enum\nfrom dataclasses import dataclass, field\n\nfrom omniconfig import configclass\n\nfrom ...utils.config import SkipBasedConfig\nfrom .search import (\n SearchBasedCalibConfig,\n SearchBasedCalibGranularity,\n SearchBasedCalibObjective,\n SearchBasedCalibStrategy,\n)\n\n__all__ = [\"ChannelOrderCalibConfig\", \"SkipBasedChannelOrderConfig\"]\n\n\n@configclass\n@dataclass\nclass ChannelOrderCalibConfig(SearchBasedCalibConfig):\n \"\"\"Configuration for channel order calibration in group quantization.\n\n Args:\n degree (`int`, *optional*, default=`2`):\n The power degree for the quantization error. Defaults to `2`.\n strategy (`SearchBasedCalibStrategy`, *optional*, default=`SearchBasedCalibStrategy.Manual`):\n The strategy for quantization calibration.\n sample_batch_size (`int`, *optional*, default=`-1`):\n The samples batch size for calibration.\n sample_size (`int`, *optional*, default=`-1`):\n The calibration sample size.\n allow_x_quant (`bool`, *optional*, default=`True`):\n Whether to allow input quantization during calibration.\n allow_w_quant (`bool`, *optional*, default=`True`):\n Whether to allow weight quantization during calibration.\n channel_metric (`ChannelMetricMode`, *optional*, default=`ChannelMetricMode.AbsNormalizedMean`):\n The mode for computing the channel importance.\n channel_index (`ChannelIndexMode`, *optional*, default=`ChannelIndexMode.Sequential`):\n The mode for ranking the channel importance.\n dynamic (`bool`, *optional*, default=`False`):\n Whether to enable dynamic channel reorder.\n \"\"\"\n\n class ChannelMetric(enum.Enum):\n \"\"\"The mode for computing the channel importance.\"\"\"\n\n InputsAbsMax = \"xMax\"\n InputsAbsMean = \"xAvg\"\n InputsRootMeanSquare = \"xRms\"\n WeightsAbsMax = \"wMax\"\n WeightsAbsMean = \"wAvg\"\n WeightsRootMeanSquare = \"wRms\"\n AbsMaxProduct = \"pMax\"\n AbsMeanProduct = \"pAvg\"\n RootMeanSquareProduct = \"pRms\"\n\n class ChannelIndex(enum.Enum):\n \"\"\"The mode for ranking the channel importance.\"\"\"\n\n Sequential = \"Seq\"\n Transpose = \"Trp\"\n\n objective: SearchBasedCalibObjective = field(init=False, default=SearchBasedCalibObjective.OutputsError)\n granularity: SearchBasedCalibGranularity = field(init=False, default=SearchBasedCalibGranularity.Layer)\n element_batch_size: int = field(init=False, default=-1)\n element_size: int = field(init=False, default=-1)\n pre_reshape: bool = field(init=False, default=True)\n allow_x_quant: bool = True\n allow_w_quant: bool = True\n channel_metric: ChannelMetric = ChannelMetric.InputsAbsMax\n channel_index: ChannelIndex = ChannelIndex.Sequential\n dynamic: bool = False\n\n def __post_init__(self) -> None:\n if self.strategy != SearchBasedCalibStrategy.Manual:\n self.strategy = SearchBasedCalibStrategy.GridSearch\n super().__post_init__()\n\n def generate_dirnames(self, *, prefix: str = \"\", **kwargs) -> list[str]:\n \"\"\"Generate the directory names of the configuration.\n\n Args:\n prefix (`str`, *optional*, default=`\"\"`):\n The prefix of the directory.\n\n Returns:\n `list[str]`:\n The directory names.\n \"\"\"\n names = super().generate_dirnames(**kwargs)\n if self.strategy == SearchBasedCalibStrategy.Manual:\n name = f\"{self.channel_metric.name}.{self.channel_index.name}\"\n else:\n name = \"search\"\n if self.dynamic:\n name += \".dynamic\"\n names.append(name)\n disallows = []\n if not self.allow_x_quant:\n disallows.append(\"x\")\n if not self.allow_w_quant:\n disallows.append(\"w\")\n if disallows:\n names.append(f\"disallow.[{'+'.join(disallows)}]\")\n if prefix:\n names = [f\"{prefix}.{name}\" for name in names]\n return names\n\n\n@configclass\n@dataclass\nclass SkipBasedChannelOrderConfig(SkipBasedConfig, ChannelOrderCalibConfig):\n \"\"\"Configuration for channel order calibration in group quantization.\n\n Args:\n degree (`int`, *optional*, default=`2`):\n The power degree for the quantization error. Defaults to `2`.\n strategy (`SearchBasedCalibStrategy`, *optional*, default=`SearchBasedCalibStrategy.Manual`):\n The strategy for quantization calibration.\n sample_batch_size (`int`, *optional*, default=`-1`):\n The samples batch size for calibration.\n sample_size (`int`, *optional*, default=`-1`):\n The calibration sample size.\n allow_x_quant (`bool`, *optional*, default=`True`):\n Whether to allow input quantization during calibration.\n allow_w_quant (`bool`, *optional*, default=`True`):\n Whether to allow weight quantization during calibration.\n channel_metric (`ChannelMetricMode`, *optional*, default=`ChannelMetricMode.AbsNormalizedMean`):\n The mode for computing the channel importance.\n channel_index (`ChannelIndexMode`, *optional*, default=`ChannelIndexMode.Sequential`):\n The mode for ranking the channel importance.\n dynamic (`bool`, *optional*, default=`False`):\n Whether to enable dynamic channel reorder.\n skips (`list[str]`, *optional*, default=`[]`):\n The keys of the modules to skip.\n \"\"\"\n\n pass\n"
},
{
"path": "deepcompressor/calib/config/rotation.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Quantization Rotation configuration.\"\"\"\n\nimport os\nimport typing as tp\nfrom dataclasses import dataclass, field\n\nimport omniconfig\nfrom omniconfig import configclass\n\n__all__ = [\"QuantRotationConfig\"]\n\n\n@configclass\n@dataclass\nclass QuantRotationConfig:\n \"\"\"Configuration for rotation quantization.\n\n Args:\n name (`str`):\n The name of the rotation quantization configuration. If `path` is provided, this is required.\n Otherwise, it is set to \"random\" if `random` is `True`, and \"hadamard\" otherwise.\n path (`str`, *optional*, default=`\"\"`):\n The path to the rotation matrix. If provided, `name` must be set.\n random (`bool`, *optional*, default=`False`):\n Whether to use random hadamard sample as rotation matrix.\n transforms (`list[str]`, *optional*, default=`[]`):\n The module keys using explicit hadamard transform.\n \"\"\"\n\n name: str = \"\"\n path: str = \"\"\n random: bool = False\n transforms: list[str] = field(default_factory=list)\n\n def __post_init__(self) -> None:\n self.transforms = sorted(set(self.transforms or []))\n if self.path and os.path.exists(self.path):\n assert self.name, \"The name of the rotation quantization configuration must be provided.\"\n self.random = False\n else:\n self.path = \"\"\n self.name = \"random\" if self.random else \"hadamard\"\n\n def generate_dirnames(self, *, prefix: str = \"\", **kwargs) -> list[str]:\n \"\"\"Get the directory names of the rotation quantization configuration.\n\n Returns:\n list[str]: The directory names of the rotation quantization configuration.\n \"\"\"\n name = self.name\n if self.transforms:\n name += f\".[{'+'.join(self.transforms)}]\"\n return [f\"{prefix}.{name}\" if prefix else name]\n\n @classmethod\n def update_get_arguments(\n cls: type[\"QuantRotationConfig\"],\n *,\n overwrites: dict[str, tp.Callable[[omniconfig.Arguments], None] | None] | None = None,\n defaults: dict[str, tp.Any] | None = None,\n ) -> tuple[dict[str, tp.Callable[[omniconfig.Arguments], None] | None], dict[str, tp.Any]]:\n \"\"\"Get the arguments for the rotation quantization configuration.\"\"\"\n overwrites = overwrites or {}\n defaults = defaults or {}\n\n collect_fn = omniconfig.ADD_PREFIX_BOOL_FIELDS(\"transform\", **defaults)\n\n def add_transforms_argument(parser):\n collect_fn(parser)\n parser.add_argument(\"--transforms\", nargs=\"+\", default=[], help=\"The keys of the modules to transform.\")\n\n overwrites.setdefault(\"transforms\", add_transforms_argument)\n return overwrites, defaults\n\n @classmethod\n def update_from_dict(\n cls: type[\"QuantRotationConfig\"], *, parsed_args: dict[str, tp.Any], overwrites: dict[str, tp.Any]\n ) -> tuple[dict[str, tp.Any], dict[str, tp.Any]]:\n \"\"\"Create a rotation quantization configuration from the parsed arguments.\"\"\"\n parsed_args.setdefault(\"transforms\", []).extend(omniconfig.COLLECT_PREFIX_BOOL_FIELDS(parsed_args, \"transform\"))\n return parsed_args, overwrites\n"
},
{
"path": "deepcompressor/calib/config/search.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Quantization calibrator configurations.\"\"\"\n\nimport enum\nfrom dataclasses import dataclass\n\nfrom omniconfig import configclass\n\nfrom ...utils.common import num2str\n\n__all__ = [\n \"SearchBasedCalibStrategy\",\n \"SearchBasedCalibGranularity\",\n \"SearchBasedCalibObjective\",\n \"SearchBasedCalibConfig\",\n]\n\n\nclass SearchBasedCalibStrategy(enum.Enum):\n \"\"\"The strategy for search-based quantization calibration.\"\"\"\n\n Manual = enum.auto()\n GridSearch = enum.auto()\n # RandomSearch = enum.auto()\n # Bayesian = enum.auto()\n # EvolutionaryAlgorithm = enum.auto()\n # EvolutionaryStrategy = enum.auto()\n\n\nclass SearchBasedCalibGranularity(enum.Enum):\n \"\"\"The granularity for search-based quantization calibration.\"\"\"\n\n Group = enum.auto()\n ChannelGroup = enum.auto()\n Layer = enum.auto()\n\n\nclass SearchBasedCalibObjective(enum.Enum):\n \"\"\"The objective for search-based quantization calibration.\"\"\"\n\n TensorError = enum.auto()\n \"\"\"minimize the quantization error of the tensor.\"\"\"\n ProductsError = enum.auto()\n \"\"\"minimize the error of the the multiplication products.\"\"\"\n OutputsError = enum.auto()\n \"\"\"minimize the error of the outputs of the evaluation module.\"\"\"\n\n\n@configclass\n@dataclass\nclass SearchBasedCalibConfig:\n \"\"\"The base configuration for search-based quantization calibration.\n\n Args:\n degree (`int`, *optional*, default=`2`):\n The power degree for the quantization error. Defaults to `2`.\n objective (`SearchBasedCalibObjective`, *optional*, default=`SearchBasedCalibObjective.OutputsError`):\n The objective for quantization calibration.\n strategy (`SearchBasedCalibStrategy`, *optional*, default=`SearchBasedCalibStrategy.Manual`):\n The strategy for quantization calibration.\n granularity (`SearchBasedCalibGranularity`, *optional*, default=`SearchBasedCalibGranularity.Layer`):\n The granularity for quantization calibration.\n element_batch_size (`int`, *optional*, default=`-1`):\n The element batch size for calibration.\n sample_batch_size (`int`, *optional*, default=`-1`):\n The samples batch size for calibration.\n element_size (`int`, *optional*, default=`-1`):\n The calibration element size.\n sample_size (`int`, *optional*, default=`-1`):\n The calibration sample size.\n pre_reshape (`bool`, *optional*, default=`True`):\n Whether to enable reshaping the tensor before calibration.\n outputs_device (`str`, *optional*, default=`\"cpu\"`):\n The device to store the precomputed outputs of the module.\n \"\"\"\n\n degree: int = 2\n objective: SearchBasedCalibObjective = SearchBasedCalibObjective.OutputsError\n strategy: SearchBasedCalibStrategy = SearchBasedCalibStrategy.Manual\n granularity: SearchBasedCalibGranularity = SearchBasedCalibGranularity.Layer\n element_batch_size: int = -1\n sample_batch_size: int = -1\n element_size: int = -1\n sample_size: int = -1\n pre_reshape: bool = True\n outputs_device: str = \"cpu\"\n\n def __post_init__(self) -> None:\n if self.outputs_device != \"cpu\":\n self.outputs_device = None\n if self.element_size != 0 or self.sample_size != 0:\n assert self.element_batch_size != 0, \"element_batch_size must not be zero\"\n assert self.sample_batch_size != 0, \"sample_batch_size must not be zero\"\n assert self.element_size != 0, \"element_size must not be zero\"\n assert self.sample_size != 0, \"sample_size must not be zero\"\n else:\n assert self.objective == SearchBasedCalibObjective.TensorError\n if self.objective == SearchBasedCalibObjective.TensorError:\n pass\n elif self.granularity == SearchBasedCalibGranularity.Layer:\n self.objective = SearchBasedCalibObjective.OutputsError\n self.element_batch_size = -1\n self.element_size = -1\n\n @property\n def needs_search(self) -> bool:\n \"\"\"Whether the search is enabled.\"\"\"\n return self.strategy != SearchBasedCalibStrategy.Manual\n\n def generate_dirnames(self, *, prefix: str = \"\", **kwargs) -> list[str]:\n \"\"\"Generate the directory names of the configuration.\n\n Args:\n prefix (`str`, *optional*, default=`\"\"`):\n The prefix of the directory.\n\n Returns:\n `list[str]`:\n The directory names.\n \"\"\"\n name = f\"{self.objective.name}.{self.strategy.name}.{self.granularity.name}.d{num2str(self.degree)}\"\n name += f\".e{num2str(self.element_size)}.s{num2str(self.sample_size)}\"\n if prefix:\n name = f\"{prefix}.{name}\"\n return [name]\n"
},
{
"path": "deepcompressor/calib/config/smooth.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Smooth quantization configuration.\"\"\"\n\nimport enum\nfrom dataclasses import dataclass, field\n\nimport omniconfig\nfrom omniconfig import configclass\n\nfrom ...utils.common import num2str\nfrom ...utils.config import SkipBasedConfig\nfrom .search import (\n SearchBasedCalibConfig,\n SearchBasedCalibGranularity,\n SearchBasedCalibObjective,\n SearchBasedCalibStrategy,\n)\n\n__all__ = [\n \"SmoothSpanMode\",\n \"SmoothCalibConfig\",\n \"SmoothAttentionCalibConfig\",\n \"SkipBasedSmoothCalibConfig\",\n \"SmoothTransfomerConfig\",\n]\n\n\nclass SmoothSpanMode(enum.Enum):\n \"\"\"The mode for computing the span used in smoothing scale calculation.\"\"\"\n\n AbsMax = enum.auto()\n RootMeanSquare = enum.auto()\n\n\n@configclass\n@dataclass\nclass SmoothCalibConfig(SearchBasedCalibConfig):\n \"\"\"Configuration for smooth quantization.\n\n Args:\n degree (`int`, *optional*, default=`2`):\n The power degree for the quantization error. Defaults to `2`.\n objective (`SearchBasedCalibObjective`, *optional*, default=`SearchBasedCalibObjective.OutputsError`):\n The objective for quantization calibration.\n strategy (`SearchBasedCalibStrategy`, *optional*, default=`SearchBasedCalibStrategy.Manual`):\n The strategy for quantization calibration.\n granularity (`SearchBasedCalibGranularity`, *optional*, default=`SearchBasedCalibGranularity.Layer`):\n The granularity for quantization calibration.\n element_batch_size (`int`, *optional*, default=`-1`):\n The element batch size for calibration.\n sample_batch_size (`int`, *optional*, default=`-1`):\n The samples batch size for calibration.\n element_size (`int`, *optional*, default=`-1`):\n The calibration element size.\n sample_size (`int`, *optional*, default=`-1`):\n The calibration sample size.\n pre_reshape (`bool`, *optional*, default=`True`):\n Whether to enable reshaping the tensor before calibration.\n outputs_device (`str`, *optional*, default=`\"cpu\"`):\n The device to store the precomputed outputs of the module.\n fuse_when_possible (`bool`, *optional*, default=`True`):\n Whether to fuse smooth scale whenever possible.\n allow_a_quant (`bool`, *optional*, default=`True`):\n Whether to allow the quantization for alpha tensor.\n allow_b_quant (`bool`, *optional*, default=`True`):\n Whether to allow the quantization for beta tensor.\n spans (`list[tuple[SmoothSpanMode, SmoothSpanMode]]`, *optional*, default=`[]`):\n The span combinations. The first element is for the alpha and the second element is for the beta.\n alpha (`float`, *optional*, default=`0.5`):\n The smoothing alpha.\n beta (`float`, *optional*, default=`-1`):\n The smoothing beta.\n num_grids (`int`, *optional*, default=`20`):\n The number of grids for grid search.\n allow_low_rank (`bool`, *optional*, default=`False`):\n Whether to allow quantization low-rank branch during calibration.\n \"\"\"\n\n fuse_when_possible: bool = True\n allow_a_quant: bool = True\n allow_b_quant: bool = True\n spans: list[tuple[SmoothSpanMode, SmoothSpanMode]] = field(\n default_factory=list,\n metadata={\n omniconfig.ARGPARSE_KWARGS: {\n \"nargs\": \"+\",\n \"type\": lambda s: tuple(SmoothSpanMode[x.split(\".\")[-1]] for x in s.split(\",\")),\n }\n },\n )\n a_spans: list[SmoothSpanMode] = field(default_factory=list, init=False)\n b_spans: list[SmoothSpanMode] = field(default_factory=list, init=False)\n alpha: float = 0.5\n beta: float = -1\n num_grids: int = 20\n allow_low_rank: bool = False\n\n def __post_init__(self) -> None: # noqa: C901\n # region remove duplicates of ranges\n _spans, _spanset, _a_spanset, _b_spanset = [], set(), set(), set()\n self.a_spans, self.b_spans = [], []\n for a_span, b_span in self.spans:\n if isinstance(a_span, str):\n a_span = SmoothSpanMode[a_span]\n if isinstance(b_span, str):\n b_span = SmoothSpanMode[b_span]\n assert isinstance(a_span, SmoothSpanMode), f\"Invalid span mode used for alpha: {a_span}\"\n assert isinstance(b_span, SmoothSpanMode), f\"Invalid span mode used for beta: {b_span}\"\n _span = (a_span, b_span)\n if _span in _spanset:\n continue\n _spans.append(_span)\n _spanset.add(_span)\n if a_span not in _a_spanset:\n _a_spanset.add(a_span)\n self.a_spans.append(a_span)\n if b_span not in _b_spanset:\n _b_spanset.add(b_span)\n self.b_spans.append(b_span)\n self.spans = _spans\n # endregion\n if self.strategy == SearchBasedCalibStrategy.Manual:\n assert len(self.spans) == 1, \"Only one span combination is allowed in manual mode\"\n assert self.alpha != 0 or self.beta != 0, \"alpha and beta cannot be both zero\"\n self.alpha, self.beta = self.get_alpha_beta_pairs()[0]\n if self.granularity == SearchBasedCalibGranularity.Group:\n self.granularity = SearchBasedCalibGranularity.ChannelGroup\n if self.allow_low_rank:\n self.granularity = SearchBasedCalibGranularity.Layer\n assert -3 <= self.alpha <= 1, \"alpha must be less than or equal to 1\"\n assert -3 <= self.beta <= 1, \"beta must be less than or equal to 1\"\n super().__post_init__()\n\n def get_alpha_beta_pairs(self) -> list[tuple[float, float]]: # noqa: C901\n \"\"\"Get the alpha and beta pairs for smooth quantization.\n\n Returns:\n `list[tuple[float, float]]`:\n The alpha and beta pair candidates.\n \"\"\"\n if self.strategy == SearchBasedCalibStrategy.Manual:\n if self.beta < 0:\n assert 0 <= self.alpha <= 1, \"alpha must be in [0, 1]\"\n return [(self.alpha, 1 - self.alpha)]\n elif self.alpha < 0:\n assert 0 <= self.beta <= 1, \"beta must be in [0, 1]\"\n return [(1 - self.beta, self.beta)]\n else:\n assert 0 <= self.alpha <= 1, \"alpha must be in [0, 1]\"\n assert 0 <= self.beta <= 1, \"beta must be in [0, 1]\"\n return [(self.alpha, self.beta)]\n choices = [i / self.num_grids for i in range(1, self.num_grids)]\n if self.alpha > 0:\n if self.beta > 0:\n return [(0, 0)] + [(alpha, alpha) for alpha in choices]\n if self.beta == 0:\n return [(0, 0)] + [(alpha, 0) for alpha in choices]\n if self.beta == -1:\n return [(0, 0)] + [(alpha, 1 - alpha) for alpha in choices]\n if self.beta == -2:\n return [(0, 0)] + [(alpha, 0) for alpha in choices] + [(alpha, 1 - alpha) for alpha in choices]\n return (\n [(0, 0)] + [(alpha, 0) for alpha in choices] + [(alpha, beta) for alpha in choices for beta in choices]\n )\n if self.alpha == 0:\n if self.beta > 0:\n return [(0, 0)] + [(0, beta) for beta in choices]\n if self.beta == 0:\n return [(0, 0)] + [(alpha, 0) for alpha in choices] + [(0, beta) for beta in choices]\n if self.beta == -1:\n return [(0, 0)] + [(0, beta) for beta in choices] + [(alpha, 1 - alpha) for alpha in choices]\n if self.beta == -2:\n return (\n [(0, 0)]\n + [(alpha, 0) for alpha in choices]\n + [(0, beta) for beta in choices]\n + [(alpha, 1 - alpha) for alpha in choices]\n )\n return (\n [(0, 0)]\n + [(alpha, 0) for alpha in choices]\n + [(0, beta) for beta in choices]\n + [(alpha, beta) for alpha in choices for beta in choices]\n )\n if self.alpha == -1:\n if self.beta > 0 or self.beta == -1:\n return [(0, 0)] + [(alpha, 1 - alpha) for alpha in choices]\n if self.beta == 0 or self.beta == -2:\n return [(0, 0)] + [(alpha, 0) for alpha in choices] + [(alpha, 1 - alpha) for alpha in choices]\n return (\n [(0, 0)] + [(alpha, 0) for alpha in choices] + [(alpha, beta) for alpha in choices for beta in choices]\n )\n if self.alpha == -2:\n if self.beta > 0 or self.beta == -1:\n return [(0, 0)] + [(0, beta) for beta in choices] + [(alpha, 1 - alpha) for alpha in choices]\n if self.beta == 0 or self.beta == -2:\n return (\n [(0, 0)]\n + [(alpha, 0) for alpha in choices]\n + [(0, beta) for beta in choices]\n + [(alpha, 1 - alpha) for alpha in choices]\n )\n return (\n [(0, 0)]\n + [(alpha, 0) for alpha in choices]\n + [(0, beta) for beta in choices]\n + [(alpha, beta) for alpha in choices for beta in choices]\n )\n if self.alpha == -3:\n if self.beta > 0:\n return (\n [(0, 0)]\n + [(0, beta) for beta in choices]\n + [(alpha, beta) for alpha in choices for beta in choices]\n )\n return (\n [(0, 0)]\n + [(0, beta) for beta in choices]\n + [(alpha, 0) for alpha in choices]\n + [(alpha, beta) for alpha in choices for beta in choices]\n )\n raise ValueError(\"Invalid alpha and beta values\")\n\n def generate_dirnames(self, *, prefix: str = \"\", **kwargs) -> list[str]:\n \"\"\"Get the directory names of the smooth quantization configuration.\n\n Args:\n prefix (`str`, *optional*, default=`\"\"`):\n The prefix of the directory.\n\n Returns:\n `list[str]`:\n The directory names of the configuration.\n \"\"\"\n names = super().generate_dirnames(**kwargs)\n names.append(\"[{}]\".format(\"+\".join(f\"a.{a_span.name}.b.{b_span.name}\" for a_span, b_span in self.spans)))\n alpha, beta = num2str(self.alpha), num2str(self.beta)\n if self.strategy == SearchBasedCalibStrategy.Manual:\n names.append(f\"a{alpha}.b{beta}\")\n elif self.alpha > 0:\n names.append(f\"g{self.num_grids}.b{beta}\")\n elif self.beta > 0:\n names.append(f\"g{self.num_grids}.a{alpha}\")\n else:\n names.append(f\"g{self.num_grids}.a{alpha}.b{beta}\")\n if self.allow_low_rank:\n names[-1] += \".lr\"\n if not self.fuse_when_possible:\n names[-1] += \".nf\"\n disallows = []\n if not self.allow_a_quant:\n disallows.append(\"a\")\n if not self.allow_b_quant:\n disallows.append(\"b\")\n if disallows:\n names.append(f\"disallow.[{'+'.join(disallows)}]\")\n if prefix:\n names = [f\"{prefix}.{name}\" for name in names]\n return names\n\n\n@configclass\n@dataclass\nclass SkipBasedSmoothCalibConfig(SkipBasedConfig, SmoothCalibConfig):\n \"\"\"Configuration for smooth quantization.\n\n Args:\n degree (`int`, *optional*, default=`2`):\n The power degree for the quantization error. Defaults to `2`.\n objective (`SearchBasedCalibObjective`, *optional*, default=`SearchBasedCalibObjective.OutputsError`):\n The objective for quantization calibration.\n strategy (`SearchBasedCalibStrategy`, *optional*, default=`SearchBasedCalibStrategy.Manual`):\n The strategy for quantization calibration.\n granularity (`SearchBasedCalibGranularity`, *optional*, default=`SearchBasedCalibGranularity.Layer`):\n The granularity for quantization calibration.\n element_batch_size (`int`, *optional*, default=`-1`):\n The element batch size for calibration.\n sample_batch_size (`int`, *optional*, default=`-1`):\n The samples batch size for calibration.\n element_size (`int`, *optional*, default=`-1`):\n The calibration element size.\n sample_size (`int`, *optional*, default=`-1`):\n The calibration sample size.\n pre_reshape (`bool`, *optional*, default=`True`):\n Whether to enable reshaping the tensor before calibration.\n outputs_device (`str`, *optional*, default=`\"cpu\"`):\n The device to store the precomputed outputs of the module.\n allow_a_quant (`bool`, *optional*, default=`True`):\n Whether to allow the quantization for alpha tensor.\n allow_b_quant (`bool`, *optional*, default=`True`):\n Whether to allow the quantization for beta tensor.\n spans (`list[tuple[SmoothSpanMode, SmoothSpanMode]]`, *optional*, default=`[]`):\n The span combinations. The first element is for the alpha and the second element is for the beta.\n alpha (`float`, *optional*, default=`0.5`):\n The smoothing alpha.\n beta (`float`, *optional*, default=`-1`):\n The smoothing beta.\n num_grids (`int`, *optional*, default=`20`):\n The number of grids for grid search.\n allow_low_rank (`bool`, *optional*, default=`False`):\n Whether to allow quantization SVD during calibration.\n skips (`list[str]`, *optional*, default=`[]`):\n The keys of the modules to skip.\n \"\"\"\n\n pass\n\n\n@configclass\n@dataclass\nclass SmoothAttentionCalibConfig(SmoothCalibConfig):\n \"\"\"Configuration for smooth quantization.\n\n Args:\n degree (`int`, *optional*, default=`2`):\n The power degree for the quantization error. Defaults to `2`.\n strategy (`SearchBasedCalibStrategy`, *optional*, default=`SearchBasedCalibStrategy.Manual`):\n The strategy for quantization calibration.\n sample_batch_size (`int`, *optional*, default=`-1`):\n The samples batch size for calibration.\n sample_size (`int`, *optional*, default=`-1`):\n The calibration sample size.\n outputs_device (`str`, *optional*, default=`\"cpu\"`):\n The device to store the precomputed outputs of the module.\n allow_a_quant (`bool`, *optional*, default=`True`):\n Whether to allow the quantization for alpha tensor.\n allow_b_quant (`bool`, *optional*, default=`True`):\n Whether to allow the quantization for beta tensor.\n spans (`list[tuple[SmoothSpanMode, SmoothSpanMode]]`, *optional*, default=`[]`):\n The span combinations. The first element is for the alpha and the second element is for the beta.\n alpha (`float`, *optional*, default=`0.5`):\n The smoothing alpha.\n beta (`float`, *optional*, default=`-1`):\n The smoothing beta.\n num_grids (`int`, *optional*, default=`20`):\n The number of grids for grid search.\n \"\"\"\n\n objective: SearchBasedCalibObjective = field(init=False, default=SearchBasedCalibObjective.OutputsError)\n granularity: SearchBasedCalibGranularity = field(init=False, default=SearchBasedCalibGranularity.Layer)\n element_batch_size: int = field(init=False, default=-1)\n element_size: int = field(init=False, default=-1)\n pre_reshape: bool = field(init=False, default=True)\n allow_low_rank: bool = field(init=False, default=False)\n\n\n@configclass\n@dataclass\nclass SmoothTransfomerConfig:\n \"\"\"Configuration for smooth quantization of transformer-based models.\n\n Args:\n proj (`SkipBasedSmoothCalibConfig` or `None`, *optional*, default=`None`):\n The smooth configuration for projections.\n attn (`SmoothAttentionCalibConfig` or `None`, *optional*, default=`None`):\n The smooth configuration for attentions.\n \"\"\"\n\n proj: SkipBasedSmoothCalibConfig | None = None\n attn: SmoothAttentionCalibConfig | None = None\n\n @property\n def enabled_proj(self) -> bool:\n \"\"\"Whether the smooth quantization is enabled for projections.\"\"\"\n return self.proj is not None\n\n @property\n def enabled_attn(self) -> bool:\n \"\"\"Whether the smooth quantization is enabled for attentions.\"\"\"\n return self.attn is not None\n\n def generate_dirnames(self, *, prefix: str = \"\", **kwargs) -> list[str]:\n \"\"\"Get the names of the smooth quantization configuration.\n\n Args:\n prefix (`str`, *optional*, default=`\"\"`):\n The prefix of the directory.\n\n Returns:\n `list[str]`:\n The names of the smooth quantization configuration\n \"\"\"\n proj_names = self.proj.generate_dirnames(prefix=\"proj\") if self.proj is not None else []\n attn_names = self.attn.generate_dirnames(prefix=\"attn\") if self.attn is not None else []\n num_names = max(len(proj_names), len(attn_names))\n names = []\n for index in range(num_names):\n name = []\n if index < len(proj_names):\n name.append(proj_names[index])\n if index < len(attn_names):\n name.append(attn_names[index])\n names.append(\"-\".join(name))\n if prefix:\n names = [f\"{prefix}.{name}\" for name in names]\n return names\n"
},
{
"path": "deepcompressor/calib/lowrank.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Quantization SVD calibration module.\"\"\"\n\nfrom dataclasses import _MISSING_TYPE, MISSING\n\nimport torch\nimport torch.nn as nn\n\nfrom ..data.common import TensorType\nfrom ..nn.patch.lowrank import LowRankBranch\nfrom ..quantizer.processor import Quantizer\nfrom ..utils import math, tools\nfrom ..utils.config import KeyEnableConfig\nfrom .config import QuantLowRankCalibConfig, SearchBasedCalibObjective\nfrom .search import SearchBasedCalibrator\n\n__all__ = [\"QuantLowRankCalibrator\"]\n\n\nclass QuantLowRankCalibrator(SearchBasedCalibrator[QuantLowRankCalibConfig, LowRankBranch]):\n \"\"\"The quantization low-rank branch calibrator.\"\"\"\n\n def __init__(\n self,\n config: QuantLowRankCalibConfig,\n w_quantizer: Quantizer,\n x_quantizer: Quantizer | None,\n develop_dtype: torch.dtype = torch.float32,\n ) -> None:\n \"\"\"Initialize the calibrator.\n\n Args:\n config (`QuantLowRankCalibConfig`):\n The configuration of the quantization low-rank branch calibrator.\n w_quantizer (`Quantizer`):\n The quantizer for weights.\n x_quantizer (`Quantizer` or `None`):\n The quantizer for inputs.\n develop_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The development data type.\n \"\"\"\n if isinstance(config, KeyEnableConfig):\n assert config.is_enabled_for(w_quantizer.key), \"The calibrator should be enabled for the quantizer.\"\n else:\n assert config.is_enabled(), \"The calibrator should be enabled.\"\n super().__init__(\n tensor_type=TensorType.Weights,\n config=config,\n w_quantizer=w_quantizer,\n x_quantizer=x_quantizer,\n y_quantizer=None,\n develop_dtype=develop_dtype,\n )\n assert self.needs_quant, \"The tensor should be quantized.\"\n self.num_iters = config.num_iters\n\n @property\n def population_size(self) -> int:\n \"\"\"Return the population size of the current iteration.\"\"\"\n return 1\n\n @property\n def allows_x_quant_for_wgts(self) -> bool:\n \"\"\"Whether the calibrator allows input quantization when tensor_type is Weights.\"\"\"\n return True\n\n @property\n def allows_w_quant_for_wgts(self) -> bool:\n \"\"\"Whether the calibrator needs weight quantization when tensor_type is Weights.\"\"\"\n return True\n\n def is_done(self) -> bool:\n \"\"\"Check if the calibration is done.\"\"\"\n return self.iter >= self.num_iters or self.early_stopped\n\n def is_last_iter(self) -> bool:\n \"\"\"Check if the current iteration is the last one.\"\"\"\n return self.iter == self.num_iters - 1\n\n def _reset(self, x_wgts: list[torch.Tensor | nn.Parameter], **kwargs) -> None: # noqa: C901\n \"\"\"Reset the calibrator.\n\n Args:\n x_wgts (`list[torch.Tensor | nn.Parameter]`):\n The weights in x-w computation.\n \"\"\"\n self.best_branch: LowRankBranch = None\n self.best_error: torch.Tensor = None\n self.error_history: list[tuple[float, float]] = []\n self.early_stopped = False\n if len(x_wgts) > 1 and not self.config.exclusive:\n self.w = torch.cat([wgt.data for wgt in x_wgts], dim=0)\n else:\n assert len(x_wgts) == 1\n self.w = x_wgts[0].data\n if self.config.compensate:\n self.qw = torch.cat(\n [\n self.w_quantizer.quantize(wgt.data, kernel=None, develop_dtype=self.develop_dtype).data\n for wgt in x_wgts\n ],\n dim=0,\n )\n else:\n self.qw = 0\n self.hat_ws: list[torch.Tensor] = [None] * len(x_wgts)\n self.ocs: list[int] = [wgt.shape[0] for wgt in x_wgts]\n\n def get_best(self) -> LowRankBranch:\n \"\"\"Get the best candidate.\n\n Returns:\n `LowRankBranch`:\n The best candidate.\n \"\"\"\n return self.best_branch\n\n def _ask(self) -> LowRankBranch:\n \"\"\"Ask for the next candidate.\n\n Returns:\n `LowRankBranch`:\n The next candidate.\n \"\"\"\n branch = LowRankBranch(\n self.w.shape[1],\n self.w.shape[0],\n rank=self.config.rank,\n weight=self.w - self.qw,\n )\n self.wgt_idx = 0\n if len(self.hat_ws) > 1:\n lw = branch.get_effective_weight().view(self.w.shape)\n rw = self.w - lw\n oc_idx = 0\n for idx, oc in enumerate(self.ocs):\n self.hat_ws[idx] = self.w_quantizer.quantize(\n rw[oc_idx : oc_idx + oc], kernel=None, develop_dtype=self.develop_dtype\n ).data\n oc_idx += oc\n self.qw = torch.cat(self.hat_ws, dim=0)\n if self.objective != SearchBasedCalibObjective.OutputsError:\n oc_idx = 0\n for idx, oc in enumerate(self.ocs):\n self.hat_ws[idx].add_(lw[oc_idx : oc_idx + oc])\n oc_idx += oc\n else:\n lw = branch.get_effective_weight().view(self.w.shape)\n self.qw = self.w_quantizer.quantize(self.w - lw, kernel=None, develop_dtype=self.develop_dtype).data\n if self.objective != SearchBasedCalibObjective.OutputsError:\n self.hat_ws = [self.qw + lw]\n else:\n self.hat_ws = [self.qw]\n return branch\n\n def _tell(self, error: list[torch.Tensor]) -> None: # noqa: C901\n \"\"\"Tell the error of the last candidate and update the best candidate.\n\n Args:\n errors (list[torch.Tensor]): The error of the last candidate.\n \"\"\"\n if len(error) > 1:\n error = [sum(error)]\n error = error[0]\n assert isinstance(error, torch.Tensor)\n assert error.numel() == 1, \"The error should only have one value.\"\n if self.best_error is None or error <= self.best_error:\n self.best_error = error\n self.best_branch = self.candidate\n elif self.config.early_stop:\n self.early_stopped = True\n if self.logger.level <= tools.logging.DEBUG:\n self.error_history.append(\n (\n math.root_(error.to(torch.float64), self.config.degree).item(),\n math.root_(self.best_error.to(torch.float64), self.config.degree).item(),\n )\n )\n if self.iter % 10 == 9 or self.is_last_iter() or self.early_stopped:\n iter_end = ((self.iter + 10) // 10) * 10\n iter_start = iter_end - 10\n iter_end = min(iter_end, self.iter + 1)\n history = self.error_history[iter_start:iter_end]\n self.logger.debug(\" - iter = [%s]\", \", \".join(f\"{i:10d}\" for i in range(iter_start, iter_end)))\n self.logger.debug(\" - error = [%s]\", \", \".join(f\"{e[0]:10.4f}\" for e in history))\n self.logger.debug(\" - best error = [%s]\", \", \".join(f\"{e[1]:10.4f}\" for e in history))\n\n def _process_x_in_xw(self, x: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor:\n if not self.needs_x_quant_for_wgts:\n return x\n return self.x_quantizer.quantize(x, channels_dim=channels_dim, develop_dtype=self.develop_dtype).data\n\n def _process_w_in_xw(self, w: torch.Tensor) -> torch.Tensor:\n hat_w = self.hat_ws[self.wgt_idx]\n self.hat_ws[self.wgt_idx] = None\n self.wgt_idx += 1\n return hat_w if self.needs_w_quant_for_wgts else w\n\n def _process_y_in_yx(self, y: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor:\n raise RuntimeError(\"_process_y_in_yx should not be called in QuantSVDCalibrator.\")\n\n def _process_x_in_yx(self, x: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor:\n raise RuntimeError(\"_process_x_in_yx should not be called in QuantSVDCalibrator.\")\n\n def _process_xw_in_yx(self, w: torch.Tensor) -> torch.Tensor:\n raise RuntimeError(\"_process_xw_in_yx should not be called in QuantSVDCalibrator.\")\n\n def _process_yw_in_yx(self, w: torch.Tensor) -> torch.Tensor:\n raise RuntimeError(\"_process_yw_in_yx should not be called in QuantSVDCalibrator.\")\n\n def _process_wgts_centric_mod(\n self, wgts: list[nn.Parameter], mods: list[nn.Module], update_state_dict: bool = True, **kwargs\n ) -> None:\n assert len(self.hat_ws) == len(wgts) == len(mods)\n shared = self.candidate\n if len(self.hat_ws) > 1:\n oc_idx = 0\n for mod, wgt, hat_w in zip(mods, wgts, self.hat_ws, strict=True):\n if update_state_dict:\n self._state_dict.append((wgt, wgt.data))\n wgt.data = hat_w\n branch = LowRankBranch(wgt.shape[1], wgt.shape[0], rank=self.config.rank)\n branch.a = shared.a\n branch.b.to(dtype=wgt.dtype, device=wgt.device)\n branch.b.weight.copy_(shared.b.weight[oc_idx : oc_idx + wgt.data.shape[0]])\n oc_idx += wgt.data.shape[0]\n self._hooks.append(branch.as_hook().register(mod))\n else:\n if update_state_dict:\n self._state_dict.append((wgts[0], wgts[0].data))\n wgts[0].data = self.hat_ws[0]\n self._hooks.append(shared.as_hook().register(mods))\n if self.needs_x_quant_for_wgts:\n self._hooks.append(self.x_quantizer.as_hook().register(mods))\n self.hat_ws = [None] * len(self.hat_ws)\n"
},
{
"path": "deepcompressor/calib/metric.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Channel-wise metric calculation module.\"\"\"\n\nimport typing as tp\n\nimport torch\n\nfrom ..data.utils.shape import infer_view_shape\n\n__all__ = [\"ChannelMetric\"]\n\n\nclass ChannelMetric:\n \"\"\"Channel-wise metric.\"\"\"\n\n @staticmethod\n def _normalize(\n tensor: torch.Tensor,\n group_shape: tp.Sequence[int],\n dtype: torch.dtype,\n ) -> torch.Tensor:\n shape, ndim = tensor.shape, tensor.ndim\n view_shape = infer_view_shape(tensor.shape, group_shape)\n # (d0, d1, d2, ...) -> (#g0, gs0, #g1, gs1, #g2, gs2, ...)\n tensor = tensor.view(view_shape).to(dtype=dtype)\n tensor_max = tensor.abs().amax(dim=list(range(1, ndim * 2, 2)), keepdim=True)\n tensor_max[tensor_max == 0] = 1\n tensor = tensor / tensor_max\n return tensor.view(shape)\n\n @staticmethod\n def _abs_max(\n tensor: torch.Tensor,\n num_channels: int,\n group_shape: tp.Sequence[int],\n device: torch.device,\n dtype: torch.dtype,\n ) -> tuple[torch.Tensor, int]:\n return (\n tensor.view(tensor.shape[0], num_channels, -1)\n .abs()\n .amax(dim=(0, 2))\n .view(-1)\n .to(dtype=dtype, device=device),\n 1,\n )\n\n @staticmethod\n def _abs_sum(\n tensor: torch.Tensor,\n num_channels: int,\n group_shape: tp.Sequence[int],\n device: torch.device,\n dtype: torch.dtype,\n ) -> tuple[torch.Tensor, int]:\n tensor = tensor.view(tensor.shape[0], num_channels, -1)\n cnt = tensor.shape[0] * tensor.shape[2]\n return tensor.abs().to(dtype=dtype).sum(dim=(0, 2)).view(-1).to(device=device), cnt\n\n @staticmethod\n def _abs_normalize_sum(\n tensor: torch.Tensor,\n num_channels: int,\n group_shape: tp.Sequence[int],\n device: torch.device,\n dtype: torch.dtype,\n ) -> tuple[torch.Tensor, int]:\n return ChannelMetric._abs_sum(\n ChannelMetric._normalize(tensor, group_shape, dtype=dtype),\n num_channels,\n group_shape,\n device=device,\n dtype=dtype,\n )\n\n @staticmethod\n def _square_sum(\n tensor: torch.Tensor,\n num_channels: int,\n group_shape: tp.Sequence[int],\n device: torch.device,\n dtype: torch.dtype,\n ) -> tuple[torch.Tensor, int]:\n tensor = tensor.view(tensor.shape[0], num_channels, -1)\n cnt = tensor.shape[0] * tensor.shape[2]\n return tensor.to(dtype=dtype).pow(2).sum(dim=(0, 2)).view(-1).to(device=device), cnt\n\n @staticmethod\n def _max_reduce(\n fn: tp.Callable[\n [torch.Tensor, int, tp.Sequence[int], torch.device, torch.dtype],\n tuple[torch.Tensor, torch.Tensor | int | float],\n ],\n tensors: tp.Sequence[torch.Tensor] | torch.Tensor,\n num_channels: int,\n group_shape: tp.Sequence[int],\n device: torch.device | str | None = None,\n dtype: torch.dtype = torch.float32,\n ) -> tuple[torch.Tensor, torch.Tensor | int | float]:\n if isinstance(tensors, torch.Tensor):\n device = device or tensors.device\n return fn(tensors, num_channels, group_shape, device, dtype)\n else:\n rst_0, rst_1 = ChannelMetric._max_reduce(fn, tensors[0], num_channels, group_shape, device, dtype)\n for tensor in tensors[1:]:\n _rst_0, _rst_1 = ChannelMetric._max_reduce(fn, tensor, num_channels, group_shape, device, dtype)\n rst_0 = torch.maximum(rst_0, _rst_0.to(device=rst_0.device))\n if isinstance(rst_1, torch.Tensor):\n rst_1 = torch.maximum(rst_1, _rst_1.to(device=rst_1.device))\n else:\n rst_1 = max(rst_1, _rst_1)\n return rst_0, rst_1\n\n @staticmethod\n def _sum_reduce(\n fn: tp.Callable[\n [torch.Tensor, int, tp.Sequence[int], torch.device, torch.dtype],\n tuple[torch.Tensor, torch.Tensor | int | float],\n ],\n tensors: tp.Sequence[torch.Tensor] | torch.Tensor,\n num_channels: int,\n group_shape: tp.Sequence[int],\n device: torch.device | str | None = None,\n dtype: torch.dtype = torch.float32,\n ) -> tuple[torch.Tensor, torch.Tensor | int | float]:\n if isinstance(tensors, torch.Tensor):\n device = device or tensors.device\n return fn(tensors.to(device), num_channels, group_shape, device, dtype)\n else:\n assert isinstance(tensors, (list, tuple))\n rst_0, rst_1 = ChannelMetric._sum_reduce(fn, tensors[0], num_channels, group_shape, device, dtype)\n for tensor in tensors[1:]:\n _rst_0, _rst_1 = ChannelMetric._sum_reduce(fn, tensor, num_channels, group_shape, device, dtype)\n rst_0 += _rst_0.to(device=rst_0.device)\n if isinstance(rst_1, torch.Tensor):\n rst_1 += _rst_1.to(device=rst_1.device)\n else:\n rst_1 += _rst_1\n return rst_0, rst_1\n\n @staticmethod\n def abs_max(\n tensors: tp.Iterable[torch.Tensor] | torch.Tensor,\n num_channels: int,\n group_shape: tp.Sequence[int],\n device: torch.device | str = None,\n dtype: torch.dtype = torch.float32,\n ) -> torch.Tensor:\n \"\"\"Get the absolute maximum of the tensors, where `R[i] = AbsMax(T[i, :])`.\"\"\"\n return ChannelMetric._max_reduce(\n ChannelMetric._abs_max, tensors, num_channels, group_shape, device=device, dtype=dtype\n )[0]\n\n @staticmethod\n def abs_mean(\n tensors: tp.Iterable[torch.Tensor] | torch.Tensor,\n num_channels: int,\n group_shape: tp.Sequence[int],\n device: torch.device | str = None,\n dtype: torch.dtype = torch.float32,\n ) -> torch.Tensor:\n \"\"\"Get the absolute mean of the tensors, where `R[i] = AbsMean(T[i, :])`.\"\"\"\n rst, cnt = ChannelMetric._sum_reduce(\n ChannelMetric._abs_sum, tensors, num_channels, group_shape, device=device, dtype=dtype\n )\n return rst.div_(cnt)\n\n @staticmethod\n def abs_normalize_mean(\n tensors: tp.Iterable[torch.Tensor] | torch.Tensor,\n num_channels: int,\n group_shape: tp.Sequence[int],\n device: torch.device | str = None,\n dtype: torch.dtype = torch.float32,\n ) -> torch.Tensor:\n \"\"\"Get the absolute group normalized mean of the tensors, where `R[i] = Mean(U[i, :])`\n and `U[i,j] = Abs(T[i, j]) / AbsMax(T[:, j]))`.\"\"\"\n rst, cnt = ChannelMetric._sum_reduce(\n ChannelMetric._abs_normalize_sum, tensors, num_channels, group_shape, device=device, dtype=dtype\n )\n return rst.div_(cnt)\n\n @staticmethod\n def root_mean_square(\n tensors: tp.Iterable[torch.Tensor] | torch.Tensor,\n num_channels: int,\n group_shape: tp.Sequence[int],\n device: torch.device | str = None,\n dtype: torch.dtype = torch.float32,\n ) -> torch.Tensor:\n \"\"\"Get the root mean square of the tensors, where `R[i] = Root(Mean(T[i, :]^2))`.\"\"\"\n rst, cnt = ChannelMetric._sum_reduce(\n ChannelMetric._square_sum, tensors, num_channels, group_shape, device=device, dtype=dtype\n )\n return rst.div_(cnt).sqrt_()\n"
},
{
"path": "deepcompressor/calib/range.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Quantization dynamic range calibration.\"\"\"\n\nimport gc\nimport typing as tp\nfrom dataclasses import _MISSING_TYPE, MISSING\n\nimport torch\nimport torch.nn as nn\n\nfrom ..data.cache import TensorsCache\nfrom ..data.common import TensorType\nfrom ..data.range import DynamicRange\nfrom ..data.scale import QuantScale\nfrom ..data.utils.shape import infer_view_shape\nfrom ..quantizer.impl.info import QuantInfo\nfrom ..quantizer.processor import Quantizer\nfrom ..utils import math, tools\nfrom .config import DynamicRangeCalibConfig, SearchBasedCalibGranularity\nfrom .search import SearchBasedCalibrator\n\n__all__ = [\"DynamicRangeCalibrator\", \"calibrate_dynamic_range\"]\n\n\nclass DynamicRangeCalibrator(SearchBasedCalibrator[DynamicRangeCalibConfig, DynamicRange]):\n \"\"\"The quantization dynamic range calibrator.\"\"\"\n\n def __init__(\n self,\n tensor_type: TensorType,\n config: DynamicRangeCalibConfig,\n static: bool,\n quantizer: Quantizer,\n pre_scale: torch.Tensor | None = None,\n ) -> None:\n \"\"\"Initialize the calibrator.\n\n Args:\n tensor_type (`TensorType`):\n The tensor type.\n config (`DynamicRangeCalibConfig`):\n The dynamic range calibration configuration.\n static (`bool`):\n Whether the dynamic range is static, i.e., whether the quantization is static.\n quantizer (`Quantizer`):\n The quantizer.\n pre_scale (`torch.Tensor` or `None`):\n The joint scale tensor of the previous quantization steps.\n \"\"\"\n super().__init__(\n tensor_type=tensor_type,\n config=config,\n w_quantizer=quantizer if tensor_type == TensorType.Weights else None,\n x_quantizer=quantizer if tensor_type == TensorType.Inputs else None,\n y_quantizer=quantizer if tensor_type == TensorType.Outputs else None,\n develop_dtype=quantizer.develop_dtype,\n )\n assert self.needs_quant, \"The tensor should be quantized.\"\n self.static = static\n self.pre_scale = pre_scale\n self.ratios = self.config.get_ratios()\n self.num_iters = len(self.ratios)\n\n @property\n def population_size(self) -> int:\n \"\"\"Return the population size of the current iteration.\"\"\"\n return len(self.ratios[self.iter])\n\n def is_clamp_based(self) -> bool:\n \"\"\"Return whether the calibration is clamp-based.\"\"\"\n return self.static or not self.config.allow_scale\n\n def _reset( # noqa: C901\n self,\n x_wgts: list[torch.Tensor | nn.Parameter],\n x_acts: TensorsCache | None,\n y_acts: TensorsCache | None,\n **kwargs,\n ) -> None:\n \"\"\"Reset the calibrator.\n\n Args:\n x_wgts (`list[torch.Tensor | nn.Parameter]`):\n The weights in x-w computation.\n x_acts (`TensorsCache` or `None`):\n The x activations in x-w computation.\n y_acts (`TensorsCache` or `None`):\n The y activations in y-x computation.\n \"\"\"\n self.base_range: DynamicRange = DynamicRange()\n self.best_range: DynamicRange = DynamicRange()\n self.best_error: torch.Tensor = None\n self.error_history: list[tuple[float, float]] = []\n self.device = None\n if self.tensor_type == TensorType.Weights:\n assert len(x_wgts) == 1, \"The weight should be a single tensor.\"\n wgts = x_wgts[0].data\n assert isinstance(wgts, torch.Tensor), \"The weight should be a tensor.\"\n tensors = [wgts]\n self.device = wgts.device\n elif self.tensor_type == TensorType.Inputs:\n assert x_acts is not None, \"The input activations should be provided.\"\n assert x_acts.num_tensors == 1, f\"Only one input is allowed, got {x_acts.num_tensors}\"\n acts = x_acts.front()\n tensors = acts.get_standardized_data(reshape=False)\n self.device = acts.orig_device\n else:\n assert y_acts is not None, \"The output activations should be provided.\"\n assert y_acts.num_tensors == 1, f\"Only one output is allowed, got {y_acts.num_tensors}\"\n acts = y_acts.front()\n tensors = acts.get_standardized_data(reshape=False)\n self.device = acts.orig_device\n shape = tensors[0].shape\n view_shape = infer_view_shape(\n shape,\n self.quantizer.config.largest_group_shape,\n skip_first_dim=self.tensor_type != TensorType.Weights,\n )\n # region get range scale shape\n self.pos_view_shape = torch.Size([1, 1, view_shape[2], *([1] * (len(view_shape) - 3))])\n self.range_shape = torch.Size([gs if i % 2 == 0 else 1 for i, gs in enumerate(view_shape)])\n if self.granularity == SearchBasedCalibGranularity.Layer:\n self.ratio_shape = self.error_shape = torch.Size((1,))\n self.ratio_view_shape = self.ratio_shape\n elif self.granularity == SearchBasedCalibGranularity.ChannelGroup:\n self.ratio_shape = self.error_shape = torch.Size((view_shape[2],))\n self.ratio_view_shape = self.pos_view_shape\n elif self.granularity == SearchBasedCalibGranularity.Group:\n self.ratio_shape = self.error_shape = torch.Size(view_shape[::2])\n self.ratio_view_shape = self.range_shape\n else:\n raise ValueError(f\"Invalid granularity: {self.granularity}\")\n assert self.ratio_shape.numel() == self.ratio_view_shape.numel()\n if self.pre_scale is not None:\n assert len(shape) * 2 == len(self.pre_scale.shape)\n self.pre_view_shape = infer_view_shape(shape, self.pre_scale.shape[1::2])\n else:\n self.pre_view_shape = torch.Size()\n # endregion\n if self.is_clamp_based():\n if self.pre_scale is not None:\n tensors = [self._preprocess_with_pre_scale(t) for t in tensors]\n tensors = [t.view(view_shape).to(dtype=self.develop_dtype) for t in tensors]\n self.base_range = DynamicRange.construct(\n tensors,\n zero_domain=self.quantizer.config.zero_domain,\n is_float_point=self.quantizer.config.quant_dtype.is_float_point,\n )\n gc.collect()\n torch.cuda.empty_cache()\n\n def get_best(self) -> DynamicRange:\n \"\"\"Get the best candidate.\n\n Returns:\n `DynamicRange`:\n The best candidate.\n \"\"\"\n if self.static:\n return DynamicRange(min=self.best_range.min, max=self.best_range.max)\n elif self.is_clamp_based():\n return DynamicRange(min=self.best_range.min, max=self.best_range.max, ratio=1.0)\n else:\n return DynamicRange(ratio=self.best_range.ratio.view(self.ratio_view_shape))\n\n def _ask(self) -> DynamicRange:\n \"\"\"Ask for the next candidate.\n\n Returns:\n `DynamicRange`:\n The next candidate.\n \"\"\"\n ratio = self.ratios[self.iter][self.candidate_id]\n if self.is_clamp_based():\n return self.base_range.scale(\n ratio=ratio,\n zero_domain=self.quantizer.config.zero_domain,\n is_float_point=self.quantizer.config.quant_dtype.is_float_point,\n )\n else:\n return DynamicRange(ratio=ratio)\n\n def _tell(self, error: list[torch.Tensor]) -> None: # noqa: C901\n \"\"\"Tell the error of the last candidate and update the best candidate.\n\n Args:\n errors (`list[torch.Tensor]`):\n The error of the last candidate.\n \"\"\"\n assert len(error) == 1, \"The error should only have one value.\"\n error = error[0]\n assert isinstance(error, torch.Tensor)\n assert error.shape == self.error_shape, f\"Error shape {error.shape} != {self.error_shape}.\"\n assert isinstance(self.candidate, DynamicRange)\n candidate_ratio = self.ratios[self.iter][self.candidate_id]\n if self.best_error is None:\n self.best_error = error\n if self.is_clamp_based():\n self.best_range.min = self.candidate.min\n self.best_range.max = self.candidate.max\n self.best_range.ratio = torch.full(\n size=self.ratio_shape, fill_value=candidate_ratio, device=self.device, dtype=self.develop_dtype\n )\n elif error.numel() > 1:\n pos = error < self.best_error\n self.best_error[pos] = error[pos]\n if self.is_clamp_based():\n if self.error_shape.numel() != self.range_shape.numel():\n pos = pos.view(self.pos_view_shape).expand(*self.range_shape)\n else:\n pos = pos.view(self.range_shape)\n self.best_range.max[pos] = self.candidate.max[pos]\n if isinstance(self.candidate.min, torch.Tensor):\n self.best_range.min[pos] = self.candidate.min[pos]\n self.best_range.ratio[pos.view(self.ratio_shape)] = candidate_ratio\n elif error < self.best_error:\n self.best_error = error\n if self.is_clamp_based():\n self.best_range.min = self.candidate.min\n self.best_range.max = self.candidate.max\n self.best_range.ratio.fill_(candidate_ratio)\n if self.logger.level <= tools.logging.DEBUG:\n self.error_history.append(\n (\n math.root_(error.to(torch.float64).sum(), self.config.degree).item(),\n math.root_(self.best_error.to(torch.float64).sum(), self.config.degree).item(),\n )\n )\n if self.is_last_candidate_in_iter():\n stype_id = self.iter\n ratios, population_size = self.ratios[stype_id], self.population_size\n for i in range(0, population_size, 5):\n self.logger.debug(\n \" - range ratio = [%s]\",\n \", \".join(f\"{ratios[j]:10.4f}\" for j in range(i, min(i + 5, population_size))),\n )\n self.logger.debug(\n \" sum error = [%s]\",\n \", \".join(f\"{self.error_history[j][0]:10.4f}\" for j in range(i, min(i + 5, population_size))),\n )\n self.logger.debug(\n \" best error = [%s]\",\n \", \".join(f\"{self.error_history[j][1]:10.4f}\" for j in range(i, min(i + 5, population_size))),\n )\n self.error_history.clear()\n if self.is_last_iter():\n self.logger.debug(\n \"+ error = [%.4f]\",\n math.root_(self.best_error.to(torch.float64).sum(), self.config.degree).item(),\n )\n\n def _preprocess_with_pre_scale(self, t: torch.Tensor) -> torch.Tensor:\n t = t.view(self.pre_view_shape)\n t = t.to(dtype=self.develop_dtype) if t.dtype != self.develop_dtype else t.clone()\n t = t.div_(self.pre_scale)\n if self.quantizer.range_bound is not None and self.quantizer.range_bound.is_set():\n t = t.clamp_(min=self.quantizer.range_bound.min, max=self.quantizer.range_bound.max)\n return t\n\n def _process_wxy(self, tensor: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor:\n shape, dtype = tensor.shape, tensor.dtype\n if self.pre_scale is not None:\n tensor = self._preprocess_with_pre_scale(tensor).view(shape)\n tensor = self.quantizer.quantize(\n tensor,\n kernel=None,\n channels_dim=channels_dim,\n dynamic_range=self.candidate,\n default_dtype=dtype,\n develop_dtype=self.develop_dtype,\n ).data\n if self.pre_scale is not None:\n tensor = tensor.view(self.pre_view_shape).mul_(self.pre_scale).to(dtype)\n tensor = tensor.view(shape)\n return tensor\n\n def _process_x_in_xw(self, x: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor:\n if self.tensor_type != TensorType.Inputs:\n return x\n return self._process_wxy(x, channels_dim)\n\n def _process_w_in_xw(self, w: torch.Tensor) -> torch.Tensor:\n if self.tensor_type != TensorType.Weights:\n return w\n return self._process_wxy(w, channels_dim=None)\n\n def _process_y_in_yx(self, y: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor:\n if self.tensor_type != TensorType.Outputs:\n return y\n return self._process_wxy(y, channels_dim)\n\n def _process_x_in_yx(self, x: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor:\n raise RuntimeError(\"_process_x_in_yx should not be called in DynamicRangeCalibrator.\")\n\n def _process_xw_in_yx(self, w: torch.Tensor) -> torch.Tensor:\n raise RuntimeError(\"_process_xw_in_yx should not be called in DynamicRangeCalibrator.\")\n\n def _process_yw_in_yx(self, w: torch.Tensor) -> torch.Tensor:\n raise RuntimeError(\"_process_yw_in_yx should not be called in DynamicRangeCalibrator.\")\n\n\ndef calibrate_dynamic_range(\n tensor_type: TensorType,\n config: DynamicRangeCalibConfig | None,\n static: bool,\n quantizer: Quantizer,\n modules: tp.Sequence[nn.Module],\n activations: TensorsCache,\n weights: tp.Sequence[nn.Parameter] | None = None,\n eval_inputs: TensorsCache | None = None,\n eval_module: nn.Module | None = None,\n eval_kwargs: dict[str, tp.Any] | None = None,\n orig_weights: tp.Sequence[tuple[nn.Parameter, torch.Tensor]] | None = None,\n orig_activations: TensorsCache | None = None,\n orig_eval_inputs: TensorsCache | None = None,\n) -> tp.Sequence[DynamicRange] | None:\n \"\"\"Calibrate the dynamic range.\n\n Args:\n tensor_type (`TensorType`):\n The tensor type.\n config (`DynamicRangeCalibConfig`):\n The quantization dynamic range calibration configuration.\n static (`bool`):\n Whether the dynamic range is static.\n quantizer (`Quantizer`):\n The quantizer.\n modules (`Sequence[nn.Module]`):\n The modules to calibrate.\n activations (`TensorsCache`):\n The inputs cache if the tensor type is not outputs, or the outputs cache if the tensor type is outputs.\n weights (`Sequence[nn.Parameter]` or `None`, *optional*, defaults to `None`):\n The weights to calibrate.\n If not provided, the weights of the modules will be used.\n eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The cache of the inputs for evaluation.\n If not provided, the `activations` cache will be used.\n eval_module (`nn.Module` or `None`, *optional*, defaults to `None`):\n The module to evaluate the quantization error.\n If not provided, the module to calibrate will be used.\n eval_kwargs (`dict[str, tp.Any]` or `None`, *optional*, defaults to `None`):\n The keyword arguments for evaluation.\n orig_weights (`Sequence[tuple[nn.Parameter, torch.Tensor]]` or `None`, *optional*, defaults to `None`):\n The original weights.\n orig_activations (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The original activations.\n orig_eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The original evaluation inputs.\n\n Returns:\n `Sequence[DynamicRange]` or `None`:\n The dynamic ranges of each quantization step.\n \"\"\"\n if config is None or not quantizer.is_enabled():\n return None\n\n decomposed_config = quantizer.config.decompose()\n num_steps = decomposed_config.num_steps\n # region dynamic range without search\n if not config.needs_search and (not static or tensor_type == TensorType.Weights):\n if config.ratio != 1.0:\n dynamic_range = DynamicRange(ratio=config.ratio)\n return tuple([dynamic_range] + [None] * (num_steps - 1))\n else:\n return None\n # endregion\n # region prepare for search\n if weights is None:\n weights = [module.weight for module in modules if hasattr(module, \"weight\")]\n if tensor_type == TensorType.Weights:\n assert len(modules) == 1, \"only one module is supported for weight quantization calibration\"\n assert len(weights) == 1, \"only one weight is supported for weight quantization calibration\"\n if eval_module is None:\n eval_module = modules[0]\n if eval_inputs is None:\n eval_inputs = activations\n else:\n assert eval_inputs is not None, \"eval_inputs is required when eval_module is provided\"\n else:\n assert activations is not None, \"activations is required for activation quantization calibration\"\n assert activations.num_tensors == 1, \"only one tensor is supported for activation quantization calibration\"\n if tensor_type != TensorType.Outputs:\n x_wgts, x_acts, x_mods, orig_x_wgts, orig_x_acts = weights, activations, modules, orig_weights, orig_activations\n y_wgts, y_acts, y_mods, orig_y_wgts, orig_y_acts = [], None, None, None, None\n else:\n x_wgts, x_acts, x_mods, orig_x_wgts, orig_x_acts = [], None, None, None, None\n y_wgts, y_acts, y_mods, orig_y_wgts, orig_y_acts = weights, activations, modules, orig_weights, orig_activations\n # endregion\n if num_steps == 1:\n dynamic_range = DynamicRangeCalibrator(\n tensor_type=tensor_type,\n config=config,\n static=static,\n quantizer=quantizer,\n ).calibrate(\n x_wgts=x_wgts,\n y_wgts=y_wgts,\n x_acts=x_acts,\n y_acts=y_acts,\n eval_inputs=eval_inputs,\n eval_module=eval_module,\n eval_kwargs=eval_kwargs,\n x_mods=x_mods,\n y_mods=y_mods,\n orig_x_wgts=orig_x_wgts,\n orig_y_wgts=orig_y_wgts,\n orig_x_acts=orig_x_acts,\n orig_y_acts=orig_y_acts,\n orig_eval_inputs=orig_eval_inputs,\n )\n return (dynamic_range,)\n # region prepare for search with progressive quantization\n if tensor_type == TensorType.Weights:\n tensor = weights[0].detach().data\n else:\n assert activations.num_tensors == 1, \"Only one tensor is supported for activation quantization\"\n acts = activations.front()\n assert len(acts.data) == 0, \"Only one tensor is supported for activation quantization\"\n tensor = acts.data[0].detach().data\n if acts.channels_dim is not None:\n tensor = tensor.reshape(-1, *tensor.shape[acts.channels_dim :])\n develop_dtype = quantizer.develop_dtype\n default_scale_dtype = quantizer.default_dtype or tensor.dtype\n develop_tensor = tensor.to(dtype=develop_dtype) if tensor.dtype != develop_dtype else tensor.clone()\n del tensor\n # endregion\n info = QuantInfo.construct(\n decomposed_config,\n tensor_shape=develop_tensor.shape,\n default_dtype=default_scale_dtype,\n quant_range=quantizer.quant_range,\n range_bound=quantizer.range_bound,\n )\n dynamic_ranges = []\n quant_scale = QuantScale()\n for step, step_info in enumerate(info.steps):\n step_quantizer = Quantizer(\n config=step_info.to_config(),\n kernel=quantizer.kernel if step == num_steps - 1 else None,\n quant_range=step_info.quant_range,\n range_bound=step_info.range_bound,\n default_dtype=quantizer.default_dtype,\n develop_dtype=quantizer.develop_dtype,\n )\n step_dynamic_range = DynamicRangeCalibrator(\n tensor_type=tensor_type,\n config=config,\n static=static,\n quantizer=step_quantizer,\n pre_scale=quant_scale.data,\n ).calibrate(\n x_wgts=x_wgts,\n y_wgts=y_wgts,\n x_acts=x_acts,\n y_acts=y_acts,\n eval_inputs=eval_inputs,\n eval_module=eval_module,\n eval_kwargs=eval_kwargs,\n x_mods=x_mods,\n y_mods=y_mods,\n orig_x_wgts=orig_x_wgts,\n orig_y_wgts=orig_y_wgts,\n )\n dynamic_ranges.append(step_dynamic_range)\n step_scale, _ = step_info.scale.quantize(\n tensor=develop_tensor.view(step_info.tensor_shape),\n dynamic_range=step_dynamic_range,\n )\n quant_scale.append(step_scale)\n if num_steps > 2 and step < num_steps - 1:\n step_quant_range = step_info.tensor_quant_range\n develop_tensor = develop_tensor.view(step_info.tensor_view_shape).div_(step_scale.data)\n develop_tensor = develop_tensor.clamp_(min=step_quant_range.min, max=step_quant_range.max)\n return tuple(dynamic_ranges)\n"
},
{
"path": "deepcompressor/calib/reorder.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Channel reordering module.\"\"\"\n\nimport gc\nimport typing as tp\nfrom dataclasses import _MISSING_TYPE, MISSING, dataclass\n\nimport torch\nimport torch.nn as nn\n\nfrom ..data.cache import TensorsCache\nfrom ..data.common import TensorType\nfrom ..quantizer.processor import Quantizer\nfrom ..utils import math, tools\nfrom ..utils.hooks import BaseInputPackager, BaseOutputPackager, BaseTensorProcessor\nfrom .config import (\n ChannelOrderCalibConfig,\n SearchBasedCalibGranularity,\n SearchBasedCalibObjective,\n SearchBasedCalibStrategy,\n)\nfrom .metric import ChannelMetric\nfrom .search import SearchBasedCalibrator\n\n__all__ = [\"ChannelOrderCalibrator\", \"ChannelReorderer\"]\n\n\n@dataclass\nclass ChannelReorderer(BaseTensorProcessor):\n \"\"\"Activation channel reordering processor.\"\"\"\n\n index: torch.Tensor\n channels_dim: int\n # region hook-related attributes\n input_packager: BaseInputPackager | None = None\n output_packager: BaseOutputPackager | None = None\n # endregion\n\n def is_enabled(self) -> bool:\n return self.index is not None\n\n def get_input_packager(self) -> BaseInputPackager | None:\n return self.input_packager\n\n def get_output_packager(self) -> BaseOutputPackager | None:\n return self.output_packager\n\n def process(self, tensor: torch.Tensor) -> torch.Tensor:\n \"\"\"Process the tensor.\n\n Args:\n tensor (torch.Tensor): The tensor to process.\n\n Returns:\n torch.Tensor: The processed tensor.\n \"\"\"\n self.index = self.index.to(device=tensor.device)\n return tensor.index_select(dim=self.channels_dim, index=self.index)\n\n\ndef get_channel_index_from_rank(\n rank: torch.Tensor,\n num_channels: int,\n num_groups: int,\n index_mode: ChannelOrderCalibConfig.ChannelIndex,\n) -> torch.Tensor:\n \"\"\"Get the index from the rank.\n\n Args:\n rank (`torch.Tensor`):\n The rank of the channels.\n num_channels (`int`):\n The number of channels.\n num_groups (`int`):\n The number of groups.\n index_mode (`ChannelOrderCalibConfig.ChannelIndex`):\n The index mode.\n\n Returns:\n `torch.Tensor`:\n The index of the channels, i.e., the order of the channels.\n \"\"\"\n if index_mode == ChannelOrderCalibConfig.ChannelIndex.Transpose:\n return rank.view(num_channels // num_groups, num_groups).t().reshape(-1)\n elif index_mode == ChannelOrderCalibConfig.ChannelIndex.Sequential:\n return rank\n else:\n raise ValueError(f\"Unsupported index mode: {index_mode}\")\n\n\ndef get_channel_metric(\n inputs: TensorsCache,\n weights: tp.Sequence[torch.Tensor],\n metric_mode: ChannelOrderCalibConfig.ChannelMetric,\n num_channels: int,\n num_heads: int = 1,\n device: torch.device | str | None = None,\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Get the metric value of the channels.\n\n Args:\n inputs (`TensorsCache`):\n The input activations.\n weights (`Sequence[torch.Tensor]`):\n The weight tensors.\n metric_mode (`ChannelOrderCalibConfig.ChannelMetric`):\n The channel metric mode.\n num_channels (`int`):\n The number of channels.\n num_heads (`int`, *optional*, defaults to `1`):\n The number of heads.\n device (`torch.device` or `str` or `None`, *optional*, defaults to `None`):\n The device of the metric value tensor.\n dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The data type of the metric value tensor.\n\n Returns:\n `torch.Tensor`:\n The metric value of the channels.\n \"\"\"\n metric_name = metric_mode.name\n if metric_name.endswith(\"Product\"):\n metric_name = metric_name[:-7]\n ipts_metric = get_channel_metric(\n inputs=inputs,\n weights=weights,\n metric_mode=ChannelOrderCalibConfig.ChannelMetric[f\"Inputs{metric_name}\"],\n num_channels=num_channels,\n num_heads=num_heads,\n device=device,\n dtype=dtype,\n )\n wgts_metric = get_channel_metric(\n inputs=inputs,\n weights=weights,\n metric_mode=ChannelOrderCalibConfig.ChannelMetric[f\"Weights{metric_name}\"],\n num_channels=num_channels,\n num_heads=num_heads,\n device=device,\n dtype=dtype,\n )\n return ipts_metric * wgts_metric\n else:\n if metric_name.startswith(\"Inputs\"):\n assert inputs.num_tensors == 1, f\"Only one input source is allowed, got {inputs.num_tensors}\"\n metric_name, tensors = metric_name[6:], inputs.front().get_standardized_data(reshape=False)\n else:\n assert metric_name.startswith(\"Weights\")\n metric_name, tensors = metric_name[7:], weights\n group_shape = [-1] * tensors[0].ndim\n group_shape[1] = num_channels // num_heads\n # convert metric name from camel case to snake case\n metric_name = \"\".join([\"_\" + c.lower() if c.isupper() else c for c in metric_name])\n metric_name = metric_name.lstrip(\"_\")\n metric_fn = getattr(ChannelMetric, metric_name)\n return metric_fn(tensors, num_channels, group_shape, device=device, dtype=dtype).view(num_channels)\n\n\ndef update_channel_metric(\n metric: torch.Tensor | None,\n inputs: TensorsCache,\n weights: tp.Sequence[torch.Tensor],\n metric_mode: ChannelOrderCalibConfig.ChannelMetric,\n num_channels: int,\n num_heads: int = 1,\n device: torch.device | str = None,\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Update the metric value of the channels.\n\n Args:\n metric (`torch.Tensor` or `None`):\n The metric value of the channels.\n inputs (`TensorsCache`):\n The input activations.\n weights (`Sequence[torch.Tensor]`):\n The weight tensors.\n metric_mode (`ChannelOrderCalibConfig.ChannelMetric`):\n The channel metric mode.\n num_channels (`int`):\n The number of channels.\n num_heads (`int`, *optional*, defaults to `1`):\n The number of heads.\n device (`torch.device` or `str`, *optional*, defaults to `None`):\n The device of the metric value tensor.\n dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The data type of the metric value tensor.\n\n Returns:\n `torch.Tensor`:\n The updated metric value of the channels.\n \"\"\"\n _metric = get_channel_metric(\n inputs=inputs,\n weights=weights,\n metric_mode=metric_mode,\n num_channels=num_channels,\n num_heads=num_heads,\n device=device,\n dtype=dtype,\n )\n if metric is None:\n return _metric\n elif \"Max\" in metric_mode.name:\n return torch.maximum(metric, _metric)\n else:\n return metric.add_(_metric)\n\n\ndef init_channel_index_from_metric(\n metric: torch.Tensor,\n /,\n metric_mode: ChannelOrderCalibConfig.ChannelMetric,\n index_mode: ChannelOrderCalibConfig.ChannelIndex,\n group_size: int,\n num_heads: int = 1,\n num_head_repeats: int = 1,\n) -> torch.Tensor:\n \"\"\"Get the index of the channels.\n\n Args:\n metric (`torch.Tensor`):\n The metric value of the channels.\n metric_mode (`ChannelOrderCalibConfig.ChannelMetric`):\n The channel metric mode.\n index_mode (`ChannelOrderCalibConfig.ChannelIndex`):\n The index mode.\n group_size (`int`):\n The quantization group size.\n num_heads (`int`, *optional*, defaults to `1`):\n The number of heads.\n num_head_repeats (`int`, *optional*, defaults to `1`):\n The number of head repeats.\n\n Returns:\n `torch.Tensor`:\n The index of the channels.\n \"\"\"\n num_channels = metric.numel()\n num_groups = num_channels // group_size\n if num_heads > 1:\n head_channels = num_channels // num_heads\n if num_head_repeats > 1:\n num_unique_heads = num_heads // num_head_repeats\n metric = metric.view(num_unique_heads, num_head_repeats, head_channels)\n metric = metric.amax(dim=1, keepdim=True) if \"Max\" in metric_mode.name else metric.sum(dim=1, keepdim=True)\n rank = metric.argsort(dim=-1).expand(num_unique_heads, num_head_repeats, -1).reshape(num_heads, -1)\n else:\n rank = metric.view(num_heads, head_channels).argsort(dim=-1)\n rank += torch.arange(0, num_channels, head_channels, dtype=torch.long, device=rank.device).view(num_heads, 1)\n index = torch.empty_like(rank)\n for head in range(num_heads):\n index[head] = get_channel_index_from_rank(\n rank[head],\n num_channels=head_channels,\n num_groups=max(num_groups // num_heads, 1),\n index_mode=index_mode,\n )\n return index.view(-1)\n else:\n rank = metric.argsort()\n return get_channel_index_from_rank(\n rank, num_channels=num_channels, num_groups=num_groups, index_mode=index_mode\n )\n\n\nclass ChannelOrderCalibrator(SearchBasedCalibrator[ChannelOrderCalibConfig, torch.Tensor]):\n \"\"\"The calibrator for quantization channel reordering.\"\"\"\n\n def __init__(\n self,\n config: ChannelOrderCalibConfig,\n weight_quantizer: Quantizer | None,\n input_quantizer: Quantizer | None,\n num_heads: int = 1,\n num_head_repeats: int = 1,\n develop_dtype: torch.dtype = torch.float32,\n ) -> None:\n \"\"\"Initialize the calibrator.\n\n Args:\n config (`ChannelOrderCalibConfig`):\n The channel order calibration configuration.\n weight_quantizer (`Quantizer` or `None`):\n The quantizer for the weights.\n input_quantizer (`Quantizer` or `None`):\n The quantizer for the inputs.\n num_heads (`int`, *optional*, defaults to `1`):\n The number of heads.\n num_head_repeats (`int`, *optional*, defaults to `1`):\n The number of head repeats.\n develop_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The development data type.\n \"\"\"\n super().__init__(\n tensor_type=TensorType.Weights,\n config=config,\n w_quantizer=weight_quantizer,\n x_quantizer=input_quantizer,\n y_quantizer=None,\n develop_dtype=develop_dtype,\n )\n assert self.config.objective == SearchBasedCalibObjective.OutputsError\n assert self.config.granularity == SearchBasedCalibGranularity.Layer\n if self.config.strategy == SearchBasedCalibStrategy.Manual:\n self.index_modes = [self.config.channel_index]\n self.metric_modes = [self.config.channel_metric]\n else:\n self.metric_modes = list(ChannelOrderCalibConfig.ChannelMetric.__members__.values())\n self.index_modes = list(ChannelOrderCalibConfig.ChannelIndex.__members__.values())\n self.num_index_modes, self.num_metric_modes = len(self.index_modes), len(self.metric_modes)\n self.num_heads = num_heads\n self.num_head_repeats = num_head_repeats\n self.metrics, self.channel_indexes = None, None\n\n @property\n def population_size(self) -> int:\n \"\"\"Get the population size.\"\"\"\n size = self.num_index_modes * self.num_metric_modes\n return (size + 1) if self.config.strategy != SearchBasedCalibStrategy.Manual else size\n\n @property\n def allows_x_quant_for_wgts(self) -> bool:\n \"\"\"Whether the calibrator needs activation quantization when tensor_type is Weights.\"\"\"\n return self.config.allow_x_quant\n\n @property\n def allows_w_quant_for_wgts(self) -> bool:\n \"\"\"Whether the calibrator needs weight quantization when tensor_type is Weights.\"\"\"\n return self.config.allow_w_quant\n\n def update_channel_metrics(self, weights: list[torch.Tensor | nn.Parameter], inputs: TensorsCache) -> None:\n \"\"\"Update the metrics of the channels.\n\n Args:\n weights (list[torch.Tensor | nn.Parameter]): The weight tensors.\n inputs (TensorsCache): The input activations.\n \"\"\"\n weights = [w.data for w in weights]\n if self.metrics is None:\n self.num_channels = weights[0].shape[1]\n self.device = weights[0].device\n self.metrics = [None] * len(self.metric_modes)\n for metric_id, metric_mode in enumerate(self.metric_modes):\n self.metrics[metric_id] = update_channel_metric(\n metric=self.metrics[metric_id],\n inputs=inputs,\n weights=weights,\n metric_mode=metric_mode,\n num_channels=self.num_channels,\n num_heads=self.num_heads,\n device=self.device,\n dtype=self.develop_dtype,\n )\n\n def init_channel_indexes(self) -> None:\n \"\"\"Initialize the indexes.\"\"\"\n if self.needs_x_quant:\n ipts_group_size = self.x_quantizer.config.smallest_group_shape[1]\n else:\n ipts_group_size = -1\n if ipts_group_size <= 0:\n ipts_group_size = self.num_channels\n if self.needs_w_quant:\n wgts_group_size = self.w_quantizer.config.smallest_group_shape[1]\n else:\n wgts_group_size = -1\n if wgts_group_size <= 0:\n wgts_group_size = self.num_channels\n group_size = min(ipts_group_size, wgts_group_size)\n self.channel_indexes = [None] + [\n init_channel_index_from_metric(\n metric,\n metric_mode=metric_mode,\n index_mode=index_mode,\n group_size=group_size,\n num_heads=self.num_heads,\n num_head_repeats=self.num_head_repeats,\n )\n for metric_mode, metric in zip(self.metric_modes, self.metrics, strict=True)\n for index_mode in self.index_modes\n ]\n self.arange = torch.arange(self.num_channels, dtype=torch.long, device=self.device)\n self.metrics = None\n gc.collect()\n torch.cuda.empty_cache()\n\n def _reset(self, x_wgts: list[torch.Tensor | nn.Parameter], x_acts: TensorsCache, **kwargs) -> None:\n \"\"\"Reset the calibrator.\n\n Args:\n x_wgts (list[list[torch.Tensor | nn.Parameter]]): Weight tensors.\n x_acts (TensorsCache): Input activations.\n \"\"\"\n if self.channel_indexes is None:\n self.update_channel_metrics(x_wgts, x_acts)\n self.init_channel_indexes()\n if self.config.strategy == SearchBasedCalibStrategy.Manual and self.channel_indexes[0] is None:\n self.channel_indexes = self.channel_indexes[1:]\n assert len(self.channel_indexes) == self.population_size\n self.baseline_errors, self.best_error, self.best_candidate_id = None, None, None\n self.error_stats_history = []\n\n def get_best(self) -> torch.Tensor:\n \"\"\"Get the best candidate.\n\n Returns:\n torch.Tensor: The best candidate.\n \"\"\"\n return self.channel_indexes[self.best_candidate_id]\n\n def _ask(self) -> torch.Tensor:\n \"\"\"Ask for the next candidate.\n\n Returns:\n torch.Tensor: The next candidate.\n \"\"\"\n channel_index = self.channel_indexes[self.candidate_id]\n channel_index_inverse = None\n if channel_index is not None:\n channel_index_inverse = torch.zeros_like(channel_index)\n channel_index_inverse[channel_index] = self.arange.to(device=channel_index.device)\n self.candidate_inverse = channel_index_inverse\n return channel_index\n\n def _tell(self, errors: list[tuple[torch.Tensor, ...]]) -> None: # noqa: C901\n \"\"\"Tell the error of the last candidate and update the best candidate.\n\n Args:\n errors (list[tuple[torch.Tensor, ...]]): The error of the last candidate.\n \"\"\"\n errors = [tuple(math.root_(e.to(torch.float64), self.config.degree) for e in error) for error in errors]\n if self.baseline_errors is None:\n self.baseline_errors = errors\n error_stats = [0, 0, 0, 0, 0]\n for baseline_error, error in zip(self.baseline_errors, errors, strict=True):\n for be, e in zip(baseline_error, error, strict=True):\n _d = e.item() - be.item()\n if e > be:\n error_stats[0] += 1\n if e < be:\n error_stats[1] -= 1\n error_stats[2] += max(_d, 0)\n error_stats[3] += min(_d, 0)\n error_stats[4] += e.item()\n if self.best_error is None or error_stats < self.best_error:\n self.best_error = error_stats\n self.best_candidate_id = self.candidate_id\n if self.logger.level <= tools.logging.DEBUG:\n self.logger.debug(\n f\"+ {self._get_metric_index_mode_str(self.candidate_id)} : {self._get_error_str(error_stats)}\"\n )\n if self.is_last_candidate_in_iter():\n self.logger.debug(f\"+ {self._get_metric_index_mode_str(self.best_candidate_id)} is the best candidate.\")\n\n def _get_error_str(self, e: list[int | float]) -> str:\n return f\"[{e[0]:+d}, {e[1]:+d}, {e[2]:>10.4f}, {e[3]:>10.4f}, {e[4]:>10.4f}]\"\n\n def _get_metric_index_mode_str(self, candidate_id: int) -> str:\n if candidate_id == 0:\n if self.config.strategy == SearchBasedCalibStrategy.Manual:\n metric_mode, index_mode = self.metric_modes[0], self.index_modes[0]\n else:\n return f\"{'baseline':>20} {'':>10}\"\n else:\n metric_id = (candidate_id - 1) % self.num_metric_modes\n index_id = (candidate_id - 1) // self.num_metric_modes\n metric_mode, index_mode = self.metric_modes[metric_id], self.index_modes[index_id]\n return f\"{metric_mode.name:>20} - {index_mode.name:>10}\"\n\n def _process_x_in_xw(self, x: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor:\n if not self.needs_x_quant_for_wgts:\n return x\n if channels_dim is MISSING:\n channels_dim = self.x_quantizer.channels_dim\n if self.candidate is not None:\n x = x.index_select(dim=channels_dim, index=self.candidate.to(x.device))\n x = self.x_quantizer.quantize(x, channels_dim=channels_dim).data\n if self.candidate is not None:\n x = x.index_select(dim=channels_dim, index=self.candidate_inverse.to(x.device))\n return x\n\n def _process_w_in_xw(self, w: torch.Tensor) -> torch.Tensor:\n if not self.needs_w_quant_for_wgts:\n return w\n if self.candidate is not None:\n w = w.index_select(dim=1, index=self.candidate.to(w.device))\n w = self.w_quantizer.quantize(w.data, kernel=None, develop_dtype=self.develop_dtype).data\n if self.candidate is not None:\n w = w.index_select(dim=1, index=self.candidate_inverse.to(w.device))\n return w\n\n def _process_x_in_yx(self, x: torch.Tensor, channels_dim: int) -> torch.Tensor:\n raise RuntimeError(\"_process_x_in_yx should not be called in ChannelOrderCalibrator.\")\n\n def _process_y_in_yx(self, x: torch.Tensor, channels_dim: int) -> torch.Tensor:\n raise RuntimeError(\"_process_y_in_yx should not be called in ChannelOrderCalibrator.\")\n\n def _process_xw_in_yx(self, w: torch.Tensor) -> torch.Tensor:\n raise RuntimeError(\"_process_xw_in_yx should not be called in ChannelOrderCalibrator.\")\n\n def _process_yw_in_yx(self, w: torch.Tensor) -> torch.Tensor:\n raise RuntimeError(\"_process_yw_in_yx should not be called in ChannelOrderCalibrator.\")\n\n def _process_wgts_centric_mod(\n self,\n wgts: list[nn.Parameter],\n mods: list[nn.Module],\n *,\n reorder_wgts: list[tuple[nn.Parameter, int]],\n reorder_ipt_mods: list[tuple[nn.Module, int, BaseInputPackager | None]],\n reorder_opt_mods: list[tuple[nn.Module, int, BaseOutputPackager | None]],\n update_state_dict: bool = True,\n **kwargs,\n ) -> None:\n channels_index = self.candidate\n if update_state_dict:\n self._state_dict.extend([(w, w.data) for w, _ in reorder_wgts])\n if channels_index is not None:\n for w, d in reorder_wgts:\n w.data = w.data.index_select(dim=d, index=channels_index.to(w.device))\n for m, channels_dim, packager in reorder_ipt_mods:\n self._hooks.append(\n ChannelReorderer(channels_index, channels_dim, input_packager=packager).as_hook().register(m)\n )\n for m, channels_dim, packager in reorder_opt_mods:\n self._hooks.append(\n ChannelReorderer(channels_index, channels_dim, output_packager=packager)\n .as_hook(is_output=True)\n .register(m)\n )\n self._candidate_backup = channels_index\n self.candidate = None # we have already reordered and thus do not need to reorder again in _process\n super()._process_wgts_centric_mod(wgts, mods, update_state_dict=False)\n\n def _recover_mod(self) -> None:\n super()._recover_mod()\n self.candidate = self._candidate_backup\n self._candidate_backup = None\n"
},
{
"path": "deepcompressor/calib/rotate.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Rotation Quantization module.\"\"\"\n\nimport typing as tp\n\nimport torch\nimport torch.nn as nn\n\nfrom ..utils.hooks import BaseInputPackager, IOHook\nfrom ..utils.math import HadamardMatrix, hardmard_transform, random_hadamard_matrix\n\n__all__ = [\n \"rotate_in_channels\",\n \"rotate_out_channels\",\n \"hadamard_in_channels\",\n \"get_rotation_matrix\",\n \"transform_rms_norm_and_linear\",\n \"transform_layer_norm_to_rms_norm\",\n \"transform_norm_and_linear\",\n]\n\n\n# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm\nclass RMSNorm(nn.Module):\n \"\"\"Root Mean Square Layer Normalization (RMSNorm).\"\"\"\n\n def __init__(self, hidden_size: int, eps=1e-6) -> None:\n \"\"\"Initialize RMSNorm.\"\"\"\n super().__init__()\n self.weight = nn.Parameter(torch.ones(hidden_size))\n self.variance_epsilon = eps\n\n def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:\n \"\"\"Apply RMSNorm normalization to hidden states.\"\"\"\n input_dtype = hidden_states.dtype\n hidden_states = hidden_states.to(torch.float32)\n variance = hidden_states.pow(2).mean(-1, keepdim=True)\n hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)\n return self.weight * hidden_states.to(input_dtype)\n\n\nclass HadamardTransformHook(IOHook):\n def __init__(\n self, rhs: torch.Tensor, lhs: torch.Tensor, lhs_k: int, scaled: bool = True, packager: BaseInputPackager = None\n ):\n super().__init__(pre=True, post=False, input_packager=packager, output_packager=None)\n self.rhs = rhs\n self.lhs = lhs\n self.lhs_k = lhs_k\n self.scaled = scaled\n\n def pre_forward(\n self,\n module: nn.Module,\n input_args: tuple[torch.Tensor, ...],\n input_kwargs: dict[str, tp.Any],\n ) -> tuple[tuple[torch.Tensor, ...], dict[str, tp.Any]]:\n tensors = self.input_packager.unpack(module, input_args, input_kwargs)\n for k, x in tensors.items():\n tensors[k] = hardmard_transform(\n x, hadamard_rhs=self.rhs, hadamard_lhs=self.lhs, lhs_k=self.lhs_k, scaled=self.scaled\n )\n return self.input_packager.repack(tensors, module, input_args, input_kwargs)\n\n\ndef rotate_in_channels(weight: nn.Parameter, /, *, rotation: torch.Tensor) -> None:\n \"\"\"Rotate the input channels of a weight matrix.\"\"\"\n shape, dtype = weight.shape, weight.dtype\n weight.data = (\n torch.matmul(weight.data.view(-1, rotation.shape[0]).to(dtype=torch.float64), rotation.to(weight.device))\n .to(dtype=dtype)\n .view(shape)\n )\n\n\ndef rotate_out_channels(weight: nn.Parameter, /, *, rotation: torch.Tensor, bias: nn.Parameter | None = None) -> None:\n \"\"\"Rotate the output channels of a weight matrix.\"\"\"\n shape, dtype = weight.shape, weight.dtype\n out_channels, head_channels = shape[0], rotation.shape[0]\n num_heads = out_channels // head_channels\n weight.data = (\n torch.matmul(\n rotation.T.to(weight.device), weight.data.view(num_heads, head_channels, -1).to(dtype=torch.float64)\n )\n .to(dtype=dtype)\n .view(shape)\n )\n if bias is not None:\n bias.data = (\n torch.matmul(\n rotation.T.to(weight.device), bias.data.view(num_heads, head_channels, -1).to(dtype=torch.float64)\n )\n .to(dtype=dtype)\n .view(-1)\n )\n\n\ndef hadamard_in_channels(\n modules: tp.Iterable[nn.Module],\n packager: BaseInputPackager = None,\n dtype: torch.dtype | None = None,\n device: torch.device | str | None = None,\n):\n \"\"\"Apply Hadamard quantization to the input channels of the modules.\"\"\"\n for module in modules:\n if isinstance(module, nn.Linear):\n in_channels = module.in_features\n device, dtype = device or module.weight.device, dtype or module.weight.dtype\n rhs_double, lhs_double, k = HadamardMatrix.get(in_channels, scale=True, dtype=torch.float64)\n module.weight.data = hardmard_transform(\n module.weight.data.to(torch.float64), rhs_double.to(device), lhs_double.to(device), k, scaled=True\n ).to(device=device, dtype=module.weight.dtype)\n del rhs_double, lhs_double, k\n rhs, lhs, k = HadamardMatrix.get(in_channels, scale=True, dtype=dtype, device=device)\n HadamardTransformHook(rhs=rhs, lhs=lhs, lhs_k=k, packager=packager).register(module)\n else:\n raise NotImplementedError(f\"Module {module} not supported!\")\n\n\ndef get_rotation_matrix(num_channels: int, random: bool = True, compatible: bool = True) -> torch.Tensor:\n \"\"\"Get a random rotation matrix for the given number of channels.\"\"\"\n if random:\n return random_hadamard_matrix(num_channels)\n else:\n rhs, lhs, k = HadamardMatrix.get(num_channels, scale=False)\n rhs = rhs.to(dtype=torch.float64)\n if k == 1:\n rotation = rhs\n elif compatible: # this is compatible with hadamard_transform\n rotation = torch.kron(lhs.T.contiguous().to(dtype=torch.float64), rhs)\n else:\n rotation = torch.kron(rhs, lhs.to(dtype=torch.float64))\n return rotation.mul_(1.0 / torch.tensor(num_channels, dtype=torch.float64).sqrt())\n\n\ndef transform_rms_norm_and_linear(norm: nn.LayerNorm | RMSNorm, next_modules: tp.Iterable[nn.Linear]) -> None:\n \"\"\"Fuse the weight multiplication of rms norm into the next adjacent linear modules.\n\n Args:\n norm (`nn.LayerNorm` or `RMSNorm`):\n normalization module.\n next_modules (`Iterable[nn.Linear]`):\n modules after the normalization module.\n \"\"\"\n ln_w = norm.weight.data.to(dtype=torch.float64)\n norm.weight.data = torch.ones_like(norm.weight.data)\n if hasattr(norm, \"bias\") and norm.bias is not None:\n ln_b = norm.bias.data.to(dtype=torch.float64)\n norm.bias = None\n else:\n ln_b = None\n for linear in next_modules:\n assert isinstance(linear, nn.Linear)\n dtype = linear.weight.dtype\n fc_w = linear.weight.data.to(dtype=torch.float64)\n ln_w = ln_w.to(fc_w.device)\n linear.weight.data = (fc_w * ln_w).to(dtype=dtype)\n if ln_b is not None:\n ln_b = ln_b.to(fc_w.device)\n if linear.bias is None:\n linear.bias = nn.Parameter(torch.zeros(linear.out_features, dtype=dtype, device=linear.weight.device))\n linear.bias.data = (linear.bias.data.to(dtype=torch.float64) + torch.matmul(fc_w, ln_b)).to(dtype=dtype)\n\n\ndef transform_layer_norm_to_rms_norm(\n parent: nn.Module,\n norm_name: str,\n prev_modules: tp.Iterable[nn.Linear],\n prev_out_channels_dims: int | tp.Iterable[int] = 0,\n) -> None:\n \"\"\"Transform LayerNorm to RMSNorm.\n\n Args:\n parent (`nn.Module`):\n Parent module that contains the normalization module.\n norm_name (`str`):\n Name of the normalization module in `parent`.\n prev_modules (`Iterable[nn.Linear]`):\n Previous adjacent linear modules.\n prev_out_channels_dims (`int` or `Iterable[int]`, *optional*, defaults to `0`):\n Output channels dimension of the previous modules' weights.\n \"\"\"\n if \".\" in norm_name:\n norm_names = norm_name.split(\".\")\n for name in norm_names[:-1]:\n parent = getattr(parent, name)\n norm_name = norm_names[-1]\n del norm_names\n norm = getattr(parent, norm_name)\n assert isinstance(norm, nn.LayerNorm)\n assert len(norm.normalized_shape) == 1, f\"LayerNorm's #dims must be 1, got {len(norm.normalized_shape)}\"\n assert norm.bias is None, \"LayerNorm's bias must be None, please call `transform_rms_norm_and_linear` in advance\"\n # region move substract mean to the previous linear modules\n assert len(prev_modules) > 0, \"No previous modules found\"\n if isinstance(prev_out_channels_dims, int):\n prev_out_channels_dims = [prev_out_channels_dims] * len(prev_modules)\n for module, dim in zip(prev_modules, prev_out_channels_dims, strict=True):\n if isinstance(module, nn.LayerNorm):\n module.bias = None\n else:\n if isinstance(module, nn.Linear):\n assert dim == 0, \"Linear module's output channels dimension is 0\"\n elif isinstance(module, nn.Embedding):\n assert dim == 1, \"Embedding module's output channels dimension is 1\"\n dtype = module.weight.dtype\n w = module.weight.data.to(dtype=torch.float64)\n module.weight.data = w.sub_(w.mean(dim=dim, keepdim=True)).to(dtype=dtype)\n if hasattr(module, \"bias\") and module.bias is not None:\n b = module.bias.data.to(dtype=torch.float64)\n module.bias.data = b.sub_(b.mean()).to(dtype=dtype)\n # endregion\n # region replace LayerNorm with RMSNorm\n rms = RMSNorm(hidden_size=norm.normalized_shape[0], eps=norm.eps)\n rms.weight.data = norm.weight.data\n setattr(parent, norm_name, rms)\n # endregion\n\n\ndef transform_norm_and_linear(\n parent: nn.Module,\n norm_name: str,\n next_modules: tp.Iterable[nn.Linear],\n prev_modules: tp.Iterable[nn.Linear] | None = None,\n prev_out_channels_dims: int | tp.Iterable[int] = 0,\n):\n \"\"\"Transform the normalization module and the next adjacent linear modules.\n\n Args:\n parent (nn.Module): Parent module.\n norm_name (str): Name of the normalization module.\n next_modules (tp.Iterable[nn.Linear]): Next adjacent linear modules.\n prev_modules (tp.Iterable[nn.Linear]): Previous adjacent linear modules.\n prev_out_channels_dims (int | tp.Iterable[int], optional): Output channels dimension of the previous modules.\n Defaults to ``0``.\n \"\"\"\n if \".\" in norm_name:\n norm_names = norm_name.split(\".\")\n for name in norm_names[:-1]:\n parent = getattr(parent, name)\n norm_name = norm_names[-1]\n del norm_names\n norm = getattr(parent, norm_name)\n transform_rms_norm_and_linear(norm, next_modules)\n if isinstance(norm, nn.LayerNorm):\n transform_layer_norm_to_rms_norm(parent, norm_name, prev_modules, prev_out_channels_dims)\n"
},
{
"path": "deepcompressor/calib/search.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Search-based uantization calibrator module.\"\"\"\n\nimport gc\nimport typing as tp\nfrom abc import ABC, abstractmethod\nfrom dataclasses import _MISSING_TYPE, MISSING\n\nimport psutil\nimport torch\nimport torch.nn as nn\nimport torch.utils.hooks\n\nfrom ..data.cache import TensorCache, TensorsCache\nfrom ..data.common import TensorType\nfrom ..data.utils.reshape import ReshapeFn\nfrom ..data.utils.shape import infer_view_shape\nfrom ..quantizer.processor import Quantizer\nfrom ..utils import tools\nfrom ..utils.hooks import Hook\nfrom .config import SearchBasedCalibConfig, SearchBasedCalibGranularity, SearchBasedCalibObjective\n\n__all__ = [\"SearchBasedCalibrator\"]\n\n\ndef _reshape_w_for_wgts(w: torch.Tensor, w_view_shape: torch.Size) -> torch.Tensor:\n # (#g0, gs0, #g1, gs1, ...)\n w = w.view(w_view_shape)\n # (#g0, gs0, #g1, gs1, ...) -> (#g0, ..., gs1, ..., gs0)\n w = w.permute(*range(0, len(w_view_shape), 2), *range(3, len(w_view_shape), 2), 1)\n # (#g0, ..., gs0, gs1, ...) -> (#g0, ..., gs1 * gs2 * ..., gs0)\n return w.reshape(*w_view_shape[::2], -1, w_view_shape[1])\n\n\ndef _reshape_x_for_wgts(x: torch.Tensor, w_view_shape: torch.Size) -> torch.Tensor:\n # x is unfolded already\n num_samples = x.shape[0]\n # (1, n, #g1, gs1, ...)\n x = x.view(1, num_samples, *w_view_shape[2:])\n # (1, n, #g1, gs1, ...) -> (1, #g1, ..., n, gs1, ...)\n x = x.permute(*range(0, len(w_view_shape), 2), *range(1, len(w_view_shape), 2))\n return x.reshape(1, *w_view_shape[2::2], num_samples, -1)\n\n\ndef _reshape_x_for_ipts(x: torch.Tensor, x_view_shape: torch.Size) -> torch.Tensor:\n # x is original tensor without unfolding\n # (#g0, gs0, #g1, gs1, ...)\n x = x.view(x_view_shape)\n # (#g0, gs0, #g1, gs1, ...) -> (#g0, #g1, ..., gs0, gs2, ..., gs1)\n x = x.permute(*range(0, len(x_view_shape), 2), 1, *range(5, len(x_view_shape), 2), 3)\n # (#g0, #g1, ..., gs0, gs2, ..., gs1) -> (#g0, #g1, ..., gs0 * gs2 * ..., gs1)\n return x.reshape(*x_view_shape[::2], -1, x_view_shape[3])\n\n\ndef _reshape_w_for_ipts(w: torch.Tensor, x_view_shape: torch.Size) -> torch.Tensor:\n return w.transpose(0, 1).reshape(1, x_view_shape[2], *([1] * (w.ndim - 2)), x_view_shape[3], -1)\n\n\n_CANDIDATE = tp.TypeVar(\"_CANDIDATE\")\n_CONFIG = tp.TypeVar(\"_CONFIG\", bound=SearchBasedCalibConfig)\n\n\nclass SearchBasedCalibrator(ABC, tp.Generic[_CONFIG, _CANDIDATE]):\n \"\"\"The base class for search-based calibration.\"\"\"\n\n config: _CONFIG\n candidate: _CANDIDATE\n\n def __init__(\n self,\n tensor_type: TensorType,\n config: _CONFIG,\n w_quantizer: Quantizer | None,\n x_quantizer: Quantizer | None,\n y_quantizer: Quantizer | None,\n develop_dtype: torch.dtype,\n ) -> None:\n \"\"\"Initialize the search-based calibrator.\n\n Args:\n tensor_type (`TensorType`):\n The tensor type.\n config (`_CONFIG`):\n The calibration configuration.\n w_quantizer (`Quantizer` or `None`):\n The w quantizer for x-w computation.\n x_quantizer (`Quantizer` or `None`):\n The x quantizer for x-w or y-x computation.\n y_quantizer (`Quantizer` or `None`):\n The y quantizer for y-x computation.\n develop_dtype (`torch.dtype`):\n The development data type.\n \"\"\"\n self.tensor_type = tensor_type\n self.config = config\n self.objective = self.config.objective\n self.granularity = self.config.granularity\n self.opts_device = None\n self.develop_dtype = develop_dtype\n self.w_quantizer = w_quantizer\n self.x_quantizer = x_quantizer\n self.y_quantizer = y_quantizer\n self.needs_w_quant = self.w_quantizer is not None and self.w_quantizer.is_enabled()\n self.needs_x_quant = self.x_quantizer is not None and self.x_quantizer.is_enabled()\n self.needs_y_quant = self.y_quantizer is not None and self.y_quantizer.is_enabled()\n self.needs_x_quant_for_wgts = self.allows_x_quant_for_wgts and self.needs_x_quant\n self.needs_w_quant_for_wgts = self.allows_w_quant_for_wgts and self.needs_w_quant\n self.needs_x_quant_for_ipts = self.allows_x_quant_for_ipts and self.needs_x_quant\n self.needs_w_quant_for_ipts = self.allows_w_quant_for_ipts and self.needs_w_quant\n self.needs_x_quant_for_opts = self.allows_x_quant_for_opts and self.needs_x_quant\n self.needs_y_quant_for_opts = self.allows_y_quant_for_opts and self.needs_y_quant\n self.needs_w_quant_for_opts = self.allows_w_quant_for_opts and self.needs_w_quant\n if self.tensor_type == TensorType.Weights:\n self.quantizer = self.w_quantizer\n self.needs_quant = self.needs_w_quant\n elif self.tensor_type == TensorType.Inputs:\n self.quantizer = self.x_quantizer\n self.needs_quant = self.needs_x_quant\n elif self.tensor_type == TensorType.Outputs:\n self.quantizer = self.y_quantizer\n self.needs_quant = self.needs_y_quant\n else:\n raise ValueError(f\"unknown tensor type: {self.tensor_type}\")\n self.num_iters = getattr(self.config, \"num_iters\", 1)\n self.logger = tools.logging.getLogger(f\"{__name__}.{self.__class__.__name__.replace('Agent', '')}\")\n\n @property\n @abstractmethod\n def population_size(self) -> int:\n \"\"\"Get the population size.\"\"\"\n ...\n\n @property\n def allows_x_quant_for_wgts(self) -> bool:\n \"\"\"Whether the calibrator allows input quantization when tensor_type is Weights.\"\"\"\n return False\n\n @property\n def allows_w_quant_for_wgts(self) -> bool:\n \"\"\"Whether the calibrator allows weight quantization when tensor_type is Weights.\"\"\"\n return True\n\n @property\n def allows_x_quant_for_ipts(self) -> bool:\n \"\"\"Whether the calibrator allows input quantization when tensor_type is Inputs.\"\"\"\n return True\n\n @property\n def allows_w_quant_for_ipts(self) -> bool:\n \"\"\"Whether the calibrator allows weight quantization when tensor_type is Inputs.\"\"\"\n return False\n\n @property\n def allows_x_quant_for_opts(self) -> bool:\n \"\"\"Whether the calibrator allows x quantization when tensor_type is Outputs.\"\"\"\n return True\n\n @property\n def allows_y_quant_for_opts(self) -> bool:\n \"\"\"Whether the calibrator allows y quantization when tensor_type is Outputs.\"\"\"\n return True\n\n @property\n def allows_w_quant_for_opts(self) -> bool:\n \"\"\"Whether the calibrator allows weight quantization when tensor_type is Outputs.\"\"\"\n return False\n\n @property\n def needs_to_pre_reshape_x_for_wgts(self) -> bool:\n \"\"\"Whether the calibrator needs to pre-reshape the inputs for weight quantization calibration.\"\"\"\n return not self.needs_x_quant_for_wgts and self.config.pre_reshape\n\n @property\n def needs_to_pre_reshape_w_for_ipts(self) -> bool:\n \"\"\"Whether the calibrator needs to pre-reshape the weights for input quantization calibration.\"\"\"\n return not self.needs_w_quant_for_ipts and self.config.pre_reshape\n\n def _reset(self, **kwargs) -> None:\n pass\n\n def reset(self, **kwargs) -> None:\n \"\"\"Reset the calibrator.\"\"\"\n self.iter = 0\n self.candidate_id = 0\n self._reset(**kwargs)\n self._state_dict: list[tuple[nn.Parameter, torch.Tensor]] = []\n self._hooks: list[Hook | torch.utils.hooks.RemovableHandle] = []\n\n def is_done(self) -> bool:\n \"\"\"Check if the calibration is done.\"\"\"\n return self.iter >= self.num_iters\n\n def is_last_iter(self) -> bool:\n \"\"\"Check if the current iteration is the last one.\"\"\"\n return self.iter == self.num_iters - 1\n\n def is_last_candidate_in_iter(self) -> bool:\n \"\"\"Check if the current candidate is the last one in the current iteration.\"\"\"\n return self.candidate_id == self.population_size - 1\n\n @abstractmethod\n def get_best(self) -> _CANDIDATE:\n \"\"\"Get the best candidate.\n\n Returns:\n `_CANDIDATE`:\n The best candidate.\n \"\"\"\n ...\n\n @abstractmethod\n def _ask(self) -> _CANDIDATE:\n \"\"\"Ask for the next candidate.\n\n Returns:\n `_CANDIDATE`:\n The next candidate.\n \"\"\"\n ...\n\n @abstractmethod\n def _tell(self, error: list[torch.Tensor]) -> None:\n \"\"\"Tell the error of the last candidate and update the best candidate.\n\n Args:\n error (`list[torch.Tensor]`):\n The error of the last candidate.\n \"\"\"\n ...\n\n def ask(self) -> _CANDIDATE:\n \"\"\"Ask for the next candidate.\n\n Returns:\n `_CANDIDATE`:\n The next candidate.\n \"\"\"\n self.candidate = self._ask()\n return self.candidate\n\n def tell(self, error: list[torch.Tensor]) -> None:\n \"\"\"Tell the error of the last candidate and update the best candidate.\n\n Args:\n error (`list[torch.Tensor]`):\n The error of the last candidate.\n \"\"\"\n self._tell(error)\n self.candidate_id += 1\n if self.candidate_id >= self.population_size:\n self.iter += 1\n self.candidate_id = 0\n\n def _parse_ipts(self, ipts: TensorsCache | None, set_device: bool = False) -> TensorsCache | None:\n if set_device:\n self.opts_device = None\n elif ipts is None:\n return None\n if self.objective == SearchBasedCalibObjective.ProductsError:\n batch_size = self.config.element_batch_size\n calib_size = self.config.element_size\n elif self.objective == SearchBasedCalibObjective.OutputsError:\n batch_size = self.config.sample_batch_size\n calib_size = self.config.sample_size\n else:\n assert self.objective == SearchBasedCalibObjective.TensorError\n batch_size = -1\n calib_size = -1\n prev_size = len(ipts.front().data)\n parsed_ipts = TensorsCache(\n {\n key: ipt.repartition(\n max_batch_size=batch_size,\n max_size=calib_size,\n standardize=self.objective == SearchBasedCalibObjective.ProductsError,\n reshape=self.tensor_type == TensorType.Weights,\n )\n for key, ipt in ipts.items()\n }\n )\n curr_size = len(parsed_ipts.front().data)\n assert all(len(ipt.data) == curr_size for ipt in parsed_ipts.values())\n if set_device and prev_size != curr_size:\n self.opts_device = self.config.outputs_device\n return parsed_ipts\n\n def _parse_args( # noqa: C901\n self,\n x_wgts: list[nn.Parameter] | None,\n y_wgts: list[nn.Parameter] | None,\n x_acts: TensorsCache | None,\n y_acts: TensorsCache | None,\n eval_inputs: TensorsCache | None,\n eval_module: nn.Module | None,\n x_mods: list[nn.Module] | None,\n y_mods: list[nn.Module] | None,\n orig_x_wgts: list[tuple[nn.Parameter, torch.Tensor]] | None,\n orig_y_wgts: list[tuple[nn.Parameter, torch.Tensor]] | None,\n orig_x_acts: TensorsCache | None,\n orig_y_acts: TensorsCache | None,\n orig_eval_inputs: TensorsCache | None,\n ) -> tuple[\n list[torch.Tensor | nn.Parameter] | None, # x_wgts\n list[torch.Tensor | nn.Parameter] | None, # y_wgts\n TensorsCache | None, # x_acts\n TensorsCache | None, # y_acts\n TensorsCache | None, # eval_inputs\n nn.Module | None, # eval_module\n list[nn.Module] | None, # x_mods\n list[nn.Module] | None, # y_mods\n list[tuple[nn.Parameter, torch.Tensor]] | None, # orig_x_wgts\n list[tuple[nn.Parameter, torch.Tensor]] | None, # orig_y_wgts\n TensorCache | None, # orig_x_acts\n TensorCache | None, # orig_y_acts\n TensorCache | None, # orig_eval_inputs\n ]:\n # region Check the types of the arguments\n if x_wgts is not None:\n assert isinstance(x_wgts, (tuple, list)), \"x_wgts should be a list\"\n assert all(isinstance(w, nn.Parameter) for w in x_wgts), \"wgts should be a list of nn.Parameter\"\n if y_wgts is not None:\n assert isinstance(y_wgts, (tuple, list)), \"y_wgts should be a list\"\n assert all(isinstance(w, nn.Parameter) for w in y_wgts), \"wgts should be a list of nn.Parameter\"\n if x_acts is not None:\n assert isinstance(x_acts, TensorsCache), \"x_acts should be a TensorsCache\"\n if y_acts is not None:\n assert isinstance(y_acts, TensorsCache), \"y_acts should be a TensorsCache\"\n if eval_inputs is not None:\n assert isinstance(eval_inputs, TensorsCache), \"eval_inputs should be a TensorsCache\"\n if x_mods is not None:\n assert isinstance(x_mods, (tuple, list)), \"x_mods should be a list\"\n if y_mods is not None:\n assert isinstance(y_mods, (tuple, list)), \"y_mods should be a list\"\n if orig_x_wgts is not None:\n assert isinstance(orig_x_wgts, (tuple, list)), \"orig_x_wgts should be a list\"\n assert all(isinstance(p, nn.Parameter) and isinstance(w, torch.Tensor) for p, w in orig_x_wgts), (\n \"orig_x_wgts should be a list of tuples of nn.Parameter and torch.Tensor\"\n )\n if x_wgts is not None:\n assert len(orig_x_wgts) >= len(x_wgts), \"orig_wgts should have at least as mtp.Any elements as wgts\"\n assert all(p is w for (p, _), w in zip(orig_x_wgts, x_wgts, strict=False)), (\n \"the parameters in orig_wgts should be in wgts in the same order\"\n )\n if orig_y_wgts is not None:\n assert isinstance(orig_y_wgts, (tuple, list)), \"orig_y_wgts should be a list\"\n assert all(isinstance(p, nn.Parameter) and isinstance(w, torch.Tensor) for p, w in orig_y_wgts), (\n \"orig_y_wgts should be a list of tuples of nn.Parameter and torch.Tensor\"\n )\n if y_wgts is not None:\n assert len(orig_y_wgts) >= len(y_wgts), \"orig_wgts should have at least as mtp.Any elements as wgts\"\n assert all(p is w for (p, _), w in zip(orig_y_wgts, y_wgts, strict=False)), (\n \"the parameters in orig_wgts should be in wgts in the same order\"\n )\n if orig_x_acts is not None:\n assert isinstance(orig_x_acts, TensorsCache), \"orig_x_acts should be a TensorsCache\"\n if orig_y_acts is not None:\n assert isinstance(orig_y_acts, TensorsCache), \"orig_y_acts should be a TensorsCache\"\n if orig_eval_inputs is not None:\n assert isinstance(orig_eval_inputs, TensorsCache), \"orig_eval_inputs should be a TensorsCache\"\n # endregion\n self.objective = self.config.objective\n self.granularity = self.config.granularity\n if self.tensor_type == TensorType.Outputs:\n # ! currently only support OutputsError and Layer granularity for Outputs\n self.objective = SearchBasedCalibObjective.OutputsError\n self.granularity = SearchBasedCalibGranularity.Layer\n if self.objective == SearchBasedCalibObjective.TensorError:\n if x_wgts is not None:\n x_wgts = [w.detach().data for w in x_wgts]\n if y_wgts is not None:\n y_wgts = [w.detach().data for w in y_wgts]\n if self.tensor_type == TensorType.Weights:\n assert x_wgts is not None, \"wgts should not be None when tensor_type is Weights\"\n elif self.tensor_type == TensorType.Inputs:\n assert x_acts is not None, \"mod_ipts should not be None when tensor_type is Inputs\"\n eval_inputs, orig_eval_inputs = x_acts, orig_x_acts\n else: # self.tensor_type == TensorType.Outputs\n assert y_acts is not None, \"opts should not be None when tensor_type is Outputs\"\n eval_inputs, orig_eval_inputs = y_acts, orig_y_acts\n eval_module = None\n elif self.objective == SearchBasedCalibObjective.ProductsError:\n assert self.tensor_type in (\n TensorType.Weights,\n TensorType.Inputs,\n ), \"tensor_type should be Weights or Inputs when objective is ProductsError\"\n assert x_wgts is not None, \"wgts should not be None when objective is ProductsError\"\n x_wgts = [w.detach().data for w in x_wgts]\n if y_wgts is not None:\n y_wgts = [w.detach().data for w in y_wgts]\n x_acts = x_acts or eval_inputs\n orig_x_acts = orig_x_acts or orig_eval_inputs\n assert x_acts is not None, \"x_acts should not be None when objective is ProductsError\"\n eval_inputs, orig_eval_inputs = x_acts, orig_x_acts\n elif self.objective == SearchBasedCalibObjective.OutputsError:\n assert eval_inputs is not None, \"eval_inputs should not be None when objective is OutputsError\"\n assert eval_module is not None, \"eval_module should not be None when OutputsError\"\n if (\n isinstance(eval_module, (nn.Linear, nn.Conv2d))\n and self.granularity.value < SearchBasedCalibGranularity.Layer.value\n and self.tensor_type != TensorType.Outputs\n ):\n self.objective = SearchBasedCalibObjective.ProductsError\n x_wgts = [w.detach().data for w in x_wgts]\n if y_wgts is not None:\n y_wgts = [w.detach().data for w in y_wgts]\n x_acts = x_acts or eval_inputs\n orig_x_acts = orig_x_acts or orig_eval_inputs\n assert x_acts is not None, \"x_acts should not be None when objective is ProductsError\"\n eval_inputs, orig_eval_inputs = x_acts, orig_x_acts\n else:\n self.objective = SearchBasedCalibObjective.OutputsError\n self.granularity = SearchBasedCalibGranularity.Layer\n else:\n raise ValueError(f\"unknown objective: {self.objective}\")\n self.logger.debug(\n f\"+ tensor_type: {self.tensor_type}, objective: {self.objective}, granularity: {self.granularity}\"\n )\n return (\n x_wgts,\n y_wgts,\n x_acts,\n y_acts,\n self._parse_ipts(eval_inputs, set_device=True),\n eval_module,\n x_mods,\n y_mods,\n orig_x_wgts,\n orig_y_wgts,\n orig_x_acts,\n orig_y_acts,\n self._parse_ipts(orig_eval_inputs),\n )\n\n # region Reshape functions for computing products\n def _reshape_w_for_wgts_centric_partial_products(self, w: torch.Tensor, *, view_shape: torch.Size) -> torch.Tensor:\n return _reshape_w_for_wgts(w, view_shape)\n\n def _reshape_x_for_wgts_centric_partial_products(\n self, x: torch.Tensor, *, view_shape: torch.Size, fn: ReshapeFn\n ) -> torch.Tensor:\n return _reshape_x_for_wgts(fn(x), view_shape)\n\n def _reshape_w_for_ipts_centric_partial_products(self, w: torch.Tensor, *, view_shape: torch.Size) -> torch.Tensor:\n return _reshape_w_for_ipts(w, view_shape)\n\n def _reshape_x_for_ipts_centric_partial_products(\n self, x: torch.Tensor, *, view_shape: torch.Size, fn: ReshapeFn = None\n ) -> torch.Tensor:\n return _reshape_x_for_ipts(x, view_shape)\n\n def _reshape_w_for_full_products(self, w: torch.Tensor, *, view_shape: torch.Size = None) -> torch.Tensor:\n return w.view(w.shape[0], -1).T\n\n def _reshape_x_for_full_products(\n self, x: torch.Tensor, *, fn: ReshapeFn, view_shape: torch.Size = None\n ) -> torch.Tensor:\n return fn(x).view(x.shape[0], -1)\n\n # endregion\n\n @abstractmethod\n def _process_x_in_xw(self, x: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor: ...\n\n @abstractmethod\n def _process_w_in_xw(self, w: torch.Tensor) -> torch.Tensor: ...\n\n @abstractmethod\n def _process_y_in_yx(self, y: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor: ...\n\n @abstractmethod\n def _process_x_in_yx(self, x: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor: ...\n\n @abstractmethod\n def _process_xw_in_yx(self, w: torch.Tensor) -> torch.Tensor: ...\n\n @abstractmethod\n def _process_yw_in_yx(self, w: torch.Tensor) -> torch.Tensor: ...\n\n def _recover_mod(self) -> None:\n for p, w in self._state_dict:\n p.data = w\n self._state_dict.clear()\n for hook in self._hooks:\n hook.remove()\n self._hooks.clear()\n\n def _process_wgts_centric_mod(\n self, wgts: list[nn.Parameter], mods: list[nn.Module], update_state_dict: bool = True, **kwargs\n ) -> None:\n if self.needs_w_quant_for_wgts:\n for w in wgts:\n if update_state_dict:\n self._state_dict.append((w, w.data))\n w.data = self._process_w_in_xw(w.data)\n if self.needs_x_quant_for_wgts:\n self._hooks.append(self.x_quantizer.as_hook(func=self._process_x_in_xw, is_output=False).register(mods))\n\n def _process_ipts_centric_mod(\n self, wgts: list[nn.Parameter], mods: list[nn.Module], update_state_dict: bool = True, **kwargs\n ) -> None:\n if self.needs_w_quant_for_ipts:\n for w in wgts:\n if update_state_dict:\n self._state_dict.append((w, w.data))\n w.data = self._process_w_in_xw(w.data)\n if self.needs_x_quant_for_ipts:\n self._hooks.append(self.x_quantizer.as_hook(self._process_x_in_xw, is_output=False).register(mods))\n\n def _process_opts_centric_mod(\n self,\n x_wgts: list[nn.Parameter],\n y_wgts: list[nn.Parameter],\n x_mods: list[nn.Module],\n y_mods: list[nn.Module],\n update_state_dict: bool = True,\n **kwargs,\n ) -> None:\n if self.needs_w_quant_for_opts:\n for w in x_wgts:\n if update_state_dict:\n self._state_dict.append((w, w.data))\n w.data = self._process_xw_in_yx(w.detach().data)\n for w in y_wgts:\n if update_state_dict:\n self._state_dict.append((w, w.data))\n w.data = self._process_yw_in_yx(w.detach().data)\n if self.needs_x_quant_for_opts:\n self._hooks.append(self.x_quantizer.as_hook(self._process_x_in_yx, is_output=True).register(x_mods))\n if self.needs_y_quant_for_opts:\n self._hooks.append(self.y_quantizer.as_hook(self._process_y_in_yx, is_output=True).register(y_mods))\n\n def calibrate(\n self,\n x_wgts: list[nn.Parameter] | None = None,\n y_wgts: list[nn.Parameter] | None = None,\n x_acts: TensorsCache | None = None,\n y_acts: TensorsCache | None = None,\n x_mods: list[nn.Module] | None = None,\n y_mods: list[nn.Module] | None = None,\n eval_inputs: TensorsCache | None = None,\n eval_module: nn.Module | None = None,\n eval_kwargs: dict[str, tp.Any] | None = None,\n orig_x_wgts: list[tuple[nn.Parameter, torch.Tensor]] | None = None,\n orig_y_wgts: list[tuple[nn.Parameter, torch.Tensor]] | None = None,\n orig_x_acts: TensorsCache | None = None,\n orig_y_acts: TensorsCache | None = None,\n orig_eval_inputs: TensorsCache | None = None,\n **kwargs,\n ) -> _CANDIDATE:\n \"\"\"Calibrate the quantization parameters.\n\n Args:\n x_wgts (`list[nn.Parameter]` or `None`, *optional*, defaults to `None`):\n The weights in x-w computation, or weights that generates x for y-x computation.\n y_wgts (`list[nn.Parameter]` or `None`, *optional*, defaults to `None`):\n The weights that generates y for y-x computation.\n x_acts (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The x activations. It should be x for x-w or y-x computation.\n y_acts (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The y activations. It should be y for y-x computation.\n eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The inputs of evaluation module `eval_module`.\n eval_module (`nn.Module` or `None`, *optional*, defaults to `None`):\n The module used for evaluation.\n x_mods (`list[nn.Module]` or `None`, *optional*, defaults to `None`):\n The modules for x activation quantization.\n It should be the modules that take in x for x-w computation,\n or the modules that generates x for y-x computation.\n y_mods (`list[nn.Module]` or `None`, *optional*, defaults to `None`):\n The modules for y activation quantization.\n It should be the modules that generates y for y-x computation.\n orig_x_wgts (`list[tuple[nn.Parameter, torch.Tensor]]` or `None`, *optional*, defaults to `None`):\n The original weights for `x_mods`.\n orig_y_wgts (`list[tuple[nn.Parameter, torch.Tensor]]` or `None`, *optional*, defaults to `None`):\n The original weights for `y_mods`.\n orig_x_acts (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The original x activations `x_acts`.\n orig_y_acts (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The original y activations `y_acts`.\n orig_eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The original inputs of evaluation module `eval_inputs`.\n eval_kwargs (`dict[str, tp.Any]` or `None`, *optional*, defaults to `None`):\n The keyword arguments for evaluation module `eval_module`.\n\n Returns:\n `_CANDIDATE`:\n The best candidate.\n \"\"\"\n tools.logging.Formatter.indent_inc()\n if self.w_quantizer is not None and self.w_quantizer.is_enabled():\n self.logger.debug(f\"+ w: {self.w_quantizer.config.quant_dtype}\")\n else:\n self.logger.debug(\"+ w: None\")\n if self.x_quantizer is not None and self.x_quantizer.is_enabled():\n self.logger.debug(f\"+ x: {self.x_quantizer.config.quant_dtype}\")\n else:\n self.logger.debug(\"+ x: None\")\n if self.y_quantizer is not None and self.y_quantizer.is_enabled():\n self.logger.debug(f\"+ y: {self.y_quantizer.config.quant_dtype}\")\n else:\n self.logger.debug(\"+ y: None\")\n (\n x_wgts,\n y_wgts,\n x_acts,\n y_acts,\n eval_inputs,\n eval_module,\n x_mods,\n y_mods,\n orig_x_wgts,\n orig_y_wgts,\n orig_x_acts,\n orig_y_acts,\n orig_eval_inputs,\n ) = self._parse_args(\n x_wgts,\n y_wgts,\n x_acts,\n y_acts,\n eval_inputs,\n eval_module,\n x_mods,\n y_mods,\n orig_x_wgts,\n orig_y_wgts,\n orig_x_acts,\n orig_y_acts,\n orig_eval_inputs,\n )\n eval_kwargs = eval_kwargs or {}\n self.logger.debug(f\"+ finished parsing calibration arguments, ram usage: {psutil.virtual_memory().percent}\")\n self.reset(\n x_wgts=x_wgts,\n y_wgts=y_wgts,\n x_acts=x_acts,\n y_acts=y_acts,\n eval_inputs=eval_inputs,\n eval_module=eval_module,\n x_mods=x_mods,\n y_mods=y_mods,\n orig_x_wgts=orig_x_wgts,\n orig_y_wgts=orig_y_wgts,\n orig_x_acts=orig_x_acts,\n orig_y_acts=orig_y_acts,\n orig_eval_inputs=orig_eval_inputs,\n eval_kwargs=eval_kwargs,\n **kwargs,\n )\n self.logger.debug(f\"+ finished resetting calibrator, ram usage: {psutil.virtual_memory().percent}\")\n gc.collect()\n torch.cuda.empty_cache()\n if self.tensor_type == TensorType.Weights:\n result = self._calibrate_wgts(\n x_wgts, eval_inputs, eval_module, x_mods, orig_x_wgts, orig_eval_inputs, eval_kwargs, **kwargs\n )\n elif self.tensor_type == TensorType.Inputs:\n result = self._calibrate_ipts(\n x_wgts, eval_inputs, eval_module, x_mods, orig_x_wgts, orig_eval_inputs, eval_kwargs, **kwargs\n )\n else:\n result = self._calibrate_opts(\n x_wgts,\n y_wgts,\n eval_inputs,\n eval_module,\n x_mods,\n y_mods,\n orig_x_wgts,\n orig_y_wgts,\n orig_eval_inputs,\n eval_kwargs,\n **kwargs,\n )\n tools.logging.Formatter.indent_dec()\n return result\n\n def _calibrate_wgts( # noqa: C901\n self,\n wgts: list[torch.Tensor | nn.Parameter],\n ipts: TensorsCache | None,\n eval_module: nn.Module | None,\n mods: list[nn.Module] | None,\n orig_wgts: list[tuple[nn.Parameter, torch.Tensor]] | None,\n orig_ipts: TensorsCache | None,\n eval_kwargs: dict[str, tp.Any],\n **kwargs,\n ) -> tp.Any:\n # region Step 1: Calculate the baseline\n if self.objective == SearchBasedCalibObjective.TensorError:\n if orig_wgts is None:\n orig_wgts = [(None, w.detach().data) for w in wgts]\n assert all(w.shape[1:] == wgts[0].shape[1:] for w in wgts)\n assert all(w.shape[1:] == wgts[0].shape[1:] for _, w in orig_wgts)\n orig_opts = None\n w_view_shapes = [infer_view_shape(w.shape, self.w_quantizer.config.largest_group_shape) for w in wgts]\n elif self.objective == SearchBasedCalibObjective.ProductsError:\n if orig_wgts is None:\n orig_wgts = [(None, w.detach().data) for w in wgts]\n assert len(orig_wgts) == len(wgts)\n assert all(w.shape[1:] == wgts[0].shape[1:] for w in wgts)\n assert all(w.shape[1:] == wgts[0].shape[1:] for _, w in orig_wgts)\n w_view_shapes = [infer_view_shape(w.shape, self.w_quantizer.config.largest_group_shape) for w in wgts]\n if self.granularity != SearchBasedCalibGranularity.Layer:\n _reshape_x = self._reshape_x_for_wgts_centric_partial_products\n _reshape_w = self._reshape_w_for_wgts_centric_partial_products\n else:\n _reshape_x = self._reshape_x_for_full_products\n _reshape_w = self._reshape_w_for_full_products\n assert isinstance(ipts, TensorsCache), \"ipts should not be None for ProductsError\"\n if orig_ipts is None:\n orig_ipts = ipts\n same_ipts = orig_ipts is ipts\n orig_ipts = TensorsCache(\n {\n key: TensorCache(\n [_reshape_x(x, view_shape=w_view_shapes[0], fn=ipt.reshape) for x in ipt.data],\n **ipt.get_factory_kwargs(channels_dim=1, reshape=ReshapeFn()),\n )\n for key, ipt in orig_ipts.items()\n },\n )\n orig_opts: dict[tuple[int, ...], torch.Tensor] = {}\n for j, (_, w) in enumerate(orig_wgts):\n w = _reshape_w(w, view_shape=w_view_shapes[j])\n for s, ipt in enumerate(orig_ipts):\n for i, x in enumerate(ipt.data):\n x = x.to(device=w.device, non_blocking=True)\n y = torch.matmul(x, w)\n y = y.view(*y.shape[:-2], y.shape[-2] * y.shape[-1])\n orig_opts[(i, s, j)] = y.to(device=self.opts_device or y.device, non_blocking=True)\n if self.needs_to_pre_reshape_x_for_wgts:\n if same_ipts:\n ipts = orig_ipts\n else:\n ipts = TensorsCache(\n {\n key: TensorCache(\n [_reshape_x(x, view_shape=w_view_shapes[0], fn=ipt.reshape) for x in ipt.data],\n **ipt.get_factory_kwargs(channels_dim=1, reshape=ReshapeFn()),\n )\n for key, ipt in ipts.items()\n }\n )\n del orig_wgts, orig_ipts, same_ipts\n elif self.objective == SearchBasedCalibObjective.OutputsError:\n w_view_shapes, _state_dict = [], []\n if orig_wgts is not None:\n _state_dict = [(p, p.data) for p, _ in orig_wgts]\n for p, w in orig_wgts:\n p.data = w.to(device=p.data.device)\n if orig_ipts is None:\n orig_ipts = ipts\n assert isinstance(orig_ipts, TensorsCache), \"orig_ipts should not be None for OutputsError\"\n orig_opts: dict[tuple[int, ...], torch.Tensor] = {}\n for i in range(len(orig_ipts.front().data)):\n ipt = orig_ipts.extract(i, eval_kwargs)\n y = eval_module(*ipt.args, **ipt.kwargs)\n y = y[0] if not isinstance(y, torch.Tensor) else y\n assert isinstance(y, torch.Tensor), \"eval_mod should return a tensor\"\n orig_opts[(i,)] = y.to(device=self.opts_device or y.device, non_blocking=True)\n del ipt, y\n for p, s in _state_dict:\n p.data = s\n del orig_wgts, orig_ipts, _state_dict\n else:\n raise ValueError(f\"Unknown objective {self.objective}\")\n gc.collect()\n torch.cuda.empty_cache()\n self.logger.debug(f\"+ finished calculating the original outputs, ram usage: {psutil.virtual_memory().percent}\")\n # endregion\n while not self.is_done():\n self.ask()\n e: list[torch.Tensor] = []\n # region Step 2: Calculate the errors\n if self.objective == SearchBasedCalibObjective.TensorError:\n assert isinstance(orig_wgts, (tuple, list))\n for w, (_, orig_w), w_view_shape in zip(wgts, orig_wgts, w_view_shapes, strict=True):\n e_w = self._process_w_in_xw(w).sub_(orig_w)\n if self.granularity == SearchBasedCalibGranularity.Group:\n e_w = e_w.view(w_view_shape).abs_().pow_(self.config.degree)\n e_w = e_w.sum(dim=tuple(range(1, len(w_view_shape), 2))).view(w_view_shape[::2])\n elif self.granularity == SearchBasedCalibGranularity.ChannelGroup:\n e_w = e_w.view(*w_view_shape[:4], -1).abs_().pow_(self.config.degree)\n e_w = e_w.sum(dim=(0, 1, 3, 4)).view(w_view_shape[2])\n elif self.granularity == SearchBasedCalibGranularity.Layer:\n e_w = e_w.abs_().pow_(self.config.degree).sum().view(-1)\n else:\n raise ValueError(f\"Unknown granularity {self.granularity}\")\n e.append(e_w)\n elif self.objective == SearchBasedCalibObjective.ProductsError:\n e = [None] * len(wgts)\n for j, w in enumerate(wgts):\n w = _reshape_w(self._process_w_in_xw(w), view_shape=w_view_shapes[j])\n for s, ipt in enumerate(ipts):\n for i, x in enumerate(ipt.data):\n x = x.to(device=w.device, non_blocking=True)\n if not self.needs_to_pre_reshape_x_for_wgts:\n x = self._process_x_in_xw(x, channels_dim=ipt.channels_dim)\n x = _reshape_x(x, view_shape=w_view_shapes[j], fn=ipt.reshape)\n y = torch.matmul(x, w)\n y = y.view(*y.shape[:-2], y.shape[-2] * y.shape[-1])\n y = y.sub_(orig_opts[(i, s, j)].to(device=w.device, non_blocking=True))\n if self.granularity == SearchBasedCalibGranularity.Group:\n y = y.to(self.develop_dtype).pow_(self.config.degree).sum(dim=-1)\n elif self.granularity == SearchBasedCalibGranularity.ChannelGroup:\n y = y.view(y.shape[0], y.shape[1], -1)\n y = y.to(self.develop_dtype).pow_(self.config.degree).sum(dim=(0, 2))\n elif self.granularity == SearchBasedCalibGranularity.Layer:\n y = y.to(self.develop_dtype).pow_(self.config.degree).sum().view(-1)\n else:\n raise ValueError(f\"Unknown granularity {self.granularity}\")\n if e[j] is None:\n e[j] = y\n else:\n e[j].add_(y)\n elif self.objective == SearchBasedCalibObjective.OutputsError:\n self._process_wgts_centric_mod(wgts=wgts, mods=mods, **kwargs)\n e = [None]\n for i in range(len(ipts.front().data)):\n ipt = ipts.extract(i, eval_kwargs)\n y = eval_module(*ipt.args, **ipt.kwargs)\n y = y[0] if not isinstance(y, torch.Tensor) else y\n assert isinstance(y, torch.Tensor), \"eval_mod should return a tensor\"\n y = (y - orig_opts[(i,)].to(device=y.device, non_blocking=True)).to(self.develop_dtype)\n y = y.pow_(self.config.degree).sum().view(-1)\n if e[0] is None:\n e[0] = y\n else:\n e[0].add_(y)\n del ipt, y\n self._recover_mod()\n else:\n raise ValueError(f\"Unknown objective {self.objective}\")\n # endregion\n self.tell(e)\n return self.get_best()\n\n def _calibrate_ipts( # noqa: C901\n self,\n wgts: list[torch.Tensor | nn.Parameter],\n ipts: TensorsCache,\n eval_module: nn.Module | None,\n mods: list[nn.Module] | None,\n orig_wgts: list[tuple[nn.Parameter, torch.Tensor]] | None,\n orig_ipts: TensorsCache | None,\n eval_kwargs: dict[str, tp.Any],\n **kwargs,\n ) -> tp.Any:\n if orig_ipts is None:\n orig_ipts = ipts\n assert ipts.num_tensors == orig_ipts.num_tensors\n assert all(\n x.shape == orig_x.shape\n for ipt, orig_ipt in zip(ipts, orig_ipts, strict=True)\n for x, orig_x in zip(ipt.data, orig_ipt.data, strict=True)\n )\n # region Step 1: Calculate the outputs\n if self.objective == SearchBasedCalibObjective.TensorError:\n assert all(x.shape == ipt.data[0].shape for ipt in ipts for x in ipt.data)\n orig_opts = None\n x_view_shapes = [\n infer_view_shape(\n ipt.data[0].view(-1, *ipt.data[0].shape[ipt.channels_dim :]).shape,\n self.x_quantizer.config.largest_group_shape,\n skip_first_dim=True,\n )\n for ipt in ipts\n ]\n del orig_wgts\n elif self.objective == SearchBasedCalibObjective.ProductsError:\n assert all(ipt.channels_dim == 1 for ipt in ipts)\n assert all(ipt.channels_dim == 1 for ipt in orig_ipts)\n assert all(x.shape[1:] == ipts.front().data[0].shape[1:] for ipt in ipts for x in ipt.data)\n if orig_wgts is None:\n orig_wgts = [(None, w.detach().data) for w in wgts]\n assert len(orig_wgts) == len(wgts)\n if self.granularity != SearchBasedCalibGranularity.Layer:\n _reshape_x = self._reshape_x_for_ipts_centric_partial_products\n _reshape_w = self._reshape_w_for_ipts_centric_partial_products\n else:\n _reshape_x = self._reshape_x_for_full_products\n _reshape_w = self._reshape_w_for_full_products\n x_view_shapes = [\n infer_view_shape(ipt.data[0].shape, self.x_quantizer.config.largest_group_shape, skip_first_dim=True)\n for ipt in ipts\n ]\n orig_opts: dict[tuple[int, ...], torch.Tensor] = {}\n for j, (_, w) in enumerate(orig_wgts):\n w = _reshape_w(w, view_shape=x_view_shapes[0])\n for s, ipt in enumerate(orig_ipts):\n for i, x in enumerate(ipt.data):\n x = x.to(device=w.device, non_blocking=True)\n x = _reshape_x(x, view_shape=x_view_shapes[s], fn=ipt.reshape)\n y = torch.matmul(x, w)\n y = y.view(*y.shape[:-2], y.shape[-2] * y.shape[-1])\n orig_opts[(i, s, j)] = y.to(device=self.opts_device or y.device, non_blocking=True)\n if self.needs_to_pre_reshape_w_for_ipts:\n for j, w in enumerate(wgts):\n wgts[j] = _reshape_w(w, view_shape=x_view_shapes[0])\n del orig_wgts, orig_ipts\n elif self.objective == SearchBasedCalibObjective.OutputsError:\n x_view_shapes, _state_dict = [], []\n if orig_wgts is not None:\n _state_dict = [(p, p.data) for p, _ in orig_wgts]\n for p, w in orig_wgts:\n p.data = w.to(device=p.data.device)\n orig_opts: dict[tuple[int, ...], torch.Tensor] = {}\n for i in range(len(orig_ipts.front().data)):\n ipt = orig_ipts.extract(i, eval_kwargs)\n y = eval_module(*ipt.args, **ipt.kwargs)\n y = y[0] if not isinstance(y, torch.Tensor) else y\n assert isinstance(y, torch.Tensor), \"eval_mod should return a tensor\"\n orig_opts[(i,)] = y.to(device=self.opts_device or y.device, non_blocking=True)\n del ipt, y\n for p, s in _state_dict:\n p.data = s\n del orig_wgts, orig_ipts, _state_dict\n else:\n raise ValueError(f\"Unknown objective {self.objective}\")\n gc.collect()\n torch.cuda.empty_cache()\n # endregion\n while not self.is_done():\n self.ask()\n e: list[torch.Tensor] = []\n # region Step 2: Calculate the outputs errors\n if self.objective == SearchBasedCalibObjective.TensorError:\n e = [None] * len(ipts)\n for s, (ipt, x_view_shape) in enumerate(zip(ipts, x_view_shapes, strict=True)):\n for x in ipt.data:\n e_x = self._process_x_in_xw(x, channels_dim=ipt.channels_dim).sub_(x)\n if self.granularity == SearchBasedCalibGranularity.Group:\n e_x = e_x.view(x_view_shape).abs_().pow_(self.config.degree)\n e_x = e_x.sum(dim=tuple(range(1, len(x_view_shape), 2)))\n if self.granularity == SearchBasedCalibGranularity.ChannelGroup:\n e_x = e_x.view(*x_view_shape[:4], -1).abs_().pow_(self.config.degree)\n e_x = e_x.sum(dim=(0, 1, 3, 4)).view(x_view_shape[2])\n elif self.granularity == SearchBasedCalibGranularity.Layer:\n e_x = e_x.abs_().pow_(self.config.degree).sum().view(-1)\n else:\n raise ValueError(f\"Unknown granularity {self.granularity}\")\n if e[s] is None:\n e[s] = e_x\n else:\n e[s].add_(e_x)\n elif self.objective == SearchBasedCalibObjective.ProductsError:\n e = [None] * len(ipts)\n for j, w in enumerate(wgts):\n if not self.needs_to_pre_reshape_w_for_ipts:\n w = self._process_w_in_xw(w)\n w = _reshape_w(w, view_shape=x_view_shapes[0])\n for s, ipt in enumerate(ipts):\n for i, x in enumerate(ipt.data):\n x = x.to(device=w.device, non_blocking=True)\n x = self._process_x_in_xw(x, channels_dim=ipt.channels_dim)\n x = _reshape_x(x, view_shape=x_view_shapes[s], fn=ipt.reshape)\n y = torch.matmul(x, w)\n y = y.view(*y.shape[:-2], y.shape[-2] * y.shape[-1])\n y = y.sub_(orig_opts[(i, s, j)].to(device=w.device, non_blocking=True))\n if self.granularity == SearchBasedCalibGranularity.Group:\n y = y.to(self.develop_dtype).pow_(self.config.degree).sum(dim=-1)\n elif self.granularity == SearchBasedCalibGranularity.ChannelGroup:\n y = y.view(y.shape[0], y.shape[1], -1)\n y = y.to(self.develop_dtype).pow_(self.config.degree).sum(dim=(0, 2))\n elif self.granularity == SearchBasedCalibGranularity.Layer:\n y = y.to(self.develop_dtype).pow_(self.config.degree).sum().view(-1)\n else:\n raise ValueError(f\"Unknown granularity {self.granularity}\")\n if e[s] is None:\n e[s] = y\n else:\n e[s].add_(y)\n elif self.objective == SearchBasedCalibObjective.OutputsError:\n self._process_ipts_centric_mod(wgts=wgts, mods=mods, **kwargs)\n e = [None]\n for i in range(len(ipts.front().data)):\n ipt = ipts.extract(i, eval_kwargs)\n y = eval_module(*ipt.args, **ipt.kwargs)\n y = y[0] if not isinstance(y, torch.Tensor) else y\n assert isinstance(y, torch.Tensor), \"eval_mod should return a tensor\"\n y = (y - orig_opts[(i,)].to(device=y.device, non_blocking=True)).to(self.develop_dtype)\n y = y.pow_(self.config.degree).sum().view(-1)\n if e[0] is None:\n e[0] = y\n else:\n e[0].add_(y)\n del ipt, y\n self._recover_mod()\n else:\n raise ValueError(f\"Unknown objective {self.objective}\")\n # endregion\n self.tell(e)\n return self.get_best()\n\n def _calibrate_opts( # noqa: C901\n self,\n x_wgts: list[torch.Tensor | nn.Parameter],\n y_wgts: list[torch.Tensor | nn.Parameter],\n eval_inputs: TensorsCache | None,\n eval_module: nn.Module | None,\n x_mods: list[nn.Module] | None,\n y_mods: list[nn.Module] | None,\n orig_x_wgts: list[tuple[nn.Parameter, torch.Tensor]] | None,\n orig_y_wgts: list[tuple[nn.Parameter, torch.Tensor]] | None,\n orig_eval_inputs: TensorsCache | None,\n eval_kwargs: dict[str, tp.Any],\n **kwargs,\n ) -> tp.Any:\n # region Step 1: Calculate the outputs\n if self.objective == SearchBasedCalibObjective.OutputsError:\n assert eval_inputs is not None, \"eval_inputs should not be None when objective is OutputsError\"\n if orig_eval_inputs is None:\n orig_eval_inputs = eval_inputs\n assert eval_inputs.num_tensors == orig_eval_inputs.num_tensors\n assert all(\n x.shape == orig_x.shape\n for key, ipt in eval_inputs.items()\n for x, orig_x in zip(ipt.data, orig_eval_inputs[key].data, strict=True)\n )\n _x_state_dict, _y_state_dict = [], []\n if orig_x_wgts is not None:\n _x_state_dict = [(p, p.data) for p, _ in orig_x_wgts]\n for p, w in orig_x_wgts:\n p.data = w.to(device=p.data.device)\n if orig_y_wgts is not None:\n _y_state_dict = [(p, p.data) for p, _ in orig_y_wgts]\n for p, w in orig_y_wgts:\n p.data = w.to(device=p.data.device)\n orig_opts: dict[tuple[int, ...], torch.Tensor] = {}\n for i in range(len(orig_eval_inputs.front().data)):\n ipt = orig_eval_inputs.extract(i, eval_kwargs)\n y = eval_module(*ipt.args, **ipt.kwargs)\n y = y[0] if not isinstance(y, torch.Tensor) else y\n assert isinstance(y, torch.Tensor), \"eval_mod should return a tensor\"\n orig_opts[(i,)] = y.to(device=self.opts_device or y.device, non_blocking=True)\n del ipt, y\n for p, s in _x_state_dict:\n p.data = s\n for p, s in _y_state_dict:\n p.data = s\n del orig_x_wgts, orig_y_wgts, orig_eval_inputs, _x_state_dict, _y_state_dict\n else:\n raise ValueError(f\"Unknown objective {self.objective}\")\n gc.collect()\n torch.cuda.empty_cache()\n # endregion\n while not self.is_done():\n self.ask()\n e: list[torch.Tensor] = []\n # region Step 2: Calculate the outputs errors\n if self.objective == SearchBasedCalibObjective.OutputsError:\n self._process_opts_centric_mod(\n x_wgts=x_wgts,\n y_wgts=y_wgts,\n x_mods=x_mods,\n y_mods=y_mods,\n **kwargs,\n )\n e = [None]\n for i in range(len(eval_inputs.front().data)):\n ipt = eval_inputs.extract(i, eval_kwargs)\n y = eval_module(*ipt.args, **ipt.kwargs)\n y = y[0] if not isinstance(y, torch.Tensor) else y\n assert isinstance(y, torch.Tensor), \"eval_mod should return a tensor\"\n y = (y - orig_opts[(i,)].to(device=y.device, non_blocking=True)).to(self.develop_dtype)\n y = y.pow_(self.config.degree).sum().view(-1)\n if e[0] is None:\n e[0] = y\n else:\n e[0].add_(y)\n del ipt, y\n self._recover_mod()\n else:\n raise ValueError(f\"Unknown objective {self.objective}\")\n # endregion\n self.tell(e)\n return self.get_best()\n"
},
{
"path": "deepcompressor/calib/smooth.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Smooth quantization module.\"\"\"\n\nimport gc\nimport typing as tp\nfrom dataclasses import _MISSING_TYPE, MISSING, dataclass\n\nimport torch\nimport torch.nn as nn\n\nfrom ..data.cache import TensorsCache\nfrom ..data.common import TensorType\nfrom ..quantizer.processor import Quantizer\nfrom ..utils import math, tools\nfrom ..utils.common import split_sequence\nfrom ..utils.hooks import BaseInputPackager, BaseOutputPackager, BaseTensorProcessor\nfrom .config import SearchBasedCalibObjective, SmoothCalibConfig, SmoothSpanMode\nfrom .metric import ChannelMetric\nfrom .search import SearchBasedCalibrator\n\n__all__ = [\n \"smooth_linear_modules\",\n \"smooth_attention\",\n \"convert_smooth_upscale_to_downscale\",\n \"ActivationSmoother\",\n \"get_smooth_scale\",\n \"get_smooth_span\",\n \"SmoothCalibrator\",\n \"SmoothLinearCalibrator\",\n \"SmoothAttentionCalibrator\",\n]\n\n\n@dataclass\nclass ActivationSmoother(BaseTensorProcessor):\n \"\"\"The quantization smoothing processor.\"\"\"\n\n smooth_scale: torch.Tensor\n channels_dim: int\n upscale: bool = False\n develop_dtype: torch.dtype | None = None\n # region hook-related attributes\n input_packager: BaseInputPackager | None = None\n output_packager: BaseOutputPackager | None = None\n # endregion\n\n def is_enabled(self) -> bool:\n return self.smooth_scale is not None\n\n def get_input_packager(self) -> BaseInputPackager | None:\n return self.input_packager\n\n def get_output_packager(self) -> BaseOutputPackager | None:\n return self.output_packager\n\n def process(self, tensor: torch.Tensor) -> torch.Tensor:\n \"\"\"Process the tensor.\n\n Args:\n tensor (`torch.Tensor`):\n The tensor to smooth.\n\n Returns:\n `torch.Tensor`:\n The smoothed tensor.\n \"\"\"\n device, dtype = tensor.device, tensor.dtype\n if self.develop_dtype is None:\n self.develop_dtype = dtype\n self.smooth_scale = self.smooth_scale.to(device=device, dtype=self.develop_dtype)\n tensor = tensor.to(dtype=self.develop_dtype)\n smooth_scale_view_shape = [1] * tensor.ndim\n smooth_scale_view_shape[self.channels_dim] = -1\n smooth_scale = self.smooth_scale.view(smooth_scale_view_shape)\n if self.upscale:\n return tensor.mul(smooth_scale).to(dtype=dtype)\n else:\n return tensor.div(smooth_scale).to(dtype=dtype)\n\n\n@torch.inference_mode()\ndef get_smooth_span(\n tensors: tp.Sequence[torch.Tensor],\n /,\n *,\n group_shape: tp.Sequence[int],\n span_mode: SmoothSpanMode,\n device: torch.device | str | None = None,\n dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Calculate the value span of tensors for calculating smoothing scale.\n\n Args:\n tensors (`Sequence[torch.Tensor]`):\n Tensors to calculate the span.\n group_shape (`Sequence[int]`):\n Quantization group shape.\n span_mode (`SmoothSpanMode`):\n The quantization smoothing span mode.\n device (`torch.device` or `str` or `None`, *optional*, defaults to `None`):\n Device to store the span.\n dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n Data type of the span.\n\n Returns:\n `torch.Tensor`:\n The span of the tensors for calculating smoothing scale.\n \"\"\"\n # convert span mode name from camel case to snake case\n range_name = \"\".join([\"_\" + c.lower() if c.isupper() else c for c in span_mode.name]).lstrip(\"_\")\n range_fn = getattr(ChannelMetric, range_name)\n r: torch.Tensor = range_fn(tensors, tensors[0].shape[1], group_shape, device=device, dtype=dtype)\n return r\n\n\n@torch.inference_mode()\ndef get_smooth_scale(*, alpha_base: torch.Tensor, beta_base: torch.Tensor, alpha: float, beta: float) -> torch.Tensor:\n \"\"\"Calculate the smoothing scale for quantization. Scale = alpha_base^alpha / beta_base^beta.\n\n Args:\n alpha_base (`torch.Tensor`):\n Base span for alpha.\n beta_base (`torch.Tensor`):\n Base span for beta.\n alpha (`float`):\n Alpha.\n beta (`float`):\n Beta.\n\n Returns:\n `torch.Tensor`:\n Smoothing scale.\n \"\"\"\n assert 0 <= alpha <= 1 and 0 <= beta <= 1, \"The smooth factors should be in [0, 1].\"\n if alpha > 0:\n smooth_scale = alpha_base.pow(alpha)\n if beta > 0:\n smooth_scale = smooth_scale.div_(beta_base.pow(beta))\n else:\n smooth_scale = beta_base.pow(-beta)\n smooth_scale[smooth_scale == 0] = 1\n if smooth_scale.isnan().any() or smooth_scale.isinf().any():\n smooth_scale = smooth_scale.fill_(1)\n assert not smooth_scale.isnan().any(), \"The smooth scale contains NaN.\"\n assert not smooth_scale.isinf().any(), \"The smooth scale contains Inf.\"\n return smooth_scale\n\n\nclass SmoothCalibrator(SearchBasedCalibrator[SmoothCalibConfig, torch.Tensor]):\n \"\"\"The quantization smoothing calibrator.\"\"\"\n\n def __init__(\n self,\n tensor_type: TensorType,\n config: SmoothCalibConfig,\n w_quantizer: Quantizer | None,\n x_quantizer: Quantizer | None,\n y_quantizer: Quantizer | None,\n num_heads: int = 1,\n num_head_repeats: int = 1,\n with_rope: bool = False,\n develop_dtype: torch.dtype = torch.float32,\n ) -> None:\n \"\"\"Initialize the calibrator.\n\n Args:\n tensor_type (`TensorType`):\n The type of tensor to quantize. Choices are ``Weights`` and ``Outputs``.\n config (`SmoothCalibConfig`):\n The quantization smoothing calibration configuration.\n w_quantizer (`Quantizer` or `None`):\n The w quantizer for x-w computation.\n x_quantizer (`Quantizer` or `None`):\n The x quantizer for x-w or y-x computation.\n y_quantizer (`Quantizer` or `None`):\n The y quantizer for y-x computation.\n num_heads (`int`, *optional*, defaults to ``1``):\n The number of heads.\n num_head_repeats (`int`, *optional*, defaults to ``1``):\n The number of head repeats.\n with_rope (`bool`, *optional*, defaults to ``False``):\n Whether rotary position embedding is used for y-x computation.\n develop_dtype (torch.dtype, *optional*, defaults to ``torch.float32``):\n The development data type.\n \"\"\"\n assert tensor_type in (TensorType.Weights, TensorType.Outputs)\n super().__init__(\n tensor_type=tensor_type,\n config=config,\n w_quantizer=w_quantizer,\n x_quantizer=x_quantizer,\n y_quantizer=y_quantizer,\n develop_dtype=develop_dtype,\n )\n self.num_heads = num_heads\n self.num_head_repeats = num_head_repeats\n self.with_rope = self.tensor_type != TensorType.Weights and with_rope\n # region set group shapes of weights, inputs and outputs\n if self.needs_w_quant:\n w_group_shape = list(self.w_quantizer.config.largest_group_shape)\n else:\n w_group_shape = [1, None, -1]\n if self.needs_x_quant:\n x_group_shape = list(self.x_quantizer.config.largest_group_shape)\n else:\n x_group_shape = [1, None, -1]\n if self.needs_y_quant:\n y_group_shape = list(self.y_quantizer.config.largest_group_shape)\n else:\n y_group_shape = [1, None, -1]\n w_group_shape[1] = x_group_shape[1] if w_group_shape[1] is None else w_group_shape[1]\n if self.tensor_type == TensorType.Weights:\n x_group_shape[1] = w_group_shape[1] if x_group_shape[1] is None else x_group_shape[1]\n else:\n x_group_shape[1] = y_group_shape[1] if x_group_shape[1] is None else x_group_shape[1]\n y_group_shape[1] = x_group_shape[1] if y_group_shape[1] is None else y_group_shape[1]\n self.w_group_shape, self.x_group_shape, self.y_group_shape = w_group_shape, x_group_shape, y_group_shape\n # endregion\n self.alpha_beta_pairs = self.config.get_alpha_beta_pairs()\n self.num_iters = 1\n\n @property\n def population_size(self) -> int:\n \"\"\"Get the population size.\"\"\"\n return len(self.alpha_beta_pairs) * len(self.span_mode_pairs)\n\n @property\n def allows_x_quant_for_wgts(self) -> bool:\n \"\"\"Whether the calibrator allows input quantization when tensor_type is Weights.\"\"\"\n return self.config.allow_a_quant\n\n @property\n def allows_w_quant_for_wgts(self) -> bool:\n \"\"\"Whether the calibrator needs weight quantization when tensor_type is Weights.\"\"\"\n return self.config.allow_b_quant\n\n @property\n def allows_w_quant_for_ipts(self) -> bool:\n \"\"\"Whether the calibrator allows weight quantization when tensor_type is Inputs.\"\"\"\n return self.config.allow_b_quant\n\n @property\n def allows_x_quant_for_opts(self) -> bool:\n \"\"\"Whether the calibrator allows x quantization when tensor_type is Outputs.\"\"\"\n return self.config.allow_b_quant\n\n @property\n def allows_y_quant_for_opts(self) -> bool:\n \"\"\"Whether the calibrator allows y quantization when tensor_type is Outputs.\"\"\"\n return self.config.allow_a_quant\n\n @property\n def allows_w_quant_for_opts(self) -> bool:\n \"\"\"Whether the calibrator allows weight quantization when tensor_type is Outputs.\"\"\"\n return False\n\n @property\n def span_mode_pairs(self) -> list[tuple[SmoothSpanMode, SmoothSpanMode]]:\n \"\"\"Get the span modes.\"\"\"\n return self.config.spans\n\n @property\n def alpha_span_modes(self) -> list[SmoothSpanMode]:\n \"\"\"Get the span modes for alpha.\"\"\"\n return self.config.a_spans\n\n @property\n def beta_span_modes(self) -> list[SmoothSpanMode]:\n \"\"\"Get the span modes for beta.\"\"\"\n return self.config.b_spans\n\n def _reset( # noqa: C901\n self,\n *,\n x_wgts: list[torch.Tensor | nn.Parameter],\n x_acts: TensorsCache,\n y_wgts: list[torch.Tensor | nn.Parameter] = None,\n y_acts: TensorsCache | None = None,\n **kwargs,\n ) -> None:\n \"\"\"Reset the calibrator.\n\n Args:\n x_wgts (`list[torch.Tensor | nn.Parameter]`):\n The weights in x-w computation, or weights that generates x for y-x computation.\n x_acts (`TensorsCache`):\n The x activations. It should be x for x-w or y-x computation.\n y_wgts (`list[torch.Tensor | nn.Parameter]` or `None`, *optional*, defaults to `None`):\n The weights that generates y for y-x computation.\n y_acts (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The y activations. It should be y for y-x computation.\n \"\"\"\n wgts_centric = self.tensor_type == TensorType.Weights\n self.num_in_channels = x_wgts[0].shape[1] if wgts_centric else x_wgts[0].shape[0]\n device = x_wgts[0].device\n if self.num_heads > 1 and self.num_head_repeats > 1:\n self.num_unique_heads = self.num_heads // self.num_head_repeats\n else:\n self.num_unique_heads = 0\n # region get x spans\n assert (\n x_acts.num_tensors == 1\n ), f\"Only one input is allowed, got {x_acts.num_tensors}=len({list(x_acts.keys())})\"\n x_tensors = x_acts.front().get_standardized_data(reshape=False)\n assert all(x.shape[1] == self.num_in_channels for x in x_tensors)\n x_spans = {}\n for span_mode in self.alpha_span_modes if wgts_centric else self.beta_span_modes:\n x_span = get_smooth_span(\n x_tensors,\n group_shape=self.x_group_shape,\n span_mode=span_mode,\n device=device,\n dtype=self.develop_dtype,\n )\n if self.num_unique_heads > 0:\n x_span = x_span.view(self.num_unique_heads, self.num_head_repeats, -1)\n x_span = (x_span.amax if \"Max\" in span_mode.name else x_span.mean)(dim=1, keepdim=True)\n x_span = x_span.expand(self.num_unique_heads, self.num_head_repeats, -1).reshape(-1)\n if self.tensor_type == TensorType.Outputs and self.with_rope:\n x_span = x_span.view(self.num_heads, 2, -1)\n x_span = (x_span.amax if \"Max\" in span_mode.name else x_span.mean)(dim=1, keepdim=True)\n x_span = x_span.expand(self.num_heads, 2, -1).reshape(-1)\n x_spans[span_mode] = x_span\n if self.logger.level <= tools.logging.DEBUG:\n self.logger.debug(\"+ x - %s\", span_mode.name)\n self.logger.debug(\"+ x = [min=%.4f, max=%.4f]\", x_span.min().item(), x_span.max().item())\n del x_tensors\n # endregion\n if wgts_centric:\n assert all(w.shape[1] == self.num_in_channels for w in x_wgts)\n w_tensors = [w.data for w in x_wgts]\n w_spans = {}\n for span_mode in self.beta_span_modes:\n w_span = get_smooth_span(\n w_tensors,\n group_shape=self.w_group_shape,\n span_mode=span_mode,\n dtype=self.develop_dtype,\n )\n if self.num_unique_heads > 0:\n w_span = w_span.view(self.num_unique_heads, self.num_head_repeats, -1)\n w_span = (w_span.amax if \"Max\" in span_mode.name else w_span.mean)(dim=1, keepdim=True)\n w_span = w_span.expand(self.num_unique_heads, self.num_head_repeats, -1).reshape(-1)\n w_spans[span_mode] = w_span\n if self.logger.level <= tools.logging.DEBUG:\n self.logger.debug(\"+ w - %s\", span_mode.name)\n self.logger.debug(\"+ w = [min=%.4f, max=%.4f]\", w_span.min().item(), w_span.max().item())\n self.span_pairs: list[tuple[torch.Tensor, torch.Tensor]] = [\n (x_spans[x_span_mode], w_spans[w_span_mode]) for x_span_mode, w_span_mode in self.span_mode_pairs\n ]\n else:\n assert y_acts.num_tensors == 1, f\"Only one output source is allowed, got {y_acts.num_tensors}\"\n if self.num_unique_heads > 0:\n num_out_channels = self.num_in_channels // self.num_head_repeats\n else:\n num_out_channels = self.num_in_channels\n assert all(w.shape[0] == self.num_in_channels for w in x_wgts)\n assert all(w.shape[0] == num_out_channels for w in y_wgts)\n y_tensors = y_acts.front().get_standardized_data(reshape=False)\n assert all(y.shape[1] == num_out_channels for y in y_tensors)\n y_spans = {}\n for span_mode in self.alpha_span_modes:\n y_span = get_smooth_span(\n y_tensors,\n group_shape=self.x_group_shape,\n span_mode=span_mode,\n device=device,\n dtype=self.develop_dtype,\n )\n if self.num_unique_heads > 0:\n y_span = y_span.view(self.num_unique_heads, 1, -1)\n y_span = y_span.expand(self.num_unique_heads, self.num_head_repeats, -1).reshape(-1)\n if self.tensor_type == TensorType.Outputs and self.with_rope:\n y_span = y_span.view(self.num_heads, 2, -1)\n y_span = (y_span.amax if \"Max\" in span_mode.name else y_span.mean)(dim=1, keepdim=True)\n y_span = y_span.expand(self.num_heads, 2, -1).reshape(-1)\n y_spans[span_mode] = y_span\n if self.logger.level <= tools.logging.DEBUG:\n self.logger.debug(\"+ y - %s\", span_mode.name)\n self.logger.debug(\"+ y = [min=%.4f, max=%.4f]\", y_span.min().item(), y_span.max().item())\n self.span_pairs: list[tuple[torch.Tensor, torch.Tensor]] = [\n (y_spans[y_span_mode], x_spans[x_span_mode]) for y_span_mode, x_span_mode in self.span_mode_pairs\n ]\n self.best_error: list[torch.Tensor] = None\n self.best_scale: torch.Tensor = None\n self.error_history: list[tuple[float, float]] = []\n\n def _split_candidate_id(self, candidate_id: int) -> tuple[int, int]:\n \"\"\"Split the candidate id into alpha_beta id and span_pair id.\n\n Args:\n candidate_id (`int`):\n The candidate id.\n\n Returns:\n `tuple[int, int]`:\n The alpha_beta id and span_mode id.\n \"\"\"\n alpha_beta_id = candidate_id % len(self.alpha_beta_pairs)\n span_pair_id = candidate_id // len(self.alpha_beta_pairs)\n return alpha_beta_id, span_pair_id\n\n def get_best(self) -> torch.Tensor:\n \"\"\"Get the best candidate.\n\n Returns:\n `torch.Tensor`:\n The best candidate.\n \"\"\"\n return self.best_scale\n\n def _ask(self) -> torch.Tensor:\n \"\"\"Ask for the next candidate.\n\n Returns:\n `torch.Tensor`:\n The next candidate.\n \"\"\"\n alpha_beta_id, span_pair_id = self._split_candidate_id(self.candidate_id)\n alpha, beta = self.alpha_beta_pairs[alpha_beta_id]\n a_span, b_span = self.span_pairs[span_pair_id]\n if alpha == 0 and beta == 0:\n scale = torch.ones_like(a_span, dtype=self.develop_dtype)\n else:\n scale = get_smooth_scale(alpha_base=a_span, beta_base=b_span, alpha=alpha, beta=beta)\n return scale\n\n def _tell(self, error: list[torch.Tensor]) -> None: # noqa: C901\n \"\"\"Tell the error of the last candidate and update the best candidate.\n\n Args:\n error (`list[torch.Tensor]`):\n The error of the last candidate.\n \"\"\"\n numel = error[0].numel()\n assert all(e.numel() == numel for e in error)\n scale = self.candidate\n self.best_error, self.best_scale = self._update_best(\n best_error=self.best_error,\n best_scale=self.best_scale,\n error=error,\n scale=scale,\n numel=numel,\n num_channels=self.num_in_channels,\n num_heads=self.num_heads,\n num_head_repeats=self.num_head_repeats,\n )\n if self.logger.level <= tools.logging.DEBUG:\n self.error_history.append(\n (\n sum(math.root_(e.to(torch.float64).sum(), self.config.degree).item() for e in error),\n sum(math.root_(b.to(torch.float64).sum(), self.config.degree).item() for b in self.best_error),\n )\n )\n if self.is_last_candidate_in_iter():\n logs: list[list[list[tuple]]] = [[] for _ in range(len(self.span_mode_pairs))]\n for i in range(self.population_size):\n c, r = self._split_candidate_id(i)\n alpha, beta = self.alpha_beta_pairs[c]\n if c % 5 == 0:\n logs[r].append([])\n logs[r][-1].append((alpha, beta, self.error_history[i][0], self.error_history[i][1]))\n for r in range(len(self.span_mode_pairs)):\n self.logger.debug(\n \" - x / w range = %s / %s\", self.span_mode_pairs[r][0].name, self.span_mode_pairs[r][1].name\n )\n for log in logs[r]:\n self.logger.debug(\n \" - alpha = [%s]\",\n \", \".join(f\"{alpha:10.4f}\" for alpha, beta, e, b in log),\n )\n self.logger.debug(\n \" - beta = [%s]\",\n \", \".join(f\"{beta:10.4f}\" for alpha, beta, e, b in log),\n )\n self.logger.debug(\n \" - sum error = [%s]\", \", \".join(f\"{e:10.4f}\" for alpha, beta, e, b in log)\n )\n self.logger.debug(\n \" - best error = [%s]\",\n \", \".join(f\"{b:10.4f}\" for alpha, beta, e, b in log),\n )\n del logs\n self.error_history.clear()\n if self.is_last_iter():\n scale = self.get_best()\n tools.logging.Formatter.indent_dec()\n self.logger.debug(\n \" + error = %.4f\",\n sum(math.root_(b.to(torch.float64).sum(), self.config.degree).item() for b in self.best_error),\n )\n self.logger.debug(\" + scale = [min=%.4f, max=%.4f]\", scale.min().item(), scale.max().item())\n tools.logging.Formatter.indent_inc()\n\n def _reshape_scale(\n self, scale: torch.Tensor, tensor: torch.Tensor, channels_dim: int, needs_reduction: bool = False\n ) -> torch.Tensor:\n if self.num_unique_heads > 0 and needs_reduction:\n scale = scale.view(self.num_unique_heads, self.num_head_repeats, -1)[:, 0, :].reshape(-1)\n shape = [1] * tensor.ndim\n shape[channels_dim] = -1\n return scale.view(shape)\n\n def _process_x_in_xw(self, x: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor:\n if not self.needs_x_quant_for_wgts:\n return x\n if channels_dim is MISSING:\n channels_dim = self.x_quantizer.channels_dim\n shape, dtype = x.shape, x.dtype\n scale = self._reshape_scale(self.candidate, x, channels_dim)\n x = x.to(dtype=self.develop_dtype) if dtype != self.develop_dtype else x.clone()\n x = x.div_(scale)\n x = self.x_quantizer.quantize(\n x, channels_dim=channels_dim, default_dtype=dtype, develop_dtype=self.develop_dtype\n ).data\n x = x.mul_(scale).to(dtype=dtype)\n return x.view(shape)\n\n def _process_w_in_xw(self, w: torch.Tensor) -> torch.Tensor:\n if not self.needs_w_quant_for_wgts:\n return w\n dtype = w.dtype\n channels_dim = 1 if self.w_quantizer.channels_dim is None else self.w_quantizer.channels_dim\n scale = self._reshape_scale(self.candidate, w, channels_dim=channels_dim)\n w = w.to(dtype=self.develop_dtype) if dtype != self.develop_dtype else w.clone()\n w = self.w_quantizer.quantize(\n w.mul_(scale), kernel=None, default_dtype=dtype, develop_dtype=self.develop_dtype\n ).data\n w = w.div_(scale).to(dtype=dtype)\n return w\n\n def _process_x_in_yx(self, x: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor:\n if not self.needs_x_quant_for_opts:\n return x\n shape, dtype = x.shape, x.dtype\n if self.objective != SearchBasedCalibObjective.OutputsError:\n if channels_dim is MISSING:\n channels_dim = self.x_quantizer.channels_dim\n scale = self._reshape_scale(self.candidate, x, channels_dim, needs_reduction=False)\n x = x.to(dtype=self.develop_dtype) if dtype != self.develop_dtype else x.clone()\n x = x.mul_(scale)\n # ! `x` is already scaled during `_process_opts_centric_mod` by scaling `xw`\n x = self.x_quantizer.quantize(\n x,\n channels_dim=channels_dim,\n default_dtype=dtype,\n develop_dtype=self.develop_dtype,\n ).data\n if self.objective != SearchBasedCalibObjective.OutputsError:\n x = x.div_(scale).to(dtype=dtype)\n return x.view(shape)\n\n def _process_y_in_yx(self, y: torch.Tensor, channels_dim: int | _MISSING_TYPE = MISSING) -> torch.Tensor:\n if not self.needs_y_quant_for_opts:\n return y\n shape, dtype = y.shape, y.dtype\n if self.objective != SearchBasedCalibObjective.OutputsError:\n if channels_dim is MISSING:\n channels_dim = self.x_quantizer.channels_dim\n scale = self._reshape_scale(self.candidate, y, channels_dim, needs_reduction=True)\n y = y.to(dtype=self.develop_dtype) if dtype != self.develop_dtype else y.clone()\n y = y.div_(scale)\n # ! `y` is already scaled during `_process_opts_centric_mod` by scaling `yw`\n y = self.y_quantizer.quantize(\n y,\n channels_dim=channels_dim,\n default_dtype=dtype,\n develop_dtype=self.develop_dtype,\n ).data\n if self.objective != SearchBasedCalibObjective.OutputsError:\n y = y.mul_(scale).to(dtype=dtype)\n return y.view(shape)\n\n def _process_xw_in_yx(self, w: torch.Tensor) -> torch.Tensor:\n raise RuntimeError(\"The method `_process_xw_in_yx` should not be called in SmoothCalibrator.\")\n\n def _process_yw_in_yx(self, w: torch.Tensor) -> torch.Tensor:\n raise RuntimeError(\"The method `_process_yw_in_yx` should not be called in SmoothCalibrator.\")\n\n def _process_wgts_centric_mod(\n self,\n wgts: list[nn.Parameter],\n mods: list[nn.Module],\n update_state_dict: bool = True,\n splits: list[int] | None = None,\n **kwargs,\n ) -> None:\n if self.needs_w_quant_for_wgts and self.config.allow_low_rank and self.w_quantizer.is_enabled_low_rank():\n assert len(wgts) == len(mods)\n for wgt in wgts:\n if update_state_dict:\n self._state_dict.append((wgt, wgt.data))\n dtype = wgt.dtype\n scale = self._reshape_scale(self.candidate, wgt.data, channels_dim=1)\n wgt.data = wgt.data.to(dtype=self.develop_dtype).mul(scale).to(dtype=dtype)\n input_packager = self.x_quantizer.get_input_packager() if self.needs_x_quant else None\n for mod in mods:\n self._hooks.append(\n ActivationSmoother(\n self.candidate,\n self.x_quantizer.channels_dim,\n develop_dtype=self.develop_dtype,\n input_packager=input_packager,\n )\n .as_hook()\n .register(mod)\n )\n if splits:\n wgts_splits: list[list[nn.Parameter]] = split_sequence(wgts, splits)\n mods_splits: list[list[nn.Module]] = split_sequence(mods, splits)\n else:\n wgts_splits, mods_splits = [wgts], [mods]\n for wgts_split, mods_split in zip(wgts_splits, mods_splits, strict=True):\n for qwgt, lowr, wgt, mod in zip(\n *self.w_quantizer.quantize_with_low_rank(wgts_split, kernel=None, develop_dtype=self.develop_dtype),\n wgts_split,\n mods_split,\n strict=True,\n ):\n wgt.data = qwgt.data\n self._hooks.append(lowr.as_hook(input_packager=input_packager).register(mod))\n if self.needs_x_quant_for_wgts:\n self._hooks.append(self.x_quantizer.as_hook().register(mod))\n else:\n super()._process_wgts_centric_mod(wgts=wgts, mods=mods, update_state_dict=update_state_dict, **kwargs)\n\n def _process_opts_centric_mod(\n self,\n x_wgts: list[nn.Parameter],\n y_wgts: list[nn.Parameter],\n x_mods: list[nn.Module],\n y_mods: list[nn.Module],\n update_state_dict: bool = True,\n **kwargs,\n ) -> None:\n for w in x_wgts:\n if update_state_dict:\n self._state_dict.append((w, w.data))\n scale = self._reshape_scale(self.candidate, w, channels_dim=0, needs_reduction=False)\n w.data = w.detach().data.to(dtype=self.develop_dtype).mul(scale).to(dtype=w.dtype)\n for w in y_wgts:\n if update_state_dict:\n self._state_dict.append((w, w.data))\n scale = self._reshape_scale(self.candidate, w, channels_dim=0, needs_reduction=True)\n w.data = w.detach().data.to(dtype=self.develop_dtype).div(scale).to(dtype=w.dtype)\n super()._process_opts_centric_mod(\n x_wgts=x_wgts,\n y_wgts=y_wgts,\n x_mods=x_mods,\n y_mods=y_mods,\n update_state_dict=False,\n **kwargs,\n )\n\n @staticmethod\n def _update_best(\n *,\n best_error: list[torch.Tensor] | None,\n best_scale: torch.Tensor,\n error: list[torch.Tensor],\n scale: torch.Tensor,\n numel: int,\n num_channels: int,\n num_heads: int,\n num_head_repeats: int,\n ) -> tuple[list[torch.Tensor], torch.Tensor]:\n if best_error is None:\n return error, scale\n elif numel == 1: # tensor wise quantization error\n if all(e <= b for b, e in zip(best_error, error, strict=True)):\n return error, scale\n return best_error, best_scale\n else: # channel group wise quantization error\n assert num_channels % numel == 0\n group_size, num_groups = num_channels // numel, numel\n needs_reduction = num_heads > 1 and num_head_repeats > 1\n if needs_reduction:\n num_head_channels = num_channels // num_heads\n num_unique_heads = num_heads // num_head_repeats\n if num_head_channels >= group_size:\n assert num_head_channels % group_size == 0\n num_groups_per_head = num_head_channels // group_size\n num_repeats = num_head_repeats\n num_unqiue_heads_per_group = 1\n else:\n assert group_size % num_head_channels == 0\n num_heads_per_group = group_size // num_head_channels\n if num_heads_per_group < num_head_repeats:\n assert num_head_repeats % num_heads_per_group == 0\n num_groups_per_head = 1\n num_repeats = num_head_repeats // num_heads_per_group\n num_unqiue_heads_per_group = 1\n else:\n assert num_heads_per_group % num_head_repeats == 0\n num_groups_per_head = 1\n num_repeats = 1\n num_unqiue_heads_per_group = num_heads_per_group // num_head_repeats\n num_uniques = num_unique_heads // num_unqiue_heads_per_group\n needs_reduction = needs_reduction and num_repeats > 1\n\n pos = torch.full((numel,), True, device=error[0][0].device)\n for b, e in zip(best_error, error, strict=True):\n if needs_reduction:\n b = b.view(num_uniques, num_repeats, num_groups_per_head).sum(dim=1, keepdim=True)\n e = e.view(num_uniques, num_repeats, num_groups_per_head).sum(dim=1, keepdim=True)\n pos = pos & (e < b).expand(num_uniques, num_repeats, num_groups_per_head).reshape_as(pos)\n else:\n pos = pos & (e < b)\n for b, e in zip(best_error, error, strict=True):\n b[pos] = e[pos]\n pos = pos.view(num_groups, 1).expand(num_groups, group_size)\n best_scale = best_scale.view(num_groups, group_size)\n best_scale[pos] = scale.view(num_groups, group_size)[pos]\n return best_error, best_scale\n\n\nclass SmoothLinearCalibrator(SmoothCalibrator):\n \"\"\"The smooth quantization calibrator for linear module.\"\"\"\n\n def __init__(\n self,\n config: SmoothCalibConfig,\n weight_quantizer: Quantizer | None,\n input_quantizer: Quantizer | None,\n num_heads: int = 1,\n num_head_repeats: int = 1,\n develop_dtype: torch.dtype = torch.float32,\n ) -> None:\n \"\"\"Initialize the calibrator.\n\n Args:\n config (`SmoothCalibConfig`):\n The quantization smoothing calibration configuration.\n weight_quantizer (`Quantizer` or `None`):\n The weight quantizer.\n input_quantizer (`Quantizer` or `None`):\n The input quantizer.\n num_heads (`int`, *optional*, defaults to `1`):\n The number of heads.\n num_head_repeats (`int`, *optional*, defaults to `1`):\n The number of head repeats.\n develop_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The development data type.\n \"\"\"\n super().__init__(\n tensor_type=TensorType.Weights,\n config=config,\n w_quantizer=weight_quantizer,\n x_quantizer=input_quantizer,\n y_quantizer=None,\n num_heads=num_heads,\n num_head_repeats=num_head_repeats,\n develop_dtype=develop_dtype,\n )\n\n\nclass SmoothAttentionCalibrator(SmoothCalibrator):\n \"\"\"The smooth quantization calibrator for attention module.\"\"\"\n\n def __init__(\n self,\n config: SmoothCalibConfig,\n query_quantizer: Quantizer | None,\n key_quantizer: Quantizer | None,\n num_heads: int = 1,\n num_head_repeats: int = 1,\n with_rope: bool = True,\n develop_dtype: torch.dtype = torch.float32,\n ) -> None:\n \"\"\"Initialize the calibrator.\n\n Args:\n config (`SmoothCalibConfig`):\n The quantization smoothing calibration configuration.\n query_quantizer (`Quantizer` or `None`):\n The query quantizer.\n key_quantizer (`Quantizer` or `None`):\n The key quantizer.\n num_heads (`int`, *optional*, defaults to `1`):\n The number of heads.\n num_head_repeats (`int`, *optional*, defaults to `1`):\n The number of head repeats.\n with_rope (`bool`, *optional*, defaults to `True`):\n Whether rotary position embedding is used.\n develop_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The development data type.\n \"\"\"\n super().__init__(\n tensor_type=TensorType.Outputs,\n config=config,\n w_quantizer=None,\n x_quantizer=query_quantizer,\n y_quantizer=key_quantizer,\n num_heads=num_heads,\n num_head_repeats=num_head_repeats,\n with_rope=with_rope,\n develop_dtype=develop_dtype,\n )\n\n def calibrate(\n self,\n q_proj_weight: nn.Parameter,\n k_proj_weight: nn.Parameter,\n queries: TensorsCache,\n keys: TensorsCache,\n query_module: nn.Module,\n key_module: nn.Module,\n eval_module: nn.Module | None = None,\n eval_inputs: TensorsCache | None = None,\n eval_kwargs: dict[str, tp.Any] | None = None,\n ) -> tp.Any:\n \"\"\"Calibrate the quantization for attention.\n\n Args:\n q_proj_weight (`nn.Parameter`):\n The query projection weight.\n k_proj_weight (`nn.Parameter`):\n The key projection weight.\n queries (`TensorsCache`):\n The query activations.\n keys (`TensorsCache`):\n The key activations.\n query_module (`nn.Module`):\n The module that generates the query activations,\n e.g., either `q_proj` for pre-rope or `q_rotary_emb` for post-rope.\n key_module (`nn.Module`):\n The module that generates the key activations,\n e.g., either `k_proj` for pre-rope or `k_rotary_emb` for post-rope.\n eval_module (`nn.Module`, *optional*):\n The evaluation module.\n eval_inputs (`TensorsCache`, *optional*):\n The evaluation inputs.\n eval_kwargs (`dict[str, tp.Any]`, *optional*):\n The evaluation keyword arguments.\n\n Returns:\n tp.Any: The evaluation result.\n \"\"\"\n return super().calibrate(\n x_wgts=[q_proj_weight],\n y_wgts=[k_proj_weight],\n x_acts=queries,\n y_acts=keys,\n x_mods=[query_module],\n y_mods=[key_module],\n eval_module=eval_module,\n eval_inputs=eval_inputs,\n eval_kwargs=eval_kwargs,\n )\n\n\ndef smooth_upscale_param(param: nn.Parameter, scale: torch.Tensor, channels_dim: int = 1) -> None:\n \"\"\"In-place smooth the parameter by upscaling.\n\n Args:\n param (`nn.Parameter`):\n The parameter to smooth.\n scale (`torch.Tensor`):\n The scale to upscale.\n channels_dim (`int`, *optional*, defaults to `1`):\n The dimension of channels.\n \"\"\"\n dtype = param.dtype\n view_shape = [1] * param.ndim\n view_shape[channels_dim] = -1\n scale = scale.to(device=param.device).view(view_shape)\n param.data = param.data.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype)\n assert not param.data.isnan().any(), \"NaN found in param when smoothing\"\n assert not param.data.isinf().any(), \"Inf found in param when smoothing\"\n\n\ndef smooth_downscale_param(param: nn.Parameter, scale: torch.Tensor, channels_dim: int = 0) -> None:\n \"\"\"In-place smooth the parameter by downscaling.\n\n Args:\n param (`nn.Parameter`):\n The parameter to smooth.\n scale (`torch.Tensor`):\n The scale to downscale.\n channels_dim (`int`, *optional*, defaults to `0`):\n The dimension of channels.\n \"\"\"\n dtype = param.dtype\n view_shape = [1] * param.ndim\n view_shape[channels_dim] = -1\n scale = scale.to(device=param.device).view(view_shape)\n param_data = param.data.to(dtype=scale.dtype)\n param_data.narrow(channels_dim, 0, scale.numel()).div_(scale)\n param.data = param_data.to(dtype=dtype)\n assert not param.data.isnan().any(), \"NaN found in param when smoothing\"\n assert not param.data.isinf().any(), \"Inf found in param when smoothing\"\n\n\ndef convert_smooth_upscale_to_downscale(\n scale: torch.Tensor, num_heads: int = 1, num_head_repeats: int = 1\n) -> torch.Tensor:\n \"\"\"Convert the upscale smooth scale to downscale smooth scale.\n\n Args:\n scale (`torch.Tensor`):\n The upscale smooth scale.\n num_heads (`int`, *optional*, defaults to `1`):\n The number of heads.\n num_head_repeats (`int`, *optional*, defaults to `1`):\n The number of head repeats.\n\n Returns:\n `torch.Tensor`:\n The downscale smooth scale.\n \"\"\"\n if num_heads > 1 and num_head_repeats > 1:\n head_channels = scale.numel() // num_heads\n num_unique_heads = num_heads // num_head_repeats\n return scale.view(num_unique_heads, num_head_repeats, head_channels)[:, 0, :].reshape(-1)\n else:\n return scale\n\n\n@torch.inference_mode()\ndef smooth_linear_modules(\n prevs: nn.Module | tp.Sequence[nn.Module] | None,\n modules: tp.Sequence[nn.Linear] | nn.Linear,\n *,\n scale: torch.Tensor | None,\n config: SmoothCalibConfig | None = None,\n weight_quantizer: Quantizer | None = None,\n input_quantizer: Quantizer | None = None,\n weights: list[nn.Parameter] | None = None,\n inputs: TensorsCache | None = None,\n eval_inputs: TensorsCache | None = None,\n eval_module: nn.Module = None,\n eval_kwargs: dict[str, tp.Any] = None,\n num_heads: int = 1,\n num_head_repeats: int = 1,\n splits: list[int] | None = None,\n extra_modules: list[nn.Linear] | None = None,\n develop_dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Smooth two consecutive modules.\n\n Args:\n prevs (`nn.Module` or `list[nn.Module]`):\n The first module(s).\n modules (`nn.Linear` or `list[nn.Linear]`):\n The second module(s).\n scale (`torch.Tensor` or `None`, *optional*, defaults to `None`):\n The smooth quantization scale.\n config (`SmoothCalibConfig` or `None`, *optional*, defaults to `None`):\n The smooth quantization configuration.\n weight_quantizer (`Quantizer` or `None`, *optional*, defaults to `None`):\n The quantizer for weights.\n input_quantizer (`Quantizer` or `None`, *optional*, defaults to `None`):\n The quantizer for inputs.\n weights (`list[nn.Parameter]` or `None`, *optional*, defaults to `None`):\n The weights of the modules. If `None`, the weights of the modules will be used.\n inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The cache of the input activations.\n eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The cache of the inputs corresponding to the `eval_module`.\n eval_module (`nn.Module`, *optional*, defaults to `None`):\n The module to evaluate the quantization error.\n eval_kwargs (`dict[str, tp.Any]`, *optional*, defaults to `None`):\n The keyword arguments for evaluation.\n num_heads (`int`, *optional*, defaults to `1`):\n The number of heads.\n num_head_repeats (`int`, *optional*, defaults to `1`):\n The number of head repeats.\n extra_modules (`list[nn.Module]` or `None`, *optional*, defaults to `None`):\n Extra modules to smooth.\n develop_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The development data type.\n\n Returns:\n `torch.Tensor`:\n The smooth quantization scale in CPU.\n \"\"\"\n if not isinstance(modules, (list, tuple)):\n modules = [modules]\n extra_modules = [] if extra_modules is None else extra_modules\n if scale is None:\n assert inputs is not None or eval_inputs is not None, \"inputs or eval_inputs must be provided\"\n scale = SmoothLinearCalibrator(\n config=config,\n weight_quantizer=weight_quantizer,\n input_quantizer=input_quantizer,\n num_heads=num_heads,\n num_head_repeats=num_head_repeats,\n develop_dtype=develop_dtype,\n ).calibrate(\n x_wgts=[module.weight for module in modules] if weights is None else weights,\n x_acts=inputs,\n x_mods=modules,\n eval_inputs=eval_inputs,\n eval_module=eval_module,\n eval_kwargs=eval_kwargs,\n splits=splits,\n )\n gc.collect()\n torch.cuda.empty_cache()\n upscale = scale\n for module in modules + extra_modules:\n upscale = upscale.to(device=module.weight.device)\n smooth_upscale_param(module.weight, upscale, channels_dim=1)\n if prevs is not None:\n downscale = convert_smooth_upscale_to_downscale(upscale, num_heads=num_heads, num_head_repeats=num_head_repeats)\n if isinstance(prevs, nn.Module):\n prevs = [prevs]\n for module in prevs:\n if module is None:\n continue\n downscale = downscale.to(device=module.weight.device)\n smooth_downscale_param(module.weight, downscale, channels_dim=0)\n if hasattr(module, \"bias\") and module.bias is not None:\n smooth_downscale_param(module.bias, downscale, channels_dim=0)\n return scale.to(device=\"cpu\")\n\n\n@torch.inference_mode()\ndef smooth_attention(\n *,\n q_proj: nn.Linear,\n k_proj: nn.Linear,\n scale: torch.Tensor | None,\n config: SmoothCalibConfig | None = None,\n query_quantizer: Quantizer | None = None,\n key_quantizer: Quantizer | None = None,\n queries: TensorsCache | None = None,\n keys: TensorsCache | None = None,\n attn_q: nn.Module | None = None,\n attn_k: nn.Module | None = None,\n eval_inputs: TensorsCache | None = None,\n eval_module: nn.Module = None,\n eval_kwargs: dict[str, tp.Any] = None,\n num_heads: int = 1,\n num_head_repeats: int = 1,\n with_rope: bool = True,\n develop_dtype: torch.dtype = torch.float32,\n) -> torch.Tensor:\n \"\"\"Smooth attention.\n\n Args:\n q_proj (`nn.Linear`):\n The query projection module.\n k_proj (`nn.Linear`):\n The key projection module.\n scale (`torch.Tensor` or `None`, *optional*, defaults to `None`):\n The smooth quantization scale.\n config (`SmoothCalibConfig` or `None`, *optional*, defaults to `None`):\n The smooth quantization configuration.\n query_quantizer (`Quantizer` or `None`, *optional*, defaults to `None`):\n The quantizer for queries.\n key_quantizer (`Quantizer` or `None`, *optional*, defaults to `None`):\n The quantizer for keys.\n queries (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The cache of the queries.\n keys (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The cache of the keys.\n attn_q (`nn.Module` or `None`, *optional*, defaults to `None`):\n The module that generates the queries.\n attn_k (`nn.Module` or `None`, *optional*, defaults to `None`):\n The module that generates the keys.\n eval_inputs (`TensorsCache` or `None`, *optional*, defaults to `None`):\n The cache of the inputs corresponding to the evaluation module.\n eval_module (`nn.Module`, *optional*, defaults to `None`):\n The module to evaluate the quantization error.\n eval_kwargs (`dict[str, tp.Any]`, *optional*, defaults to `None`):\n The keyword arguments for evaluation.\n num_heads (`int`, *optional*, defaults to `1`):\n The number of heads.\n num_head_repeats (`int`, *optional*, defaults to `1`):\n The number of head repeats.\n with_rope (`bool`, *optional*, defaults to `True`):\n Whether quantization is applied after rotary position embedding.\n develop_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):\n The development data type.\n\n Returns:\n `torch.Tensor`:\n The smooth quantization scale in CPU.\n \"\"\"\n if scale is None:\n assert queries is not None and keys is not None and eval_inputs is not None\n assert attn_q is not None and attn_k is not None, \"modules must be provided\"\n scale = SmoothAttentionCalibrator(\n config=config,\n query_quantizer=query_quantizer,\n key_quantizer=key_quantizer,\n num_heads=num_heads,\n num_head_repeats=num_head_repeats,\n with_rope=with_rope,\n develop_dtype=develop_dtype,\n ).calibrate(\n q_proj_weight=q_proj.weight,\n k_proj_weight=k_proj.weight,\n queries=queries,\n keys=keys,\n query_module=attn_q,\n key_module=attn_k,\n eval_inputs=eval_inputs,\n eval_module=eval_module,\n eval_kwargs=eval_kwargs,\n )\n gc.collect()\n torch.cuda.empty_cache()\n upscale = scale.to(device=q_proj.weight.device)\n smooth_upscale_param(q_proj.weight, upscale, channels_dim=0)\n downscale = convert_smooth_upscale_to_downscale(upscale, num_heads=num_heads, num_head_repeats=num_head_repeats)\n smooth_downscale_param(k_proj.weight, downscale, channels_dim=0)\n return scale.to(device=\"cpu\")\n"
},
{
"path": "deepcompressor/csrc/load.py",
"content": "# -*- coding: utf-8 -*-\n\"\"\"Deepcompressor Extension.\"\"\"\n\nimport os\n\nfrom torch.utils.cpp_extension import load\n\n__all__ = [\"_C\"]\n\ndirpath = os.path.dirname(__file__)\n\n_C = load(\n name=\"deepcompressor_C\",\n sources=[f\"{dirpath}/pybind.cpp\", f\"{dirpath}/quantize/quantize.cu\"],\n extra_cflags=[\"-g\", \"-O3\", \"-fopenmp\", \"-lgomp\", \"-std=c++20\"],\n extra_cuda_cflags=[\n \"-O3\",\n \"-std=c++20\",\n \"-U__CUDA_NO_HALF_OPERATORS__\",\n \"-U__CUDA_NO_HALF_CONVERSIONS__\",\n \"-U__CUDA_NO_HALF2_OPERATORS__\",\n \"-U__CUDA_NO_HALF2_CONVERSIONS__\",\n \"-U__CUDA_NO_BFLOAT16_OPERATORS__\",\n \"-U__CUDA_NO_BFLOAT16_CONVERSIONS__\",\n \"-U__CUDA_NO_BFLOAT162_OPERATORS__\",\n \"-U__CUDA_NO_BFLOAT162_CONVERSIONS__\",\n \"--expt-relaxed-constexpr\",\n \"--expt-extended-lambda\",\n \"--use_fast_math\",\n \"--ptxas-options=--allow-expensive-optimizations=true\",\n \"--threads=8\",\n ],\n)\n"
},
{
"path": "deepcompressor/csrc/pybind.cpp",
"content": "#include \n#include \n\n#include \"quantize/quantize.h\"\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"round_to_nearest_in_codebook_cuda\", &round_to_nearest_in_codebook_cuda,\n py::arg(\"tensor\"), py::arg(\"codebook\"), py::arg(\"inplace\") = false,\n py::arg(\"bnb\") = false, \"RTN with codebook (CUDA)\");\n}\n"
},
{
"path": "deepcompressor/csrc/quantize/quantize.cu",
"content": "\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"quantize.h\"\n\n// The following code is adapted from the bitsandbytes library:\n// https://github.com/bitsandbytes-foundation/bitsandbytes/blob/main/csrc/kernels.cu#L232\ntemplate \n__device__ __forceinline__\n typename std::conditional::type\n bnb_nearest_neighbor(float_t x, float_t *codebook, const int C)\n{\n int mid = (C >> 1) - 1;\n int hi = C - 1;\n int lo = 0;\n\n float_t lval = codebook[lo];\n float_t hval = codebook[hi];\n float_t mval = codebook[mid];\n\n for (int step = (C >> 2); step > 0; step >>= 1)\n {\n if (x > mval)\n {\n lo = mid;\n lval = mval;\n mid += step;\n }\n else\n {\n hi = mid;\n hval = mval;\n mid -= step;\n }\n mval = codebook[mid];\n }\n\n if (x > mval)\n {\n if constexpr (ret_val)\n {\n return (x - mval > hval - x) ? hval : mval;\n }\n else\n {\n return (x - mval > hval - x) ? hi : mid;\n }\n }\n else\n {\n if constexpr (ret_val)\n {\n return (x - lval < mval - x) ? lval : mval;\n }\n else\n {\n return (x - lval < mval - x) ? lo : mid;\n }\n }\n}\n\ntemplate \n__device__ __forceinline__\n typename std::conditional::type\n nearest_neighbor(float_t x, const float_t *codebook, int C)\n{\n int lo = 0;\n int bit = 1 << (31 - __clz(C));\n\n float_t lval = codebook[lo];\n while (bit)\n {\n int next = lo | bit;\n float_t nval = codebook[next];\n bool pred = next < C && nval <= x;\n lo = pred ? next : lo;\n lval = pred ? nval : lval;\n bit >>= 1;\n }\n\n int hi = lo + (lo < C - 1);\n float_t hval = codebook[hi];\n\n if constexpr (ret_val)\n {\n return (x + x < lval + hval) ? lval : hval;\n }\n else\n {\n return (x + x < lval + hval) ? lo : hi;\n }\n}\n\n// CUDA kernel: Each thread processes one element from x and finds the nearest\n// codebook entry. The codebook (of size C < 256) is first loaded into shared\n// memory.\ntemplate \n__global__ void round_to_nearest_in_codebook_kernel(\n const float_t *__restrict__ x, const float_t *__restrict__ codebook,\n float_t *__restrict__ y, const int N, const int C)\n{\n // Use a shared memory array for the codebook.\n __shared__ float_t s_codebook[256];\n\n // Have the first few threads load the codebook into shared memory.\n for (int i = threadIdx.x; i < C; i += blockDim.x)\n {\n s_codebook[i] = codebook[i];\n }\n __syncthreads();\n\n // Global index for the element processed by this thread.\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < N)\n {\n if constexpr (bnb)\n {\n y[idx] = bnb_nearest_neighbor(x[idx], s_codebook, C);\n }\n else\n {\n y[idx] = nearest_neighbor(x[idx], s_codebook, C);\n }\n }\n}\n\ntorch::Tensor round_to_nearest_in_codebook_cuda(torch::Tensor tensor,\n torch::Tensor codebook,\n bool inplace, bool bnb)\n{\n auto x = tensor.contiguous();\n auto c = codebook.contiguous();\n auto y = inplace ? x : torch::empty_like(tensor);\n const int N = x.numel();\n const int C = c.numel();\n const int threads = 256;\n const int blocks = (N + threads - 1) / threads;\n AT_DISPATCH_FLOATING_TYPES(\n tensor.scalar_type(), \"round_to_nearest_in_codebook_cuda\", [&]\n {\n if (bnb && (C & (C - 1)) == 0) {\n round_to_nearest_in_codebook_kernel\n <<>>(x.data_ptr(),\n c.data_ptr(),\n y.data_ptr