├── .dev-scripts
├── .gitignore
├── README.md
├── basic_tests.sh
├── extract_install_cmd.py
├── publish-docker-internal.sh
├── publish-with-docker.sh
└── publish.sh
├── .gitignore
├── CHANGELOG.md
├── CONTRIBUTING.md
├── LICENSE
├── README.md
├── bitorch_engine
├── __init__.py
├── extensions
│ └── __init__.py
├── functions
│ ├── __init__.py
│ └── cuda
│ │ ├── __init__.py
│ │ ├── extension.py
│ │ ├── functions.py
│ │ ├── functions_cuda.cpp
│ │ └── functions_cuda_kernel.cu
├── layers
│ ├── __init__.py
│ ├── qconv
│ │ ├── __init__.py
│ │ ├── binary
│ │ │ ├── __init__.py
│ │ │ ├── cpp
│ │ │ │ ├── __init__.py
│ │ │ │ ├── binary_conv.cpp
│ │ │ │ ├── extension.py
│ │ │ │ └── layer.py
│ │ │ ├── cutlass
│ │ │ │ ├── __init__.py
│ │ │ │ ├── binary_conv2d_cutlass.cpp
│ │ │ │ ├── binary_conv2d_cutlass_kernel.cu
│ │ │ │ ├── extension.py
│ │ │ │ └── layer.py
│ │ │ └── layer.py
│ │ └── nbit
│ │ │ ├── __init__.py
│ │ │ ├── cutlass
│ │ │ ├── __init__.py
│ │ │ ├── extension.py
│ │ │ ├── layer.py
│ │ │ ├── q4_conv_cutlass.cpp
│ │ │ └── q4_conv_cutlass_kernel.cu
│ │ │ └── layer.py
│ ├── qembedding
│ │ ├── __init__.py
│ │ └── binary
│ │ │ ├── __init__.py
│ │ │ └── layer.py
│ ├── qlinear
│ │ ├── __init__.py
│ │ ├── binary
│ │ │ ├── __init__.py
│ │ │ ├── binary_implementation.py
│ │ │ ├── cpp
│ │ │ │ ├── __init__.py
│ │ │ │ ├── binary_linear.cpp
│ │ │ │ ├── extension.py
│ │ │ │ └── layer.py
│ │ │ ├── cuda
│ │ │ │ ├── __init__.py
│ │ │ │ ├── binary_linear_cuda.cpp
│ │ │ │ ├── binary_linear_cuda_kernel.cu
│ │ │ │ ├── bmm.py
│ │ │ │ ├── extension.py
│ │ │ │ └── layer.py
│ │ │ ├── cutlass
│ │ │ │ ├── __init__.py
│ │ │ │ ├── binary_linear_cutlass.cpp
│ │ │ │ ├── binary_linear_cutlass_kernel.cu
│ │ │ │ ├── binary_linear_cutlass_kernel.h
│ │ │ │ ├── extension.py
│ │ │ │ ├── kernel_selection.h
│ │ │ │ └── layer.py
│ │ │ └── layer.py
│ │ ├── layer.py
│ │ ├── nbit
│ │ │ ├── __init__.py
│ │ │ ├── cuda
│ │ │ │ ├── __init__.py
│ │ │ │ ├── exl2
│ │ │ │ │ ├── compat.cuh
│ │ │ │ │ ├── config.h
│ │ │ │ │ ├── kernel_select.cuh
│ │ │ │ │ ├── matrix_view.cuh
│ │ │ │ │ ├── q_gemm_kernel.cuh
│ │ │ │ │ ├── q_gemm_kernel_gptq.cuh
│ │ │ │ │ ├── quant
│ │ │ │ │ │ ├── qdq_2.cuh
│ │ │ │ │ │ ├── qdq_3.cuh
│ │ │ │ │ │ ├── qdq_4.cuh
│ │ │ │ │ │ ├── qdq_5.cuh
│ │ │ │ │ │ ├── qdq_6.cuh
│ │ │ │ │ │ ├── qdq_8.cuh
│ │ │ │ │ │ └── qdq_util.cuh
│ │ │ │ │ └── util.cuh
│ │ │ │ ├── extension.py
│ │ │ │ ├── mbwq_layer.py
│ │ │ │ ├── mbwq_linear_cuda_kernel.cu
│ │ │ │ ├── mpq_layer.py
│ │ │ │ ├── mpq_linear_cuda_kernel.cu
│ │ │ │ ├── q_linear_cuda.cpp
│ │ │ │ └── utils.py
│ │ │ ├── cutlass
│ │ │ │ ├── __init__.py
│ │ │ │ ├── extension.py
│ │ │ │ ├── q4_layer.py
│ │ │ │ ├── q4_linear_cutlass_kernel.cu
│ │ │ │ ├── q8_layer.py
│ │ │ │ ├── q8_linear_cutlass_kernel.cu
│ │ │ │ └── q_linear_cutlass.cpp
│ │ │ ├── layer.py
│ │ │ └── mps
│ │ │ │ ├── __init__.py
│ │ │ │ ├── extension.py
│ │ │ │ ├── mlx_bindings.cpp
│ │ │ │ ├── mpq_layer.py
│ │ │ │ ├── mpq_linear_mlx.cpp
│ │ │ │ └── mpq_linear_mlx.h
│ │ └── qlinear_implementation.py
│ └── qmha
│ │ ├── __init__.py
│ │ └── binary
│ │ ├── __init__.py
│ │ └── layer.py
├── optim
│ ├── __init__.py
│ ├── diode_beta.py
│ └── galore_projector.py
└── utils
│ ├── __init__.py
│ ├── arch_helper.py
│ ├── convert.py
│ ├── cpp_extension.py
│ ├── cuda_extension.py
│ ├── cutlass_path.py
│ ├── mlx_extension.py
│ ├── mlx_path.py
│ ├── model_helper.py
│ ├── quant_operators.py
│ └── safe_import.py
├── docker
├── Dockerfile
├── README.md
└── build_scripts
│ └── install_modified_pytorch.sh
├── docs
├── .gitignore
├── Makefile
├── README.md
├── make_docs.sh
├── requirements.txt
├── scripts
│ └── convert_docs.py
└── source
│ ├── _templates
│ ├── class.rst
│ └── module.rst
│ ├── build_options.rst
│ ├── conf.py
│ ├── documentation.rst
│ ├── index.rst
│ └── installation.rst
├── examples
├── __init__.py
├── mnist-lightning
│ ├── main.py
│ ├── mlp.py
│ └── requirements.txt
└── mnist
│ ├── README.md
│ ├── __init__.py
│ ├── datasets
│ ├── __init__.py
│ ├── base.py
│ ├── dummy_dataset.py
│ └── mnist.py
│ ├── requirements.txt
│ └── train_mnist.py
├── licenses
├── LICENSE.GPTQ-for-LLaMa.txt
├── LICENSE.cutlass.txt
├── LICENSE.exllamav2.txt
├── LICENSE.mlx.txt
├── LICENSE.pytorch.txt
└── LICENSE.tcbnn.txt
├── requirements-dev.txt
├── requirements.txt
├── setup.py
├── structure.md
├── tests
├── __init__.py
├── functions
│ ├── __init__.py
│ └── test_quant_ops.py
├── layers
│ ├── __init__.py
│ ├── test_binary_conv.py
│ ├── test_binary_embedding.py
│ ├── test_binary_linear.py
│ ├── test_custom_binary_linear.py
│ ├── test_nbit_conv.py
│ ├── test_nbit_linear.py
│ ├── test_nbit_linear_mixbits.py
│ ├── test_nbit_linear_mps.py
│ └── util.py
└── util
│ ├── __init__.py
│ └── binary_mse_loss.py
└── version.txt
/.dev-scripts/.gitignore:
--------------------------------------------------------------------------------
1 | test_*.sh
2 |
--------------------------------------------------------------------------------
/.dev-scripts/README.md:
--------------------------------------------------------------------------------
1 | # Development Scripts
2 |
3 | To publish a binary release with docker (for CUDA), run (replace engine version and cuda version):
4 | ```bash
5 | ./.dev-scripts/publish-with-docker.sh v0.2.3 11.8
6 | ```
7 |
8 | ## Fixing RPaths
9 |
10 | See [this github issue](https://github.com/pytorch/builder/issues/468).
11 |
--------------------------------------------------------------------------------
/.dev-scripts/basic_tests.sh:
--------------------------------------------------------------------------------
1 |
2 | # bash code to run after installation to test for correct package installation
3 |
4 | if [ -n "${SKIP_TESTS}" ]; then
5 | echo "Skipping tests. Done."
6 | exit 0
7 | fi
8 |
9 | # try basic importing, should detect errors of .so loading
10 | python -c "from bitorch_engine.layers.qlinear.binary.cpp import BinaryLinearCPP"
11 | python -c "from bitorch_engine.layers.qembedding.binary import BinaryEmbeddingCuda"
12 | python -c "from bitorch_engine.layers.qlinear.nbit.cutlass import Q4LinearCutlass, Q8LinearCutlass, Q4MatMul"
13 | python -c "from bitorch_engine.layers.qlinear.nbit.cuda import MPQLinearCuda, MBWQLinearCuda"
14 | python -c "from bitorch_engine.layers.qlinear.nbit.cuda.utils import pack_fp_weight, unpack_qweight"
15 | echo "Imports successful!"
16 |
17 | set +o errexit
18 | echo "Testing..."
19 | (
20 | rm -rf bitorch_install_tmp_test_dir
21 | mkdir bitorch_install_tmp_test_dir
22 | cd bitorch_install_tmp_test_dir
23 | git clone https://github.com/GreenBitAI/bitorch-engine.git --depth 1 --branch "v0.2.6" bitorch_engine_git
24 | mv bitorch_engine_git/tests .
25 | pip install pytest numpy
26 | pytest tests/layers/test_nbit_linear.py
27 | pytest tests/layers/test_nbit_linear_mixbits.py
28 | pytest tests/functions/test_quant_ops.py
29 | pytest tests/layers/test_binary_linear.py
30 | )
31 | rm -rf bitorch_install_tmp_test_dir
32 |
--------------------------------------------------------------------------------
/.dev-scripts/extract_install_cmd.py:
--------------------------------------------------------------------------------
1 | # take caution: everything is quite hardcoded here
2 | # any changes to the readme could break this code
3 | # run it from root directory: python extract_install_cmd.py path/to/custom/torch-xxx.whl
4 |
5 | import argparse
6 | parser = argparse.ArgumentParser()
7 | parser.add_argument("custom_pytorch_path", help="Path to custom PyTorch wheel")
8 | parser.add_argument("custom_bitorch_engine_path", help="Path to built bitorch engine wheel file")
9 | args = parser.parse_args()
10 |
11 | BLOCK_HEADER_START_BINARY = "### Binary Release"
12 | BLOCK_HEADER_START_FROM_SOURCE = "#### Conda on Linux"
13 | BLOCK_END = "##########"
14 |
15 | with open("README.md") as infile:
16 | content = infile.readlines()
17 |
18 | with open(".dev-scripts/basic_tests.sh") as infile:
19 | test_appendix = infile.readlines()
20 |
21 |
22 | def write_file(filepath, main_content):
23 | with open(filepath, "w") as outfile:
24 | outfile.write(FILE_INTRO)
25 | outfile.writelines(main_content)
26 | outfile.writelines(test_appendix)
27 |
28 |
29 | source_local_install_instructions = []
30 | source_global_install_instructions = []
31 | binary_local_install_instructions = []
32 | binary_global_install_instructions = []
33 |
34 | in_code_block = False
35 | reading_instructions = False
36 | insert_block_pause = False
37 | instruction_type = ""
38 |
39 | FILE_INTRO = """#!/usr/bin/env bash
40 |
41 | trap exit INT
42 | set -o errexit
43 | set -o xtrace
44 |
45 | """
46 | EXTRA_CONDA_INSTRUCTION = """# extra step for bash script (not required in a proper command line):
47 | eval "$(conda shell.bash hook)"
48 | """
49 |
50 |
51 | for line in content:
52 | if line.startswith("```"):
53 | in_code_block = not in_code_block
54 | continue
55 | if line.startswith(BLOCK_HEADER_START_FROM_SOURCE):
56 | reading_instructions = True
57 | instruction_type = "source-global"
58 | BLOCK_END = BLOCK_HEADER_START_FROM_SOURCE.split()[0]
59 | continue
60 | if line.startswith(BLOCK_HEADER_START_BINARY):
61 | reading_instructions = True
62 | instruction_type = "binary-global"
63 | BLOCK_END = BLOCK_HEADER_START_BINARY.split()[0]
64 | continue
65 | if line.startswith(""):
66 | if "source" in instruction_type:
67 | instruction_type = "source-local"
68 | if "binary" in instruction_type:
69 | instruction_type = "binary-local"
70 | continue
71 | if line.startswith("
" in lines[line_num + 2]:
108 | if self.verbose:
109 | print("Active:", line)
110 | self.active = True
111 | if " " in line:
112 | if self.verbose:
113 | print("Inactive:", line)
114 | self.active = False
115 | self.replace_line(lines, line_num, "")
116 |
117 | def internal_process_line(self, lines, line_num):
118 | stripped_line = lines[line_num].strip()
119 | if stripped_line in ["", ".. raw:: html", " ", ""]:
120 | self.replace_line(lines, line_num, "")
121 | return
122 | if stripped_line == "":
123 | summary_line = line_num + 1
124 | while lines[summary_line].strip() == "":
125 | summary_line += 1
126 | self.replace_line(lines, line_num, ".. dropdown:: " + lines[summary_line])
127 | for i in range(line_num + 1, summary_line + 1):
128 | self.replace_line(lines, i, "")
129 | return
130 | self.replace_line(lines, line_num, " " + lines[line_num])
131 |
132 | def main(args):
133 | # set build_options_separate=False to integrate build options into installation part:
134 | build_options_separate = True
135 | verbose = False
136 |
137 | install_section = SectionExtractor("Installation", verbose=verbose)
138 | build_section = SectionExtractor("Build options", modify_by=-1 if build_options_separate else 0, verbose=verbose)
139 | section_collapse_fixer = SectionCollapseFixer(verbose=verbose)
140 | replacer = Replacer([
141 | [
142 | r"`docker readme `__",
143 | r"`docker readme `__",
144 | ],
145 | ], verbose=verbose)
146 |
147 | content = None
148 | with open(args.input, "r") as f:
149 | content = f.readlines()
150 | if content is None:
151 | return
152 |
153 | i = 0
154 | while i < len(content):
155 | for converter in [install_section, build_section, section_collapse_fixer, replacer]:
156 | converter.process_line(content, i)
157 | i += 1
158 |
159 | install_section.write_to_file(content, "docs/source/installation.rst")
160 | if build_options_separate:
161 | build_section.write_to_file(content, "docs/source/build_options.rst")
162 | else:
163 | build_section.write_to_file(content, "docs/source/installation_test.rst", mode="a")
164 |
165 |
166 | if __name__ == "__main__":
167 | parser = argparse.ArgumentParser()
168 | parser.add_argument("input", help="Input file")
169 | main(parser.parse_args())
170 |
--------------------------------------------------------------------------------
/docs/source/_templates/class.rst:
--------------------------------------------------------------------------------
1 | {{ fullname | escape | underline}}
2 |
3 | .. currentmodule:: {{ module }}
4 |
5 | .. autoclass:: {{ objname }}
6 | :members:
7 | :no-inherited-members:
8 | :special-members: __call__, __add__, __mul__, forward, __init__
9 |
10 | {% block methods %}
11 | {% if methods %}
12 | .. rubric:: {{ _('Methods') }}
13 |
14 | .. autosummary::
15 | :nosignatures:
16 | {% for item in methods %}
17 |
18 | {%- if item.startswith('__init__') %}
19 | {%- if item not in inherited_members %}
20 | ~{{ name }}.{{ item }}
21 | {%- endif -%}
22 | {%- endif -%}
23 | {%- if not item.startswith('_') %}
24 | {%- if item not in inherited_members %}
25 | ~{{ name }}.{{ item }}
26 | {%- endif -%}
27 | {%- endif -%}
28 | {%- endfor %}
29 | {% endif %}
30 | {% endblock %}
31 |
32 | {% block attributes %}
33 | {% if attributes %}
34 | .. rubric:: {{ _('Attributes') }}
35 |
36 | .. autosummary::
37 | {% for item in attributes %}
38 | {%- if item not in inherited_members %}
39 | ~{{ name }}.{{ item }}
40 | {%- endif -%}
41 | {%- endfor %}
42 | {% endif %}
43 | {% endblock %}
44 |
--------------------------------------------------------------------------------
/docs/source/_templates/module.rst:
--------------------------------------------------------------------------------
1 | {{ fullname | escape | underline}}
2 |
3 | .. automodule:: {{ fullname }}
4 |
5 | {% block attributes %}
6 | {% if attributes %}
7 | .. rubric:: Module attributes
8 |
9 | .. autosummary::
10 | :toctree:
11 | {% for item in attributes %}
12 | {{ item }}
13 | {%- endfor %}
14 | {% endif %}
15 | {% endblock %}
16 |
17 | {% block functions %}
18 | {% if functions and fullname != 'bitorch_engine' %}
19 | .. rubric:: {{ _('Functions') }}
20 |
21 | .. autosummary::
22 | :toctree:
23 | :nosignatures:
24 | {% for item in functions %}
25 | {{ item }}
26 | {%- endfor %}
27 | {% endif %}
28 | {% endblock %}
29 |
30 | {% block classes %}
31 | {% if classes %}
32 | .. rubric:: {{ _('Classes') }}
33 |
34 | .. autosummary::
35 | :template: class.rst
36 | :toctree:
37 | {% for item in classes %}
38 | {{ item }}
39 | {%- endfor %}
40 | {% endif %}
41 | {% endblock %}
42 |
43 | {% block exceptions %}
44 | {% if exceptions %}
45 | .. rubric:: {{ _('Exceptions') }}
46 |
47 | .. autosummary::
48 | :toctree:
49 | {% for item in exceptions %}
50 | {{ item }}
51 | {%- endfor %}
52 | {% endif %}
53 | {% endblock %}
54 |
55 | {% block modules %}
56 | {% if modules %}
57 | .. rubric:: {{ _('Modules') }}
58 |
59 | .. autosummary::
60 | :toctree:
61 | :template: module.rst
62 | :recursive:
63 | {% for item in modules %}
64 | {% if not item.endswith('.extensions') %}{{ item }}{%- endif %}
65 | {%- endfor %}
66 | {% endif %}
67 | {% endblock %}
68 |
--------------------------------------------------------------------------------
/docs/source/build_options.rst:
--------------------------------------------------------------------------------
1 | Build options
2 | =============
3 |
4 | Building Specific Extensions
5 | ----------------------------
6 |
7 | While developing, a specific cpp/cuda extension can be (re-)build, by
8 | using the environment variable ``BIE_BUILD_ONLY``, like so:
9 |
10 | .. code:: bash
11 |
12 | BIE_BUILD_ONLY="bitorch_engine/layers/qlinear/binary/cpp" pip install -e . -v
13 |
14 | It needs to a relative path to one extension directory.
15 |
16 | Building for a Specific CUDA Architecture
17 | -----------------------------------------
18 |
19 | To build for a different CUDA Arch, use the environment variable
20 | ``BIE_CUDA_ARCH`` (e.g. use ‘sm_75’, ‘sm_80’, ‘sm_86’):
21 |
22 | .. code:: bash
23 |
24 | BIE_CUDA_ARCH="sm_86" pip install -e . -v
25 |
26 | Force Building CUDA Modules
27 | ---------------------------
28 |
29 | If you have CUDA development libraries installed, but
30 | ``torch.cuda.is_available()`` is False, e.g. in HPC or docker
31 | environments, you can still build the extensions that depend on CUDA, by
32 | setting ``BIE_FORCE_CUDA="true"``:
33 |
34 | .. code:: bash
35 |
36 | BIE_FORCE_CUDA="true" pip install -e . -v
37 |
38 | Skip Library File Building
39 | --------------------------
40 |
41 | If you just want to avoid rebuilding any files, you can set
42 | ``BIE_SKIP_BUILD``:
43 |
44 | .. code:: bash
45 |
46 | BIE_SKIP_BUILD="true" python3 -m build --no-isolation --wheel
47 |
48 | This would create a wheel and package ``.so`` files without trying to
49 | rebuild them.
50 |
51 |
--------------------------------------------------------------------------------
/docs/source/conf.py:
--------------------------------------------------------------------------------
1 | # Configuration file for the Sphinx documentation builder.
2 | #
3 | # For the full list of built-in configuration values, see the documentation:
4 | # https://www.sphinx-doc.org/en/master/usage/configuration.html
5 |
6 | import os
7 | import sys
8 | from pathlib import Path
9 |
10 | sys.path.insert(0, os.path.abspath('../../'))
11 |
12 | # -- Project information -----------------------------------------------------
13 | # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
14 |
15 | project = 'Bitorch Engine'
16 | copyright = '2024, Haojin Yang, Joseph Bethge, Nianhui Guo, Maximilian Schulze, Hong Guo, Paul Mattes'
17 | author = 'Haojin Yang, Joseph Bethge, Nianhui Guo, Maximilian Schulze, Hong Guo, Paul Mattes'
18 |
19 | root_path = Path(__file__).resolve().parent.parent.parent
20 | release = "unknown"
21 | version_file = root_path / "version.txt"
22 | if version_file.exists():
23 | with open(root_path / "version.txt") as handle:
24 | version_content = handle.read().strip()
25 | if version_content:
26 | release = version_content
27 |
28 |
29 | # -- General configuration ---------------------------------------------------
30 | # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
31 |
32 | extensions = [
33 | "sphinx.ext.autodoc",
34 | "sphinx.ext.autosummary",
35 | "sphinx.ext.doctest",
36 | "sphinx.ext.todo",
37 | "sphinx.ext.mathjax",
38 | "sphinx.ext.ifconfig",
39 | "sphinx.ext.viewcode",
40 | "sphinx.ext.githubpages",
41 | "sphinx.ext.napoleon",
42 | "autodocsumm",
43 | "sphinx.ext.intersphinx",
44 | "sphinx_toolbox.code",
45 | "sphinx_toolbox.collapse",
46 | "sphinx_design",
47 | ]
48 |
49 | # List of patterns, relative to source directory, that match files and
50 | # directories to ignore when looking for source files.
51 | # This pattern also affects html_static_path and html_extra_path.
52 | exclude_patterns = ['.so', '_build', 'Thumbs.db', '.DS_Store']
53 |
54 | # Add any paths that contain templates here, relative to this directory.
55 | templates_path = ['_templates']
56 | autosummary_generate = True
57 |
58 | # -- Options for HTML output -------------------------------------------------
59 | # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
60 |
61 | import sphinx_rtd_theme
62 | # html_theme = 'alabaster'
63 | html_theme = 'sphinx_rtd_theme'
64 | # html_static_path = ['_static']
65 |
--------------------------------------------------------------------------------
/docs/source/documentation.rst:
--------------------------------------------------------------------------------
1 | Full Documentation
2 | ==================
3 |
4 | .. autosummary::
5 | :toctree: _autosummary
6 | :template: module.rst
7 | :recursive:
8 |
9 | bitorch_engine
10 |
--------------------------------------------------------------------------------
/docs/source/index.rst:
--------------------------------------------------------------------------------
1 |
2 | Welcome to Bitorch Engine's documentation!
3 | ==========================================
4 |
5 | Welcome to the documentation of Bitorch Engine (BIE): a cutting-edge computation library for neural networks that enhances PyTorch by integrating specialized layers and functions for Low-Bit quantized neural network operations. This is where you can find all the information you need about how to use BIE.
6 |
7 | Building on the foundational strengths of Bitorch Engine, the technology has been employed in pioneering projects that push the boundaries of neural network training and inference. For instance,
8 |
9 | - `green-bit-llm-trainer `_: In this project, BIE represents a significant leap in the field of Large Language Model (LLM) fine-tuning. Unlike traditional approaches that either quantize a fully trained model or introduce a few additional trainable parameters for `LoRA `_ style fine-tuning, this project innovates by directly fine-tuning the quantized parameters of LLMs. This paradigm shift allows for the full-scale quantization fine-tuning of LLMs, ensuring that the training process tightly integrates with the quantization schema from the outset.
10 | - `green-bit-llm-inference `_ also showcases the BIE's adeptness at supporting inference for models quantized from 4 to 2-bits without any significant loss in accuracy compared to the original 32 or 16-bits models. It stands as a testament to BIE's capability to maintain the delicate balance between model size, computational efficiency, and accuracy, addressing one of the key challenges in deploying sophisticated neural networks in resource-constrained environments.
11 |
12 | All changes are tracked in the `changelog `_.
13 |
14 |
15 | .. toctree::
16 | :maxdepth: 4
17 | :caption: Contents:
18 |
19 | installation
20 | build_options
21 | documentation
22 |
23 | Indices and tables
24 | ==================
25 |
26 | * :ref:`genindex`
27 | * :ref:`modindex`
28 | * :ref:`search`
29 |
30 | Enjoy exploring our documentation!
31 |
--------------------------------------------------------------------------------
/examples/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/GreenBitAI/bitorch-engine/0009b1f8178e86df6433383879c0005a49c0cc92/examples/__init__.py
--------------------------------------------------------------------------------
/examples/mnist-lightning/mlp.py:
--------------------------------------------------------------------------------
1 |
2 | import torch.nn as nn
3 | import torch.nn.functional as F
4 |
5 |
6 | class MLP(nn.Module):
7 | def __init__(self, hidden_features=1024, num_layers=2, *args, **kwargs):
8 | super().__init__(*args, **kwargs)
9 | assert 2 <= num_layers <= 3, "MLP currently only supports 2 or 3 layers."
10 | self.num_layers = num_layers
11 | self.flatten = nn.Flatten()
12 |
13 | self.first_linear = nn.Linear(28*28, hidden_features)
14 | self.act1 = nn.PReLU()
15 | self.bn1 = nn.BatchNorm1d(num_features=hidden_features)
16 |
17 | self.linear1 = nn.Linear(hidden_features, hidden_features)
18 | self.act2 = nn.PReLU()
19 | self.bn2 = nn.BatchNorm1d(num_features=hidden_features)
20 |
21 | if num_layers > 2:
22 | self.linear2 = nn.Linear(hidden_features, hidden_features)
23 | self.act3 = nn.PReLU()
24 | self.bn3 = nn.BatchNorm1d(num_features=hidden_features)
25 |
26 | self.last_linear = nn.Linear(hidden_features, 10)
27 |
28 | def forward(self, x):
29 | x = self.flatten(x)
30 |
31 | x = self.first_linear(x)
32 | x = self.bn1(x)
33 | x = self.act1(x)
34 |
35 | x = self.linear1(x)
36 | x = self.bn2(x)
37 | x = self.act2(x)
38 |
39 | if self.num_layers > 2:
40 | x = self.linear2(x)
41 | x = self.bn3(x)
42 | x = self.act3(x)
43 |
44 | return self.last_linear(x)
45 |
46 |
47 | class SequentialMLP(nn.Module):
48 | def __init__(self, hidden_features = 2048, *args, **kwargs):
49 | super().__init__(*args, **kwargs)
50 | self.flatten = nn.Flatten()
51 | self.first_linear = nn.Linear(28*28, hidden_features)
52 | self.body = nn.Sequential(
53 | nn.BatchNorm1d(num_features=hidden_features),
54 | nn.ReLU(),
55 | nn.Linear(hidden_features, hidden_features),
56 | nn.BatchNorm1d(num_features=hidden_features),
57 | nn.ReLU(),
58 | nn.Linear(hidden_features, hidden_features),
59 | nn.BatchNorm1d(num_features=hidden_features),
60 | nn.ReLU(),
61 | )
62 | self.last_linear = nn.Linear(hidden_features, 10)
63 |
64 | def forward(self, x):
65 | x = self.flatten(x)
66 | x = self.first_linear(x)
67 | x = self.body(x)
68 | return self.last_linear(x)
69 |
--------------------------------------------------------------------------------
/examples/mnist-lightning/requirements.txt:
--------------------------------------------------------------------------------
1 | bitorch
2 | galore-torch
3 | lightning
4 |
--------------------------------------------------------------------------------
/examples/mnist/README.md:
--------------------------------------------------------------------------------
1 | # Example for MNIST
2 |
3 | In this example script we train a simple model for the MNIST dataset using [Bitorch Engine](https://github.com/GreenBitAI/bitorch-engine).
4 |
5 | First the requirements for this example need to be installed:
6 | ```bash
7 | pip install -r requirements.txt
8 | ```
9 |
10 | Then you can run the following to train an MLP with 3 layers (one of which is a binary layer),
11 | or add `--help` for more arguments:
12 | ```bash
13 | python train_mnist.py --epochs 10 --model q_mlp --log-interval 100
14 | ```
15 |
--------------------------------------------------------------------------------
/examples/mnist/__init__.py:
--------------------------------------------------------------------------------
1 | """
2 | This package contains an example for training an image classification model on the MNIST data set with BITorch
3 | and deploying it with the inference engine.
4 | """
5 |
--------------------------------------------------------------------------------
/examples/mnist/datasets/__init__.py:
--------------------------------------------------------------------------------
1 | """
2 | This submodule contains data preparation code for some of the datasets used with our models,
3 | i.e. MNIST, CIFAR 10 and 100 and ImageNet.
4 | """
5 |
6 | from typing import List, Type
7 |
8 | from .base import BasicDataset
9 | from .mnist import MNIST
10 |
11 | __all__ = [
12 | "BasicDataset",
13 | "dataset_from_name",
14 | "dataset_names",
15 | "MNIST",
16 | ]
17 |
18 |
19 | def dataset_from_name(name: str) -> Type[BasicDataset]:
20 | """returns the dataset to which the name belongs to (name has to be the value of the datasets
21 | name-attribute)
22 |
23 | Args:
24 | name (str): name of the dataset
25 |
26 | Raises:
27 | ValueError: raised if no dataset under that name was found
28 |
29 | Returns:
30 | dataset: the dataset
31 | """
32 | for dataset_class in [MNIST]:
33 | if dataset_class.name == name:
34 | return dataset_class
35 | raise Exception(f"unknown dataset: {name}")
36 |
37 |
38 | def dataset_names() -> List[str]:
39 | """getter for list of dataset names for argparse
40 |
41 | Returns:
42 | List: the dataset names
43 | """
44 | return [dataset_class.name for dataset_class in [MNIST]]
45 |
--------------------------------------------------------------------------------
/examples/mnist/datasets/base.py:
--------------------------------------------------------------------------------
1 | import logging
2 | import os
3 | from pathlib import Path
4 | from typing import Optional, Tuple, Any
5 |
6 | import torch
7 | from torch.utils.data import Dataset
8 | from torchvision.transforms import transforms
9 |
10 | from .dummy_dataset import DummyDataset
11 |
12 |
13 | class BasicDataset(Dataset):
14 | name = "None"
15 | num_classes = 0
16 | shape = (0, 0, 0, 0)
17 | mean: Any = None
18 | std_dev: Any = None
19 | num_train_samples = 0
20 | num_val_samples = 0
21 |
22 | def __init__(self, train: bool, root_directory: Optional[str] = None, download: bool = False) -> None:
23 | """initializes the dataset.
24 |
25 | Args:
26 | train (bool): whether the train or test dataset is wanted
27 | root_directory (str): path to main dataset storage directory
28 | download (bool): whether train/test should be downloaded if it does not exist
29 |
30 | Returns:
31 | Dataset: the created test/train dataset
32 | """
33 | super(BasicDataset, self).__init__()
34 | self.is_train = train
35 | self._download = download
36 | self.root_directory = self.get_dataset_root_directory(root_directory)
37 | self.dataset = self.get_dataset(download)
38 |
39 | @classmethod
40 | def get_train_and_test(cls, root_directory: str, download: bool = False) -> Tuple["BasicDataset", "BasicDataset"]:
41 | """creates a pair of train and test dataset.
42 |
43 | Returns:
44 | Tuple: the train and test dataset
45 | """
46 | return cls(True, root_directory, download), cls(False, root_directory, download)
47 |
48 | @classmethod
49 | def get_dummy_train_and_test_datasets(cls) -> Tuple[DummyDataset, DummyDataset]:
50 | train_set = DummyDataset(cls.shape, cls.num_classes, cls.num_train_samples) # type: ignore
51 | val_set = DummyDataset(cls.shape, cls.num_classes, cls.num_val_samples) # type: ignore
52 | return train_set, val_set
53 |
54 | def get_dataset_root_directory(self, root_directory_argument: Optional[str]) -> Path:
55 | """chooses the dataset root directory based on the passed argument or environment variables.
56 |
57 | Returns:
58 | Tuple: the train and test dataset
59 | """
60 | if root_directory_argument is not None:
61 | return Path(root_directory_argument)
62 |
63 | environment_variable_name = f"{self.name.upper()}_HOME"
64 | if os.environ.get(environment_variable_name) is not None:
65 | return Path(os.environ.get(environment_variable_name)) # type: ignore
66 | if os.environ.get("BITORCH_DATA_HOME") is not None:
67 | return Path(os.environ.get("BITORCH_DATA_HOME")) / self.name # type: ignore
68 |
69 | environment_variable_hint = (
70 | f" To change this, set '{environment_variable_name}' or 'BITORCH_DATA_HOME' "
71 | f"(in the latter case, the data resides in the folder '{self.name}' in BITORCH_DATA_HOME)."
72 | f" Some datasets can be downloaded by adding the --download command line argument."
73 | )
74 | if self._download:
75 | logging.warning("Dataset is being downloaded to the directory './data'." + environment_variable_hint)
76 | return Path("./data")
77 | else:
78 | raise ValueError(f"Dataset {self.name} not found." + environment_variable_hint)
79 |
80 | def get_dataset(self, download: bool) -> Dataset:
81 | """creates the actual dataset
82 |
83 | Args:
84 | download (bool): toggles if train/test shall be downloaded if possible
85 |
86 | Raises:
87 | NotImplementedError: thrown, because this method needs to be overwritten by subclasses
88 |
89 | Returns:
90 | Dataset: the created test/train dataset
91 | """
92 | raise NotImplementedError()
93 |
94 | def get_transform(self) -> Any:
95 | if self.is_train:
96 | return self.train_transform()
97 | return self.test_transform()
98 |
99 | @classmethod
100 | def test_transform(cls) -> Any:
101 | """get the transform for the test data.
102 |
103 | Returns:
104 | transform: the transform pipeline
105 | """
106 | return transforms.Compose([transforms.ToTensor(), cls.get_normalize_transform()])
107 |
108 | @classmethod
109 | def train_transform(cls) -> Any:
110 | """get the transform for the training data.
111 |
112 | Returns:
113 | transform: the transform pipeline
114 | """
115 | return transforms.Compose([transforms.ToTensor(), cls.get_normalize_transform()])
116 |
117 | @classmethod
118 | def get_normalize_transform(cls) -> transforms.Normalize:
119 | return transforms.Normalize(cls.mean, cls.std_dev)
120 |
121 | def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]: # type: ignore
122 | """returns the item at the given index of the dataset.
123 |
124 | Args:
125 | index (int): requested index
126 |
127 | Returns:
128 | Tuple[torch.Tensor, torch.Tensor]: data and label at the specified index
129 | """
130 | return self.dataset[index]
131 |
132 | def __len__(self) -> int:
133 | return len(self.dataset) # type: ignore
134 |
135 | def num_samples(self) -> int:
136 | """returns the (theoretical) dataset size."""
137 | return self.num_train_samples if self.is_train else self.num_val_samples
138 |
--------------------------------------------------------------------------------
/examples/mnist/datasets/dummy_dataset.py:
--------------------------------------------------------------------------------
1 | from torch.utils.data import Dataset
2 | import torch
3 | from typing import Tuple
4 |
5 |
6 | class DummyDataset(Dataset):
7 | """An iterator that produces repeated dummy data.
8 | Args:
9 | data_sample: a data sample that should be produced at each step.
10 | batch_size: the batch size for storing.
11 | sample_count: number of `data` samples in the dummy dataset.
12 | """
13 |
14 | def __init__(self, data_shape: torch.Size, num_classes: int, sample_count: int) -> None:
15 | self._data_sample = torch.zeros(data_shape)
16 | self._class_sample = torch.zeros((num_classes,), dtype=torch.int64)
17 | self._sample_count = sample_count
18 |
19 | def __len__(self) -> int:
20 | return self._sample_count
21 |
22 | def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
23 | return self._data_sample, self._class_sample
24 |
--------------------------------------------------------------------------------
/examples/mnist/datasets/mnist.py:
--------------------------------------------------------------------------------
1 | from torch.utils.data import Dataset
2 | from torchvision.datasets import mnist
3 |
4 | from .base import BasicDataset
5 |
6 |
7 | class MNIST(BasicDataset):
8 | name = "mnist"
9 | num_classes = 10
10 | shape = (1, 1, 28, 28)
11 |
12 | mean = (0.1307,)
13 | std_dev = (0.3081,)
14 | num_train_samples = 60000
15 | num_val_samples = 10000
16 |
17 | def get_dataset(self, download: bool = True) -> Dataset:
18 | return mnist.MNIST(
19 | root=self.root_directory,
20 | train=self.is_train,
21 | transform=self.get_transform(),
22 | download=download,
23 | )
24 |
--------------------------------------------------------------------------------
/examples/mnist/requirements.txt:
--------------------------------------------------------------------------------
1 | bitorch
2 |
--------------------------------------------------------------------------------
/licenses/LICENSE.cutlass.txt:
--------------------------------------------------------------------------------
1 | Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2 | SPDX-License-Identifier: BSD-3-Clause
3 |
4 | Redistribution and use in source and binary forms, with or without
5 | modification, are permitted provided that the following conditions are met:
6 |
7 | 1. Redistributions of source code must retain the above copyright notice, this
8 | list of conditions and the following disclaimer.
9 |
10 | 2. Redistributions in binary form must reproduce the above copyright notice,
11 | this list of conditions and the following disclaimer in the documentation
12 | and/or other materials provided with the distribution.
13 |
14 | 3. Neither the name of the copyright holder nor the names of its
15 | contributors may be used to endorse or promote products derived from
16 | this software without specific prior written permission.
17 |
18 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
19 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
20 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
21 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
22 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
23 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
24 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
25 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
26 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
27 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
28 |
--------------------------------------------------------------------------------
/licenses/LICENSE.exllamav2.txt:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Permission is hereby granted, free of charge, to any person obtaining a copy
4 | of this software and associated documentation files (the "Software"), to deal
5 | in the Software without restriction, including without limitation the rights
6 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7 | copies of the Software, and to permit persons to whom the Software is
8 | furnished to do so, subject to the following conditions:
9 |
10 | The above copyright notice and this permission notice shall be included in all
11 | copies or substantial portions of the Software.
12 |
13 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
19 | SOFTWARE.
--------------------------------------------------------------------------------
/licenses/LICENSE.mlx.txt:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright © 2023 Apple Inc.
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software without restriction, including without limitation the rights
8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 |
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 |
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 |
--------------------------------------------------------------------------------
/licenses/LICENSE.pytorch.txt:
--------------------------------------------------------------------------------
1 | From PyTorch:
2 |
3 | Copyright (c) 2016- Facebook, Inc (Adam Paszke)
4 | Copyright (c) 2014- Facebook, Inc (Soumith Chintala)
5 | Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
6 | Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu)
7 | Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
8 | Copyright (c) 2011-2013 NYU (Clement Farabet)
9 | Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
10 | Copyright (c) 2006 Idiap Research Institute (Samy Bengio)
11 | Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
12 |
13 | From Caffe2:
14 |
15 | Copyright (c) 2016-present, Facebook Inc. All rights reserved.
16 |
17 | All contributions by Facebook:
18 | Copyright (c) 2016 Facebook Inc.
19 |
20 | All contributions by Google:
21 | Copyright (c) 2015 Google Inc.
22 | All rights reserved.
23 |
24 | All contributions by Yangqing Jia:
25 | Copyright (c) 2015 Yangqing Jia
26 | All rights reserved.
27 |
28 | All contributions by Kakao Brain:
29 | Copyright 2019-2020 Kakao Brain
30 |
31 | All contributions by Cruise LLC:
32 | Copyright (c) 2022 Cruise LLC.
33 | All rights reserved.
34 |
35 | All contributions from Caffe:
36 | Copyright(c) 2013, 2014, 2015, the respective contributors
37 | All rights reserved.
38 |
39 | All other contributions:
40 | Copyright(c) 2015, 2016 the respective contributors
41 | All rights reserved.
42 |
43 | Caffe2 uses a copyright model similar to Caffe: each contributor holds
44 | copyright over their contributions to Caffe2. The project versioning records
45 | all such contribution and copyright details. If a contributor wants to further
46 | mark their specific copyright on a particular contribution, they should
47 | indicate their copyright solely in the commit message of the change when it is
48 | committed.
49 |
50 | All rights reserved.
51 |
52 | Redistribution and use in source and binary forms, with or without
53 | modification, are permitted provided that the following conditions are met:
54 |
55 | 1. Redistributions of source code must retain the above copyright
56 | notice, this list of conditions and the following disclaimer.
57 |
58 | 2. Redistributions in binary form must reproduce the above copyright
59 | notice, this list of conditions and the following disclaimer in the
60 | documentation and/or other materials provided with the distribution.
61 |
62 | 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America
63 | and IDIAP Research Institute nor the names of its contributors may be
64 | used to endorse or promote products derived from this software without
65 | specific prior written permission.
66 |
67 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
68 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
69 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
70 | ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
71 | LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
72 | CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
73 | SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
74 | INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
75 | CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
76 | ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
77 | POSSIBILITY OF SUCH DAMAGE.
78 |
--------------------------------------------------------------------------------
/licenses/LICENSE.tcbnn.txt:
--------------------------------------------------------------------------------
1 | TCBNN:
2 | Accelerating Binarized Neural Networks via Bit-Tensor-Cores in Turing GPUs
3 |
4 | 06/30/2020 by Ang Li from High-Performance-Computing Group,
5 | ACMD, PCSD, Pacific Northwest National Laboratory (PNNL),
6 | Richland, WA, 99354, USA.
7 |
8 |
9 | Copyright © 2020, Battelle Memorial Institute
10 |
11 | 1.Battelle Memorial Institute (hereinafter Battelle) hereby grants permission
12 | to any person or entity lawfully obtaining a copy of this software and associated
13 | documentation files (hereinafter “the Software”) to redistribute and use the
14 | Software in source and binary forms, with or without modification. Such person
15 | or entity may use, copy, modify, merge, publish, distribute, sublicense, and/or
16 | sell copies of the Software, and may permit others to do so, subject to the
17 | following conditions:
18 |
19 | - Redistributions of source code must retain the above copyright notice, this list
20 | of conditions and the following disclaimers.
21 |
22 | - Redistributions in binary form must reproduce the above copyright notice, this list
23 | of conditions and the following disclaimer in the documentation and/or other materials
24 | provided with the distribution.
25 |
26 | - Other than as used herein, neither the name Battelle Memorial Institute or Battelle
27 | may be used in any form whatsoever without the express written consent of Battelle.
28 |
29 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY
30 | EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
31 | OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
32 | SHALL BATTELLE OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
33 | EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
34 | SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
35 | HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
36 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
37 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
38 |
39 |
--------------------------------------------------------------------------------
/requirements-dev.txt:
--------------------------------------------------------------------------------
1 | pytest
2 | numpy
3 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | bitorch
2 | torch
3 | py-cpuinfo>=9.0.0
4 | setuptools~=69.0
5 |
--------------------------------------------------------------------------------
/structure.md:
--------------------------------------------------------------------------------
1 |
2 | .
3 | ├── ...
4 | ├── layers # Low-bit layers
5 | │ ├── qconv # quantized Convolotional Layer
6 | │ ├── qembedding # currently only 1-bit embedding layer supported
7 | │ ├── ...
8 | │ └── qlinear
9 | │ ├── binary (1-bit)
10 | │ │ ├── cpp # x86 CPU
11 | │ │ ├── cuda # Nvidia GPU
12 | │ │ └── cutlass # Nvidia GPU
13 | │ └── n-bit (2/4/8-bit)
14 | │ ├── mps # Apple GPU
15 | │ ├── cuda # Nvidia GPU, e.g., weight-only quantized LLMs
16 | │ └── cutlass # Nvidia GPU, e.g., quantization aware training for both activation and weight
17 | └── optim
18 | │ └── DiodeMix # dedicated optimizer for low-bit quantized model
19 | ├── functions
20 | └── ...
--------------------------------------------------------------------------------
/tests/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/GreenBitAI/bitorch-engine/0009b1f8178e86df6433383879c0005a49c0cc92/tests/__init__.py
--------------------------------------------------------------------------------
/tests/functions/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/GreenBitAI/bitorch-engine/0009b1f8178e86df6433383879c0005a49c0cc92/tests/functions/__init__.py
--------------------------------------------------------------------------------
/tests/layers/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/GreenBitAI/bitorch-engine/0009b1f8178e86df6433383879c0005a49c0cc92/tests/layers/__init__.py
--------------------------------------------------------------------------------
/tests/layers/test_custom_binary_linear.py:
--------------------------------------------------------------------------------
1 | import pytest
2 | import torch
3 | from torch import nn
4 | from torch.autograd import Function
5 | from torch.nn import Parameter
6 |
7 |
8 | class LinearFunction(Function):
9 | @staticmethod
10 | # ctx is the first argument to forward
11 | def forward(ctx, input, weight, bias=None):
12 | # The forward pass can use ctx.
13 | ctx.save_for_backward(input, weight, bias)
14 | # TODO: instead of converting to float and use torch's mm we should use binary mm instead
15 | output = input.float().mm(weight.t().float())
16 | if bias is not None:
17 | output += bias.unsqueeze(0).expand_as(output)
18 | return output
19 |
20 | @staticmethod
21 | def backward(ctx, grad_output):
22 | input, weight, bias = ctx.saved_tensors
23 | grad_input = grad_weight = grad_bias = None
24 |
25 | if ctx.needs_input_grad[0]:
26 | # TODO: instead of converting to float and use torch's mm we should use our binary mm instead
27 | grad_input = grad_output.mm(weight.float())
28 | if ctx.needs_input_grad[1]:
29 | # TODO: instead of converting to float and use torch's mm we should use our binary mm instead
30 | grad_weight = grad_output.t().mm(input.float())
31 | if bias is not None and ctx.needs_input_grad[2]:
32 | grad_bias = grad_output.sum(0)
33 |
34 | print("Weight grads calculated: ", grad_weight)
35 |
36 | # manually need to convert to our required formulation
37 | def binarize_grad(x):
38 | # 1.0 = True
39 | # 0.0 = False
40 | return torch.where(x >= 0.0, 1.0, 0.0)
41 |
42 | # TODO: we could inline this later
43 | grad_input = binarize_grad(grad_input)
44 | grad_weight = binarize_grad(grad_weight)
45 |
46 | return grad_input, grad_weight, grad_bias
47 |
48 |
49 | # Option 2: wrap in a function, to support default args and keyword args.
50 | def linear(input, weight, bias=None):
51 | return LinearFunction.apply(input, weight, bias)
52 |
53 |
54 | class TLinear(nn.Module):
55 | def __init__(self, input_features, output_features, bias=True):
56 | super().__init__()
57 | self.input_features = input_features
58 | self.output_features = output_features
59 |
60 | # nn.Parameter is a special kind of Tensor, that will get
61 | # automatically registered as Module's parameter once it's assigned
62 | # as an attribute. Parameters and buffers need to be registered, or
63 | # they won't appear in .parameters() (doesn't apply to buffers), and
64 | # won't be converted when e.g. .cuda() is called. You can use
65 | # .register_buffer() to register buffers.
66 | # nn.Parameters require gradients by default.
67 | self.weight = nn.Parameter(torch.empty(output_features, input_features))
68 | if bias:
69 | self.bias = nn.Parameter(torch.empty(output_features))
70 | else:
71 | # You should always register all possible parameters, but the
72 | # optional ones can be None if you want.
73 | self.register_parameter('bias', None)
74 |
75 | # Not a very smart way to initialize weights
76 | nn.init.uniform_(self.weight, -0.1, 0.1)
77 | if self.bias is not None:
78 | nn.init.uniform_(self.bias, -0.1, 0.1)
79 |
80 | def forward(self, input):
81 | # See the autograd section for explanation of what happens here.
82 | return LinearFunction.apply(input, self.weight, self.bias)
83 |
84 | def extra_repr(self):
85 | # (Optional)Set the extra information about this module. You can test
86 | # it by printing an object of this class.
87 | return 'input_features={}, output_features={}, bias={}'.format(
88 | self.input_features, self.output_features, self.bias is not None
89 | )
90 |
91 |
92 | def test_q_linear_with_binary_weights():
93 | torch.manual_seed(42)
94 | # import pydevd_pycharm
95 | # pydevd_pycharm.settrace('localhost', port=11004, stdoutToServer=True, stderrToServer=True)
96 |
97 | batch_size = 10
98 | num_inputs = 32
99 | num_outputs = 64
100 |
101 | # option 1: set dtype
102 | # currently not possible, check uniform bounds fails:
103 | # layer = TstLinear(num_inputs, num_outputs, bias=False, dtype=torch.bool)
104 |
105 | # option 2: manually replace weight:
106 | # currently possible
107 | # but we have to
108 | layer = TLinear(num_inputs, num_outputs, bias=False)
109 | layer.weight = Parameter(torch.rand((num_outputs, num_inputs)) > 0.5, requires_grad=True)
110 |
111 | input = Parameter(torch.rand((batch_size, num_inputs)) > 0.5, requires_grad=True)
112 |
113 | result = layer(input)
114 | print(result)
115 | print("Result shape: ", result.size())
116 |
117 | print("Grad before backward: ", layer.weight.grad)
118 | mse_loss = nn.MSELoss()
119 | dummy_loss = mse_loss(result, torch.ones_like(result) * 10)
120 | dummy_loss.backward()
121 | print("Grad after backward: ", layer.weight.grad)
122 | # technically it works, but we only get boolean gradients (all true currently)
123 | # could be fixed with custom backward pass?
124 |
--------------------------------------------------------------------------------
/tests/layers/test_nbit_conv.py:
--------------------------------------------------------------------------------
1 | import pytest
2 | import torch
3 | import numpy as np
4 | import time
5 | from bitorch.layers import QConv2d
6 | from tests.layers.util import get_cuda_test_device
7 | from bitorch_engine.layers.qconv.nbit.cutlass import Q4Conv2dCutlass
8 |
9 | """
10 | Test nbit inference layers
11 | """
12 |
13 | # lower print threshold
14 | torch.set_printoptions(threshold=100)
15 |
16 | def to_device(data: torch.Tensor, device: torch.device) -> torch.Tensor:
17 | if isinstance(data, (list, tuple)):
18 | return [to_device(x, device) for x in data]
19 | return data.to(device)
20 |
21 | TEST_BATCH_SIZE = [1, 32, 64, 128]
22 | # Input shape: (batch size, num of input channels, h, w)
23 | TEST_INPUT_DATA = [
24 | ((64, 64, 64), [64, 32],
25 | {"kernel_size": 3, "padding": 2, "stride": 2, "dilation": 1}),
26 | ((64, 56, 56), [64, 256],
27 | {"kernel_size": 1, "padding": 0, "stride": 1, "dilation": 1}),
28 | ((64, 56, 56), [64, 64],
29 | {"kernel_size": 3, "padding": 0, "stride": 1, "dilation": 1}),
30 | ((256, 56, 56), [256, 64],
31 | {"kernel_size": 1, "padding": 0, "stride": 1, "dilation": 1}),
32 | ((256, 56, 56), [256, 128],
33 | {"kernel_size": 1, "padding": 0, "stride": 1, "dilation": 1}),
34 | ((128, 28, 28), [128, 128],
35 | {"kernel_size": 3, "padding": 0, "stride": 1, "dilation": 1}),
36 | ((128, 28, 28), [128, 512],
37 | {"kernel_size": 1, "padding": 0, "stride": 1, "dilation": 1}),
38 | ((256, 14, 14), [256, 256],
39 | {"kernel_size": 3, "padding": 2, "stride": 1, "dilation": 1}),
40 | ((512, 7, 7), [512, 512],
41 | {"kernel_size": 3, "padding": 0, "stride": 1, "dilation": 1}),
42 | ]
43 |
44 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available, skipping CUDA-related test")
45 | @pytest.mark.parametrize("input_shape, args, kwargs", TEST_INPUT_DATA)
46 | @pytest.mark.parametrize("BS", TEST_BATCH_SIZE)
47 | def test_q4_conv_cuda(input_shape, args, kwargs, BS):
48 | print("\n")
49 | print("batch_size: {}, input_shape:{}, input/output nums:{}, other args:{}"
50 | .format(BS, input_shape, args, kwargs))
51 | input_shape = (BS, input_shape[0], input_shape[1], input_shape[2])
52 | input = np.random.uniform(-1, 1, input_shape)
53 | input_tensor = torch.tensor(input).float()
54 |
55 | num_runs = 10
56 |
57 | # to gpu
58 | device = get_cuda_test_device()
59 | torch.cuda.set_device(device)
60 | input_tensor_cuda = to_device(input_tensor, device)
61 |
62 | layer = QConv2d(*args, **kwargs)
63 | layer.to(device)
64 | start_time = time.time()
65 | for i in range(num_runs):
66 | result_bitorch = layer(input_tensor_cuda)
67 | time_engine = time.time() - start_time
68 | print("bitorch binary-conv: %.6f s" % (time_engine / num_runs))
69 |
70 | padding = kwargs["padding"]
71 | kernel_size = kwargs["kernel_size"]
72 | stride = kwargs["stride"]
73 | dilation = kwargs["dilation"]
74 | nbit_conv_layer = Q4Conv2dCutlass(in_channels=int(args[0]),
75 | out_channels=args[1],
76 | kernel_size=kernel_size,
77 | stride=stride,
78 | padding=padding,
79 | dilation=dilation,
80 | device=device)
81 | nbit_conv_layer.to(device)
82 | nbit_conv_layer.prepare_params()
83 | result = nbit_conv_layer(input_tensor_cuda)
84 | # print(result)
85 |
86 | grad_input_data = np.random.uniform(-1, 1, result.shape)
87 | grad_input_tensor = torch.tensor(grad_input_data).float()
88 | grad_input_tensor_cuda = to_device(grad_input_tensor, device)
89 |
90 | # use quantized weight for inference
91 | nbit_conv_layer.generate_quantized_weight(qweight_only=True)
92 | nbit_conv_layer.eval()
93 | start_time = time.time()
94 | for i in range(num_runs):
95 | result_quantized_w = nbit_conv_layer(input_tensor_cuda)
96 | time_engine = time.time() - start_time
97 | print("engine q4-conv forward (CUTLASS): %.6f s" % (time_engine / num_runs))
98 | # print(result_quantized_w)
99 | assert torch.equal(result, result_quantized_w)
100 |
101 | start_time = time.time()
102 | for i in range(num_runs):
103 | result.backward(grad_input_tensor_cuda, retain_graph=True)
104 | torch.cuda.synchronize()
105 | time_engine = time.time() - start_time
106 | print("bitorch-engine q4-conv backward (CUTLASS): %.6f s" % (time_engine/num_runs))
107 |
--------------------------------------------------------------------------------
/tests/layers/test_nbit_linear_mixbits.py:
--------------------------------------------------------------------------------
1 | import time, math
2 | import pytest
3 | import torch
4 | import json
5 |
6 | from bitorch_engine.layers.qlinear.nbit.cuda import MBWQLinearCuda
7 | from tests.layers.util import get_cuda_test_device, get_packed_info, get_q_groups
8 |
9 |
10 | """
11 | Test nbit inference layers
12 | """
13 |
14 | # lower print threshold
15 | torch.set_printoptions(threshold=100)
16 |
17 |
18 | def to_device(data: torch.Tensor, device: torch.device) -> torch.Tensor:
19 | if isinstance(data, (list, tuple)):
20 | return [to_device(x, device) for x in data]
21 | return data.to(device)
22 |
23 |
24 | # ========= testing data ========== #
25 |
26 | test_json_string = '{' \
27 | '"q_proj": { "group_size": { "4": 32, "2": 32 }, "bits": [ 4, 2 ], "bits_prop": [ 0.75, 0.25 ], "scale_bits": 4 }, ' \
28 | '"k_proj": { "group_size": { "4": 32, "2": 32 }, "bits": [ 4, 2 ], "bits_prop": [ 0.25, 0.75 ], "scale_bits": 4 } ' \
29 | '}'
30 |
31 |
32 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available, skipping CUDA-related test")
33 | @pytest.mark.parametrize("num_input_features", [128, 4096])
34 | @pytest.mark.parametrize("num_hidden_fc", [128, 4096])
35 | @pytest.mark.parametrize("batch_size", [1, 2])
36 | def test_mbwq_linear_exl2_cuda(num_input_features, num_hidden_fc, batch_size):
37 |
38 | # support 4-bit und bf16 only!
39 | dtype = torch.half
40 | num_runs = 1
41 | device = get_cuda_test_device()
42 | torch.cuda.set_device(device)
43 |
44 | # creating test data
45 | input_data = torch.normal(0, 1, size=(batch_size, num_input_features), requires_grad=False, dtype=dtype)
46 | # to gpu
47 | input_data_cuda = to_device(input_data, device)
48 |
49 | # Parsing the JSON string into a Python dictionary
50 | gbe_strategy = json.loads(test_json_string)
51 |
52 | for key, value in gbe_strategy.items():
53 | # read attribute information from predefined json config
54 | groups, rows = get_packed_info(num_input_features, value["bits"], value["bits_prop"], value["group_size"])
55 |
56 | print("\nM:{}, N:{}, K:{}, bits:{}, group_size:{}, bits_prop:{}, groups:{}, packed_rows:{}."
57 | .format(batch_size, num_hidden_fc, num_input_features, str(value["bits"]), str(value["group_size"]), str(value["bits_prop"]),
58 | groups, rows))
59 |
60 | # creating int weights
61 | random_int4_tensor = torch.randint(0, 2 ** 16 - 1, size=(rows, num_hidden_fc))
62 | int_weight_cuda = to_device(random_int4_tensor, device)
63 |
64 | mbwq_linear_layer = MBWQLinearCuda(in_channels=num_input_features, out_channels=num_hidden_fc,
65 | w_bit=4, dtype=dtype, group_size=32,
66 | dq_group_size=1, use_gba_quant=True,
67 | asym=False, dq_mode=2, use_mbw=True,
68 | groups=groups, rows_packed=rows)
69 |
70 | mbwq_linear_layer.set_qweight_data(int_weight_cuda)
71 | # random scales and zeros
72 | scales = torch.randn_like(mbwq_linear_layer.scales).half()
73 | zeros = torch.randn_like(mbwq_linear_layer.zeros).half()
74 | mbwq_linear_layer.set_scales(scales)
75 | mbwq_linear_layer.set_zeros(zeros)
76 | mbwq_linear_layer.q_perm = torch.tensor([i for i in range(num_input_features)], dtype=torch.short)
77 |
78 | # get q_groups
79 | q_groups = get_q_groups(groups, value["bits"], value["group_size"], num_input_features, value["bits_prop"])
80 |
81 | assert mbwq_linear_layer.q_groups.numel() == len(q_groups)
82 |
83 | mbwq_linear_layer.q_groups = torch.Tensor(q_groups).to(torch.short)
84 |
85 | # to device
86 | mbwq_linear_layer.to(device)
87 |
88 | # will perform qweight layout transformation
89 | mbwq_linear_layer.prepare_params()
90 |
91 | # Testing fp weight reconstruction
92 | reconstructed_fp_weights = MBWQLinearCuda.exl2fp_weight(mbwq_linear_layer.qweight, mbwq_linear_layer.scales,
93 | mbwq_linear_layer.zeros, mbwq_linear_layer.q_perm,
94 | mbwq_linear_layer.q_group_map, mbwq_linear_layer.rows).to(dtype)
95 |
96 | # pytorch result:
97 | result_pt = torch.matmul(input_data_cuda.mul(mbwq_linear_layer.channel_scale), reconstructed_fp_weights)
98 |
99 | # Testing inference output
100 | start_time = time.time()
101 | for i in range(num_runs):
102 | result = mbwq_linear_layer(input_data_cuda)
103 | torch.cuda.synchronize()
104 | time_engine = time.time() - start_time
105 |
106 | print("bitorch-engine mbwq_linear (CUDA) run time: %.6f s" % (time_engine/num_runs))
107 |
108 | assert torch.all(torch.isclose(result, result_pt, rtol=2, atol=2, equal_nan=False))
--------------------------------------------------------------------------------
/tests/layers/test_nbit_linear_mps.py:
--------------------------------------------------------------------------------
1 | import time
2 |
3 | import pytest
4 | import torch
5 | import numpy as np
6 |
7 | from bitorch_engine.layers.qlinear.nbit.mps import MPQLinearMlx
8 | from tests.layers.util import get_mps_test_device
9 |
10 | """
11 | Test nbit inference layers
12 | """
13 |
14 | # lower print threshold
15 | torch.set_printoptions(threshold=100)
16 |
17 |
18 | def to_device(data: torch.Tensor, device: torch.device) -> torch.Tensor:
19 | if isinstance(data, (list, tuple)):
20 | return [to_device(x, device) for x in data]
21 | return data.to(device)
22 |
23 |
24 | @pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available, skipping MPS-related test")
25 | @pytest.mark.parametrize("w_bit", [2, 4, 8])
26 | @pytest.mark.parametrize("dtype", [torch.half])
27 | @pytest.mark.parametrize("num_input_features", [1024, 8192])
28 | @pytest.mark.parametrize("num_hidden_fc", [1024, 8192])
29 | @pytest.mark.parametrize("batch_size", [64])
30 | @pytest.mark.parametrize("group_size", [32, 128])
31 | @pytest.mark.parametrize("llama_v", [2])
32 | @pytest.mark.parametrize("model_size", ["1.1b", "7b", "30b", "70b"])
33 | def test_mpq_linear_mps(w_bit, dtype, num_input_features, num_hidden_fc, batch_size, group_size, llama_v, model_size):
34 | import mlx.core
35 |
36 | if ((w_bit == 1 or w_bit == 8) and group_size == 32) \
37 | or (model_size in ["30b", "70b"] and w_bit != 2) \
38 | or (model_size in ["1.1b"] and w_bit != 2):
39 | pytest.skip()
40 |
41 | # === following configuration from low_bit_llama https://github.com/GreenBitAI/low_bit_llama === #
42 | double_groupsize = -1
43 | if group_size == 32 and model_size not in ["1.1b", "1.1B"]:
44 | asym = True
45 | else:
46 | asym = False
47 |
48 | if w_bit == 2:
49 | if asym:
50 | double_groupsize = -1
51 | else:
52 | if group_size == 32:
53 | double_groupsize = 32
54 | else:
55 | if llama_v == 1 and model_size not in ["30b", "30B"]:
56 | double_groupsize = 64
57 | else:
58 | double_groupsize = 32
59 | else:
60 | if model_size in ["3b", "3B"]:
61 | double_groupsize = 64
62 | elif model_size in ["7b", "7B"]:
63 | double_groupsize = 256
64 |
65 | v1 = (llama_v == 1) and model_size in ["7b", "7B"]
66 |
67 | if w_bit not in [2, 4, 8]:
68 | print("w_bit not supported")
69 | pytest.skip()
70 | if group_size not in [32, 64, 128]:
71 | print("group_size not supported")
72 | pytest.skip()
73 | if asym:
74 | print("asym not supported")
75 | pytest.skip()
76 | # =============================================================================================== #
77 |
78 | print("\nM:{}, N:{}, K:{}, bits:{}, dtype:{}, groupsize:{}, llama_v:{}, model_size:{}."
79 | .format(batch_size, num_hidden_fc, num_input_features, w_bit, dtype, group_size, llama_v, model_size))
80 |
81 | num_runs = 10
82 | # creating test data
83 | input_data = torch.normal(0, 1, size=(batch_size, num_input_features), requires_grad=False, dtype=dtype)
84 | grad_input_data = torch.normal(0, 1, size=(batch_size, num_hidden_fc), requires_grad=False, dtype=dtype)
85 |
86 | # to gpu
87 | device = get_mps_test_device()
88 | input_data_mps = to_device(input_data, device)
89 | grad_input_data_mps = to_device(grad_input_data, device)
90 |
91 | time_engine = 0
92 | for i in range(num_runs):
93 | b_linear = torch.nn.Linear(num_input_features, num_hidden_fc, bias=False, dtype=dtype)
94 | b_linear.to(device)
95 | start_time = time.time()
96 | result = b_linear(input_data_mps)
97 | time_engine += time.time() - start_time
98 | print(f"pytorch linear forward run time (device {device}): {time_engine / num_runs:.6f} s")
99 |
100 | mpq_linear_layer = MPQLinearMlx(in_channels=num_input_features,
101 | out_channels=num_hidden_fc,
102 | w_bit=w_bit,
103 | dtype=dtype,
104 | group_size=group_size,
105 | dq_group_size=double_groupsize,
106 | dq_mode=1 if v1 else 2,
107 | use_gba_quant=True,
108 | asym=asym,
109 | requires_grad=False)
110 |
111 | mpq_linear_layer.to(device)
112 | mpq_linear_layer.prepare_params()
113 |
114 | time_engine = 0
115 | time_mlx = 0
116 | for i in range(num_runs):
117 | random_int_tensor = torch.randint(1, 1000, size=mpq_linear_layer.qweight.shape,
118 | dtype=mpq_linear_layer.qweight.dtype, device=mpq_linear_layer.qweight.device)
119 | mpq_linear_layer.set_qweight_data(random_int_tensor)
120 | start_time = time.time()
121 | result1 = mpq_linear_layer(input_data_mps)
122 | time_engine += time.time() - start_time
123 |
124 | qweight_mlx = mlx.core.array(mpq_linear_layer.qweight.numpy().astype(np.uint32))
125 | scales_mlx = mlx.core.array(mpq_linear_layer.scales.numpy())
126 | zeros_mlx = mlx.core.array(mpq_linear_layer.zeros.numpy())
127 |
128 | start_time = time.time()
129 | input_mlx = mlx.core.array(input_data_mps.cpu().numpy())
130 | mlx_matmul = mlx.core.quantized_matmul(
131 | input_mlx,
132 | qweight_mlx,
133 | scales=scales_mlx,
134 | biases=zeros_mlx,
135 | transpose=True,
136 | group_size=group_size,
137 | bits=w_bit)
138 | mlx_matmul = torch.from_numpy(np.array(mlx_matmul)).to('mps')
139 | time_mlx += time.time() - start_time
140 | assert mlx_matmul.equal(result1), "mps matmul failed"
141 | print(f"bitorch-engine mpq linear forward (MPS) run time: {time_engine / num_runs:.6f} s")
142 | print(f"Mlx (MPS) run time: {time_mlx / num_runs:.6f} s")
--------------------------------------------------------------------------------
/tests/layers/util.py:
--------------------------------------------------------------------------------
1 | import os
2 | import math
3 | from typing import Union, List
4 |
5 | import torch
6 |
7 |
8 | def activate_remote_pycharm_debug(port: int = 11004):
9 | import pydevd_pycharm
10 | pydevd_pycharm.settrace('localhost', port=port, stdoutToServer=True, stderrToServer=True)
11 |
12 |
13 | def to_device(data: torch.Tensor, device: torch.device) -> Union[torch.Tensor, List[torch.Tensor]]:
14 | if isinstance(data, (list, tuple)):
15 | return [to_device(x, device) for x in data]
16 | return data.to(device)
17 |
18 |
19 | def get_cuda_test_device_id():
20 | return int(os.environ.get("BIE_DEVICE", "0"))
21 |
22 |
23 | def get_cuda_test_device():
24 | return torch.device(f"cuda:{get_cuda_test_device_id()}")
25 |
26 |
27 | def get_mps_test_device():
28 | return torch.device("mps")
29 |
30 |
31 | def get_packed_info(channels, n_bits, bits_prop, bits_group_size):
32 | groups = 0
33 | rows = 0
34 | bits_channel = []
35 | for idx in range(len(bits_prop)):
36 | if idx < len(bits_prop) - 1:
37 | minimal_channels = list(bits_group_size.values())[idx]
38 | channel_pre_pack = max(1, int(channels * (bits_prop[idx])) // minimal_channels) * minimal_channels
39 | bits_channel.append(channel_pre_pack)
40 | groups += channel_pre_pack // minimal_channels
41 | rows += channel_pre_pack // 32 * n_bits[idx]
42 | else:
43 | minimal_channels = list(bits_group_size.values())[idx]
44 | channel_pre_pack = channels - sum(bits_channel)
45 | bits_channel.append(channel_pre_pack)
46 | groups += channel_pre_pack // minimal_channels
47 | rows += channel_pre_pack // 32 * n_bits[idx]
48 |
49 | return groups, rows
50 |
51 |
52 | def get_q_groups(groups, n_bits, group_size, channels, bits_prop):
53 | qgroups = []
54 | bits_column_end_index = []
55 |
56 | for idx in range(len(bits_prop)):
57 | if idx < len(bits_prop) - 1:
58 | minimal_columns = list(group_size.values())[idx]
59 | columns_index = max(1, int(channels * (
60 | bits_prop[idx])) // minimal_columns) * minimal_columns # TODO: determine the minimal bits columns
61 | if idx > 0:
62 | columns_index += bits_column_end_index[-1]
63 | bits_column_end_index.append(columns_index)
64 | else:
65 | bits_column_end_index.append(channels)
66 |
67 | for bits_idx, bits in enumerate(n_bits):
68 | if bits_idx == 0:
69 | rows_per_bit = bits_column_end_index[bits_idx]
70 | else:
71 | rows_per_bit = bits_column_end_index[bits_idx] - bits_column_end_index[bits_idx - 1]
72 |
73 | gs = group_size[str(bits)]
74 | groups_per_bit = rows_per_bit // gs
75 |
76 | for group in range(groups_per_bit):
77 | qgroups.append(bits) # record bits per group
78 | qgroups.append(0)
79 |
80 | out_row = 0
81 | rem_rows = channels
82 | for i in range(groups):
83 | bits = qgroups[2 * i]
84 | gs = group_size[str(bits)]
85 |
86 | rows_per_group = min(gs, rem_rows) # rows per group before packing
87 | wpqr = 32 / bits # INT32 elements per group for packing
88 | qrows = math.ceil(rows_per_group / wpqr) # rows per group after packing
89 | qgroups[2 * i + 1] = out_row # record packed rows start idx per group
90 |
91 | out_row += qrows
92 |
93 | return qgroups
94 |
95 |
96 | def pack_rows_4_pytorch(input_tensor, rows, columns):
97 | # Calculate the number of output columns and store 4 columns per 32 bits, so use columns * 4 / 32
98 | out_columns = columns * 4 // 32
99 |
100 | # Initialize the output tensor, the size is [rows, out_columns], the data type is uint32
101 | output_tensor = torch.zeros((rows, out_columns), dtype=torch.int64, device=input_tensor.device)
102 |
103 | # To simulate the behavior of the CUDA kernel, we need to perform the same operation for each output element
104 | # This is implemented using the broadcast and indexing functions of PyTorch
105 | for row in range(rows):
106 | for out_column in range(out_columns):
107 | packed = 0
108 | for i in range(8):
109 | # Simulate operations in the CUDA kernel
110 | x = input_tensor[row, out_column * 8 + i].item() - 1
111 | packed |= (x << (i * 4))
112 | output_tensor[row, out_column] = packed
113 | # Use bitwise operators to extract the unsigned lower 32 bits
114 | # Note: 0xFFFFFFFF is a 32-bit all-1 mask, which is used to extract the lower 32 bits
115 | output_tensor = (output_tensor & 0xFFFFFFFF).to(torch.int32)
116 | return output_tensor
--------------------------------------------------------------------------------
/tests/util/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/GreenBitAI/bitorch-engine/0009b1f8178e86df6433383879c0005a49c0cc92/tests/util/__init__.py
--------------------------------------------------------------------------------
/tests/util/binary_mse_loss.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | def binary_mse_loss(output, target):
4 | print("Output:", output)
5 | print("Target:", target)
6 | loss = torch.mean((output.to(torch.float32) - target.to(torch.float32))**2)
7 | return loss
8 |
--------------------------------------------------------------------------------
/version.txt:
--------------------------------------------------------------------------------
1 | 0.2.6
2 |
--------------------------------------------------------------------------------