Compare commits
4 Commits
readme-dev
..
paper
| Author | SHA1 | Date | |
|---|---|---|---|
 Eddie-Wang1120
|
3dcfd14628 | ||
|
Eddie-Wang1120
|
0ab05d6f64 | ||
|
Eddie-Wang1120
|
61e37b5430 | ||
|
Eddie-Wang1120
|
4c736e3728 |
@@ -34,7 +34,6 @@ nppBackup
|
||||
|
||||
# Models
|
||||
models/*
|
||||
gpu/checkpoints/*
|
||||
|
||||
# Python
|
||||
|
||||
|
||||
+2
-2
@@ -1,4 +1,4 @@
|
||||
[submodule "3rdparty/llama.cpp"]
|
||||
path = 3rdparty/llama.cpp
|
||||
url = https://github.com/Eddie-Wang1120/llama.cpp.git
|
||||
branch = merge-dev
|
||||
url = git@github.com:Eddie-Wang1120/llama.cpp.git
|
||||
branch = pp
|
||||
|
||||
Vendored
+1
-1
Submodule 3rdparty/llama.cpp updated: 5eb47b7210...5095a95664
+5
-5
@@ -14,6 +14,7 @@ set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
|
||||
# option list
|
||||
option(BITNET_ARM_TL1 "bitnet.cpp: use tl1 on arm platform" OFF)
|
||||
option(BITNET_X86_TL2 "bitnet.cpp: use tl2 on x86 platform" OFF)
|
||||
option(BITNET_TL2_LOSS "bitnet.cpp: use tl2 on x86 platform" OFF)
|
||||
|
||||
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED true)
|
||||
@@ -24,6 +25,7 @@ set(THREADS_PREFER_PTHREAD_FLAG ON)
|
||||
# override ggml options
|
||||
set(GGML_BITNET_ARM_TL1 ${BITNET_ARM_TL1})
|
||||
set(GGML_BITNET_X86_TL2 ${BITNET_X86_TL2})
|
||||
set(GGML_BITNET_TL2_LOSS ${BITNET_TL2_LOSS})
|
||||
|
||||
if (GGML_BITNET_ARM_TL1)
|
||||
add_compile_definitions(GGML_BITNET_ARM_TL1)
|
||||
@@ -31,15 +33,13 @@ endif()
|
||||
if (GGML_BITNET_X86_TL2)
|
||||
add_compile_definitions(GGML_BITNET_X86_TL2)
|
||||
endif()
|
||||
|
||||
if (CMAKE_C_COMPILER_ID STREQUAL "GNU" OR CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
|
||||
add_compile_options(-fpermissive)
|
||||
if (GGML_BITNET_TL2_LOSS)
|
||||
add_compile_definitions(GGML_BITNET_TL2_LOSS)
|
||||
endif()
|
||||
|
||||
find_package(Threads REQUIRED)
|
||||
|
||||
add_subdirectory(src)
|
||||
set(LLAMA_BUILD_SERVER ON CACHE BOOL "Build llama.cpp server" FORCE)
|
||||
add_subdirectory(3rdparty/llama.cpp)
|
||||
|
||||
# install
|
||||
@@ -75,4 +75,4 @@ install(FILES ${CMAKE_CURRENT_BINARY_DIR}/LlamaConfig.cmake
|
||||
DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/Llama)
|
||||
|
||||
set_target_properties(llama PROPERTIES PUBLIC_HEADER ${CMAKE_CURRENT_SOURCE_DIR}/llama.h)
|
||||
install(TARGETS llama LIBRARY PUBLIC_HEADER)
|
||||
install(TARGETS llama LIBRARY PUBLIC_HEADER)
|
||||
@@ -2,16 +2,10 @@
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
|
||||
|
||||
[<img src="./assets/header_model_release.png" alt="BitNet Model on Hugging Face" width="800"/>](https://huggingface.co/microsoft/BitNet-b1.58-2B-4T)
|
||||
bitnet.cpp is the official inference framework for 1-bit LLMs (e.g., BitNet b1.58). It offers a suite of optimized kernels, that support **fast** and **lossless** inference of 1.58-bit models on CPU (with NPU and GPU support coming next).
|
||||
|
||||
Try it out via this [demo](https://bitnet-demo.azurewebsites.net/), or build and run it on your own [CPU](https://github.com/microsoft/BitNet?tab=readme-ov-file#build-from-source) or [GPU](https://github.com/microsoft/BitNet/blob/main/gpu/README.md).
|
||||
|
||||
bitnet.cpp is the official inference framework for 1-bit LLMs (e.g., BitNet b1.58). It offers a suite of optimized kernels, that support **fast** and **lossless** inference of 1.58-bit models on CPU and GPU (NPU support will coming next).
|
||||
|
||||
The first release of bitnet.cpp is to support inference on CPUs. bitnet.cpp achieves speedups of **1.37x** to **5.07x** on ARM CPUs, with larger models experiencing greater performance gains. Additionally, it reduces energy consumption by **55.4%** to **70.0%**, further boosting overall efficiency. On x86 CPUs, speedups range from **2.37x** to **6.17x** with energy reductions between **71.9%** to **82.2%**. Furthermore, bitnet.cpp can run a 100B BitNet b1.58 model on a single CPU, achieving speeds comparable to human reading (5-7 tokens per second), significantly enhancing the potential for running LLMs on local devices. Please refer to the [technical report](https://arxiv.org/abs/2410.16144) for more details.
|
||||
|
||||
<img src="./assets/m2_performance.jpg" alt="m2_performance" width="800"/>
|
||||
<img src="./assets/intel_performance.jpg" alt="m2_performance" width="800"/>
|
||||
<img src="./assets/f_compa.png" alt="performance" width="800"/>
|
||||
<!-- <img src="./assets/intel_performance.jpg" alt="m2_performance" width="800"/> -->
|
||||
|
||||
>The tested models are dummy setups used in a research context to demonstrate the inference performance of bitnet.cpp.
|
||||
|
||||
@@ -22,9 +16,7 @@ A demo of bitnet.cpp running a BitNet b1.58 3B model on Apple M2:
|
||||
https://github.com/user-attachments/assets/7f46b736-edec-4828-b809-4be780a3e5b1
|
||||
|
||||
## What's New:
|
||||
- 05/20/2025 [BitNet Official GPU inference kernel](https://github.com/microsoft/BitNet/blob/main/gpu/README.md) 
|
||||
- 04/14/2025 [BitNet Official 2B Parameter Model on Hugging Face](https://huggingface.co/microsoft/BitNet-b1.58-2B-4T)
|
||||
- 02/18/2025 [Bitnet.cpp: Efficient Edge Inference for Ternary LLMs](https://arxiv.org/abs/2502.11880)
|
||||
- 02/18/2025 [Bitnet.cpp: Efficient Edge Inference for Ternary LLMs](https://arxiv.org/abs/2502.11880) 
|
||||
- 11/08/2024 [BitNet a4.8: 4-bit Activations for 1-bit LLMs](https://arxiv.org/abs/2411.04965)
|
||||
- 10/21/2024 [1-bit AI Infra: Part 1.1, Fast and Lossless BitNet b1.58 Inference on CPUs](https://arxiv.org/abs/2410.16144)
|
||||
- 10/17/2024 bitnet.cpp 1.0 released.
|
||||
@@ -35,38 +27,9 @@ https://github.com/user-attachments/assets/7f46b736-edec-4828-b809-4be780a3e5b1
|
||||
## Acknowledgements
|
||||
|
||||
This project is based on the [llama.cpp](https://github.com/ggerganov/llama.cpp) framework. We would like to thank all the authors for their contributions to the open-source community. Also, bitnet.cpp's kernels are built on top of the Lookup Table methodologies pioneered in [T-MAC](https://github.com/microsoft/T-MAC/). For inference of general low-bit LLMs beyond ternary models, we recommend using T-MAC.
|
||||
## Official Models
|
||||
<table>
|
||||
</tr>
|
||||
<tr>
|
||||
<th rowspan="2">Model</th>
|
||||
<th rowspan="2">Parameters</th>
|
||||
<th rowspan="2">CPU</th>
|
||||
<th colspan="3">Kernel</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<th>I2_S</th>
|
||||
<th>TL1</th>
|
||||
<th>TL2</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/microsoft/BitNet-b1.58-2B-4T">BitNet-b1.58-2B-4T</a></td>
|
||||
<td rowspan="2">2.4B</td>
|
||||
<td>x86</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>✅</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
## Supported Models
|
||||
❗️**We use existing 1-bit LLMs available on [Hugging Face](https://huggingface.co/) to demonstrate the inference capabilities of bitnet.cpp. We hope the release of bitnet.cpp will inspire the development of 1-bit LLMs in large-scale settings in terms of model size and training tokens.**
|
||||
❗️**We use existing 1-bit LLMs available on [Hugging Face](https://huggingface.co/) to demonstrate the inference capabilities of bitnet.cpp. These models are neither trained nor released by Microsoft. We hope the release of bitnet.cpp will inspire the development of 1-bit LLMs in large-scale settings in terms of model size and training tokens.**
|
||||
|
||||
<table>
|
||||
</tr>
|
||||
@@ -78,8 +41,9 @@ This project is based on the [llama.cpp](https://github.com/ggerganov/llama.cpp)
|
||||
</tr>
|
||||
<tr>
|
||||
<th>I2_S</th>
|
||||
<th>TL1</th>
|
||||
<th>TL2</th>
|
||||
<th>TL1(TL1_1)</th>
|
||||
<th>TL2(TL2_1)</th>
|
||||
<th>TL2-Loss(TL2_0)</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/1bitLLM/bitnet_b1_58-large">bitnet_b1_58-large</a></td>
|
||||
@@ -88,12 +52,14 @@ This project is based on the [llama.cpp](https://github.com/ggerganov/llama.cpp)
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>✅</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/1bitLLM/bitnet_b1_58-3B">bitnet_b1_58-3B</a></td>
|
||||
@@ -102,12 +68,14 @@ This project is based on the [llama.cpp](https://github.com/ggerganov/llama.cpp)
|
||||
<td>❌</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/HF1BitLLM/Llama3-8B-1.58-100B-tokens">Llama3-8B-1.58-100B-tokens</a></td>
|
||||
@@ -116,12 +84,14 @@ This project is based on the [llama.cpp](https://github.com/ggerganov/llama.cpp)
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>✅</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/collections/tiiuae/falcon3-67605ae03578be86e4e87026">Falcon3 Family</a></td>
|
||||
@@ -130,12 +100,14 @@ This project is based on the [llama.cpp](https://github.com/ggerganov/llama.cpp)
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>✅</td>
|
||||
<td>✅</td>
|
||||
<td>❌</td>
|
||||
<td>✅</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
@@ -161,11 +133,11 @@ This project is based on the [llama.cpp](https://github.com/ggerganov/llama.cpp)
|
||||
### Build from source
|
||||
|
||||
> [!IMPORTANT]
|
||||
> If you are using Windows, please remember to always use a Developer Command Prompt / PowerShell for VS2022 for the following commands. Please refer to the FAQs below if you see any issues.
|
||||
> If you are using Windows, please remember to always use a Developer Command Prompt / PowerShell for VS2022 for the following commands
|
||||
|
||||
1. Clone the repo
|
||||
```bash
|
||||
git clone --recursive https://github.com/microsoft/BitNet.git
|
||||
git clone --recursive -b paper https://github.com/microsoft/BitNet.git
|
||||
cd BitNet
|
||||
```
|
||||
2. Install the dependencies
|
||||
@@ -178,13 +150,15 @@ pip install -r requirements.txt
|
||||
```
|
||||
3. Build the project
|
||||
```bash
|
||||
# Manually download the model and run with local path
|
||||
huggingface-cli download microsoft/BitNet-b1.58-2B-4T-gguf --local-dir models/BitNet-b1.58-2B-4T
|
||||
python setup_env.py -md models/BitNet-b1.58-2B-4T -q i2_s
|
||||
# Download the model from Hugging Face, convert it to quantized gguf format, and build the project
|
||||
python setup_env.py --hf-repo 1bitLLM/bitnet_b1_58-large -q i2_s
|
||||
|
||||
# Or you can manually download the model and run with local path
|
||||
huggingface-cli download 1bitLLM/bitnet_b1_58-large --local-dir models/bitnet_b1_58-large
|
||||
python setup_env.py -md models/bitnet_b1_58-large -q i2_s
|
||||
```
|
||||
<pre>
|
||||
usage: setup_env.py [-h] [--hf-repo {1bitLLM/bitnet_b1_58-large,1bitLLM/bitnet_b1_58-3B,HF1BitLLM/Llama3-8B-1.58-100B-tokens,tiiuae/Falcon3-1B-Instruct-1.58bit,tiiuae/Falcon3-3B-Instruct-1.58bit,tiiuae/Falcon3-7B-Instruct-1.58bit,tiiuae/Falcon3-10B-Instruct-1.58bit}] [--model-dir MODEL_DIR] [--log-dir LOG_DIR] [--quant-type {i2_s,tl1}] [--quant-embd]
|
||||
usage: setup_env.py [-h] [--hf-repo {1bitLLM/bitnet_b1_58-large,1bitLLM/bitnet_b1_58-3B,HF1BitLLM/Llama3-8B-1.58-100B-tokens,tiiuae/Falcon3-1B-Instruct-1.58bit,tiiuae/Falcon3-3B-Instruct-1.58bit,tiiuae/Falcon3-7B-Instruct-1.58bit,tiiuae/Falcon3-10B-Instruct-1.58bit}] [--model-dir MODEL_DIR] [--log-dir LOG_DIR] [--quant-type {i2_s,tl1,tl2,tl2-loss}] [--quant-embd]
|
||||
[--use-pretuned]
|
||||
|
||||
Setup the environment for running inference
|
||||
@@ -197,7 +171,7 @@ optional arguments:
|
||||
Directory to save/load the model
|
||||
--log-dir LOG_DIR, -ld LOG_DIR
|
||||
Directory to save the logging info
|
||||
--quant-type {i2_s,tl1}, -q {i2_s,tl1}
|
||||
--quant-type {i2_s,tl1,tl2,tl2-loss}, -q {i2_s,tl1,tl2,tl2-loss}
|
||||
Quantization type
|
||||
--quant-embd Quantize the embeddings to f16
|
||||
--use-pretuned, -p Use the pretuned kernel parameters
|
||||
@@ -206,7 +180,7 @@ optional arguments:
|
||||
### Basic usage
|
||||
```bash
|
||||
# Run inference with the quantized model
|
||||
python run_inference.py -m models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf -p "You are a helpful assistant" -cnv
|
||||
python run_inference.py -m models/Falcon3-7B-Instruct-1.58bit/ggml-model-i2_s.gguf -p "You are a helpful assistant" -cnv
|
||||
```
|
||||
<pre>
|
||||
usage: run_inference.py [-h] [-m MODEL] [-n N_PREDICT] -p PROMPT [-t THREADS] [-c CTX_SIZE] [-temp TEMPERATURE] [-cnv]
|
||||
@@ -278,36 +252,5 @@ python utils/generate-dummy-bitnet-model.py models/bitnet_b1_58-large --outfile
|
||||
# Run benchmark with the generated model, use -m to specify the model path, -p to specify the prompt processed, -n to specify the number of token to generate
|
||||
python utils/e2e_benchmark.py -m models/dummy-bitnet-125m.tl1.gguf -p 512 -n 128
|
||||
```
|
||||
### FAQ (Frequently Asked Questions)📌
|
||||
|
||||
#### Q1: The build dies with errors building llama.cpp due to issues with std::chrono in log.cpp?
|
||||
|
||||
**A:**
|
||||
This is an issue introduced in recent version of llama.cpp. Please refer to this [commit](https://github.com/tinglou/llama.cpp/commit/4e3db1e3d78cc1bcd22bcb3af54bd2a4628dd323) in the [discussion](https://github.com/abetlen/llama-cpp-python/issues/1942) to fix this issue.
|
||||
|
||||
#### Q2: How to build with clang in conda environment on windows?
|
||||
|
||||
**A:**
|
||||
Before building the project, verify your clang installation and access to Visual Studio tools by running:
|
||||
```
|
||||
clang -v
|
||||
```
|
||||
|
||||
This command checks that you are using the correct version of clang and that the Visual Studio tools are available. If you see an error message such as:
|
||||
```
|
||||
'clang' is not recognized as an internal or external command, operable program or batch file.
|
||||
```
|
||||
|
||||
It indicates that your command line window is not properly initialized for Visual Studio tools.
|
||||
|
||||
• If you are using Command Prompt, run:
|
||||
```
|
||||
"C:\Program Files\Microsoft Visual Studio\2022\Professional\Common7\Tools\VsDevCmd.bat" -startdir=none -arch=x64 -host_arch=x64
|
||||
```
|
||||
|
||||
• If you are using Windows PowerShell, run the following commands:
|
||||
```
|
||||
Import-Module "C:\Program Files\Microsoft Visual Studio\2022\Professional\Common7\Tools\Microsoft.VisualStudio.DevShell.dll" Enter-VsDevShell 3f0e31ad -SkipAutomaticLocation -DevCmdArguments "-arch=x64 -host_arch=x64"
|
||||
```
|
||||
|
||||
These steps will initialize your environment and allow you to use the correct Visual Studio tools.
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 84 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 14 KiB |
@@ -1,93 +0,0 @@
|
||||
# BitNet Inference Kernel
|
||||
|
||||
This repository provides a highly efficient GEMV kernel implementation for the BitNet model, optimized for W2A8 inference — 2-bit weights and 8-bit activations. It is tailored for use with the [BitNet-b1.58-2B-4T](https://arxiv.org/abs/2504.12285) model.
|
||||
|
||||
## Features
|
||||
|
||||
- Support for W2A8 (2-bit weight × 8-bit activation) GEMV computation
|
||||
- Custom CUDA kernels with low-latency execution
|
||||
- Optimizations for memory access, decoding, and compute throughput
|
||||
|
||||
## Usage
|
||||
|
||||
Installation and kernel performance tests:
|
||||
|
||||
```bash
|
||||
# (Recommended) Create a new conda environment
|
||||
conda create --name bitnet-gpu "python<3.13"
|
||||
conda activate bitnet-gpu
|
||||
|
||||
# Install dependencies
|
||||
pip install -r requirements.txt
|
||||
|
||||
# Build the kernel
|
||||
cd bitnet_kernels
|
||||
bash compile.sh
|
||||
cd ..
|
||||
|
||||
# Run performance tests
|
||||
python test.py
|
||||
```
|
||||
|
||||
End-to-end inference:
|
||||
|
||||
```bash
|
||||
# Download and convert the BitNet-b1.58-2B model
|
||||
mkdir checkpoints
|
||||
huggingface-cli download microsoft/bitnet-b1.58-2B-4T-bf16 --local-dir ./checkpoints/bitnet-b1.58-2B-4T-bf16
|
||||
python ./convert_safetensors.py --safetensors_file ./checkpoints/bitnet-b1.58-2B-4T-bf16/model.safetensors --output checkpoints/model_state.pt --model_name 2B
|
||||
python ./convert_checkpoint.py --input ./checkpoints/model_state.pt
|
||||
rm ./checkpoints/model_state.pt
|
||||
|
||||
# Inference
|
||||
python3 ./generate.py ./checkpoints/ --interactive --chat_format
|
||||
```
|
||||
|
||||
## Optimizations
|
||||
|
||||
### Weight Permutation
|
||||
|
||||
The weight matrix is divided into 16×32 blocks to optimize memory access patterns.
|
||||
|
||||
Within each block, values are stored contiguously in memory and permuted to facilitate efficient access and processing.
|
||||
|
||||
See `convert_checkpoint.py` for details.
|
||||
|
||||
### Fast Decoding
|
||||
|
||||
Every 16 two-bit values are packed into a single 32-bit integer using the following interleaving pattern:
|
||||
```
|
||||
[0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15]
|
||||
```
|
||||
|
||||
This layout is designed to accelerate decoding by enabling efficient extraction of 4 values at a time into `int8`.
|
||||
|
||||
### `dp4a` Instruction
|
||||
|
||||
We use the `dp4a` instruction to accelerate low-precision dot product operations.
|
||||
|
||||
This instruction performs a dot product between two 4-element vectors (each stored in a 32-bit word as 8-bit integers) and accumulates the result into a 32-bit integer.
|
||||
|
||||
It significantly improves GEMV throughput when processing quantized weights and activations.
|
||||
|
||||
|
||||
## Performance
|
||||
|
||||
Kernel performance (tested on NVIDIA A100 40GB GPU):
|
||||
|
||||
| Shape (N×K) | W2A8 Latency (us) | BF16 Latency (us) | Speedup Ratio |
|
||||
|---------------------|-------------------|-------------------|----------------------|
|
||||
| 2560 × 2560 | 13.32 | 18.32 | 1.38 |
|
||||
| 3840 × 2560 | 14.90 | 18.87 | 1.27 |
|
||||
| 13824 × 2560 | 18.75 | 59.51 | 3.17 |
|
||||
| 2560 × 6912 | 14.49 | 37.78 | 2.61 |
|
||||
| 3200 × 3200 | 14.61 | 19.08 | 1.31 |
|
||||
| 4800 × 3200 | 13.09 | 21.84 | 1.67 |
|
||||
| 3200 × 10240 | 19.64 | 60.79 | 3.10 |
|
||||
| 20480 × 3200 | 30.99 | 112.39 | 3.63 |
|
||||
|
||||
Generation throughput:
|
||||
|
||||
| BF16 (tokens/s) | W2A8 (tokens/s) | Speedup Ratio |
|
||||
|---|---|---|
|
||||
| 10.9 | 213.3 | 19.6 |
|
||||
@@ -1,37 +0,0 @@
|
||||
#include "bitnet_kernels.h"
|
||||
|
||||
extern "C" void bitlinear_int8xint2(int8_t* input0, int8_t* input1, __nv_bfloat16* output0, __nv_bfloat16* s, __nv_bfloat16* ws, int M, int N, int K, cudaStream_t stream){
|
||||
if (M == 1 && N == 3840 && K == 2560){
|
||||
ladder_int8xint2_kernel<1, 3840, 2560, 3, 8, 16><<<dim3(240, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if (M == 1 && N == 2560 && K == 2560){
|
||||
ladder_int8xint2_kernel<1, 2560, 2560, 1, 8, 16><<<dim3(160, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if (M == 1 && N == 13824 && K == 2560){
|
||||
ladder_int8xint2_kernel<1, 13824, 2560, 2, 8, 16><<<dim3(864, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if (M == 1 && N == 2560 && K == 6912){
|
||||
ladder_int8xint2_kernel<1, 2560, 6912, 1, 8, 16><<<dim3(160, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if(M == 1 && N == 4800 && K == 3200){
|
||||
ladder_int8xint2_kernel<1, 4800, 3200, 6, 8, 16><<<dim3(300, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if(M == 1 && N == 3200 && K == 3200){
|
||||
ladder_int8xint2_kernel<1, 3200, 3200, 1, 8, 16><<<dim3(200, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if(M == 1 && N == 20480 && K == 3200){
|
||||
ladder_int8xint2_kernel<1, 20480, 3200, 2, 8, 16><<<dim3(1280, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if(M == 1 && N == 3200 && K == 10240){
|
||||
ladder_int8xint2_kernel<1, 3200, 10240, 1, 8, 16><<<dim3(200, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if(M == 1 && N == 5120 && K == 27648){
|
||||
ladder_int8xint2_kernel<1, 5120, 27648, 1, 8, 16><<<dim3(320, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else if(M == 1 && N == 55296 && K == 5120){
|
||||
ladder_int8xint2_kernel<1, 55296, 5120, 1, 8, 16><<<dim3(3456, 1, 1), dim3(8, 16, 1), 0, stream>>>(input0, input1, output0, s, ws);
|
||||
}
|
||||
else{
|
||||
std::cout << "required ladder gemm kernel: M " << M << ", N " << N << ", K " << K << std::endl;
|
||||
}
|
||||
}
|
||||
@@ -1,83 +0,0 @@
|
||||
#include <cuda_runtime.h>
|
||||
#include <math_constants.h>
|
||||
#include <math.h>
|
||||
#include <mma.h>
|
||||
#include <iostream>
|
||||
#include <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
|
||||
#if (((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 4)) || (__CUDACC_VER_MAJOR__ > 11))
|
||||
#define TVM_ENABLE_L2_PREFETCH 1
|
||||
#else
|
||||
#define TVM_ENABLE_L2_PREFETCH 0
|
||||
#endif
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 800
|
||||
#define TVM_ENBALE_EFFICIENT_SMEM_PTR_CAST 1
|
||||
#else
|
||||
#define TVM_ENBALE_EFFICIENT_SMEM_PTR_CAST 0
|
||||
#endif
|
||||
|
||||
template <typename T1, typename T2>
|
||||
__device__ void decode_i2s_to_i8s(T1 *_i2s, T2 *_i8s, const int N = 16)
|
||||
{
|
||||
// convert 8 int2b_t to 8 int8b_t -> 2 int32
|
||||
uint *i8s = reinterpret_cast<uint *>(_i8s);
|
||||
|
||||
// i2s = {e0, e4, e8, e12, e1, e5, e9, e13, e2, e6, e10, e14, e3, e7, e11, e15}
|
||||
uint const i2s = *_i2s;
|
||||
|
||||
static constexpr uint immLut = (0xf0 & 0xcc) | 0xaa; // 0b11101010
|
||||
static constexpr uint BOTTOM_MASK = 0x03030303; // 0xf -> 0b11 select 0,3
|
||||
static constexpr uint I4s_TO_I8s_MAGIC_NUM = 0x00000000;
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < (N / 4); i++)
|
||||
{
|
||||
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
|
||||
: "=r"(i8s[i])
|
||||
: "r"(i2s >> (2 * i)), "n"(BOTTOM_MASK), "n"(I4s_TO_I8s_MAGIC_NUM), "n"(immLut));
|
||||
i8s[i] = __vsubss4(i8s[i], 0x02020202);
|
||||
}
|
||||
}
|
||||
|
||||
template <int M, int N, int K, int ws_num, int K_block_size, int N_block_size>
|
||||
__global__ void __launch_bounds__(128) ladder_int8xint2_kernel(int8_t* __restrict__ A, int8_t* __restrict__ B, __nv_bfloat16* __restrict__ dtype_transform, __nv_bfloat16* __restrict__ s, __nv_bfloat16* __restrict__ ws) {
|
||||
constexpr int K_per_loop = 16;
|
||||
constexpr int wmma_K = 32;
|
||||
constexpr int wmma_N = 16;
|
||||
int in_thread_C_local[1];
|
||||
signed char A_local[K_per_loop];
|
||||
int B_reshape_local[1];
|
||||
signed char B_decode_local[K_per_loop];
|
||||
int red_buf0[1];
|
||||
in_thread_C_local[0] = 0;
|
||||
#pragma unroll
|
||||
for (int k_0 = 0; k_0 < K/(K_per_loop * K_block_size); ++k_0) {
|
||||
*(int4*)(A_local + 0) = *(int4*)(A + ((k_0 * K_per_loop * K_block_size) + (((int)threadIdx.x) * K_per_loop)));
|
||||
B_reshape_local[0] = *(int*)(B +
|
||||
(((int)blockIdx.x) * N_block_size * K / 4) +
|
||||
(k_0 * K_block_size * K_per_loop * wmma_N / 4) +
|
||||
((((int)threadIdx.x) >> 1) * wmma_K * wmma_N / 4) +
|
||||
((((int)threadIdx.y) >> 3) * (wmma_K * wmma_N / 2) / 4) +
|
||||
((((int)threadIdx.x) & 1) * (wmma_K * wmma_N / 4) / 4) +
|
||||
((((int)threadIdx.y) & 7) * (wmma_K / 2) / 4)
|
||||
);
|
||||
decode_i2s_to_i8s(B_reshape_local, B_decode_local, 16);
|
||||
#pragma unroll
|
||||
for (int k_2_0 = 0; k_2_0 < 4; ++k_2_0) {
|
||||
in_thread_C_local[0] = __dp4a(*(int *)&A_local[((k_2_0 * 4))],*(int *)&B_decode_local[((k_2_0 * 4))], in_thread_C_local[0]);
|
||||
}
|
||||
}
|
||||
red_buf0[0] = in_thread_C_local[0];
|
||||
#pragma unroll
|
||||
for (int offset = K_block_size/2; offset > 0; offset /= 2) {
|
||||
red_buf0[0] += __shfl_down_sync(__activemask(), red_buf0[0], offset, K_block_size);
|
||||
}
|
||||
int out_idx = ((((int)blockIdx.x) * N_block_size) + ((int)threadIdx.y));
|
||||
int ws_idx = out_idx / (N / ws_num);
|
||||
if (threadIdx.x == 0)
|
||||
dtype_transform[out_idx] = (__nv_bfloat16)(((float)red_buf0[0])/(float)s[0]*(float)ws[ws_idx]);
|
||||
}
|
||||
@@ -1,3 +0,0 @@
|
||||
nvcc -std=c++17 -Xcudafe --diag_suppress=177 --compiler-options -fPIC -lineinfo --shared bitnet_kernels.cu -lcuda -gencode=arch=compute_80,code=compute_80 -o libbitnet.so
|
||||
|
||||
|
||||
@@ -1,13 +0,0 @@
|
||||
from setuptools import setup
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
|
||||
setup(
|
||||
name='bitlinear_cpp',
|
||||
ext_modules=[
|
||||
CUDAExtension('bitlinear_cuda', [
|
||||
'bitnet_kernels.cu',
|
||||
])
|
||||
],
|
||||
cmdclass={
|
||||
'build_ext': BuildExtension
|
||||
})
|
||||
@@ -1,100 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from dataclasses import dataclass
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from safetensors.torch import save_file
|
||||
import model
|
||||
from pack_weight import convert_weight_int8_to_int2
|
||||
|
||||
@torch.inference_mode()
|
||||
def convert_ts_checkpoint(
|
||||
*,
|
||||
input_path: str = "",
|
||||
) -> None:
|
||||
|
||||
config = model.ModelArgs()
|
||||
print(f"Model config {config.__dict__}")
|
||||
|
||||
def quant_weight_int8(weight):
|
||||
s = 1.0 / weight.abs().mean().clamp_(min=1e-5)
|
||||
new_weight = (weight * s).round().clamp(-1, 1).to(torch.int8)
|
||||
new_scale = (1.0 / s).to(torch.bfloat16)
|
||||
return new_weight, new_scale.reshape(1)
|
||||
|
||||
def quant_weight_fp16(weight):
|
||||
s = 1.0 / weight.abs().mean().clamp_(min=1e-5)
|
||||
new_weight = (weight * s).round().clamp(-1, 1) / s
|
||||
return new_weight
|
||||
|
||||
def convert_int8_to_int2(weight):
|
||||
return convert_weight_int8_to_int2(weight)
|
||||
|
||||
merged_result = torch.load(input_path, map_location="cpu", mmap=True)
|
||||
int2_result = {}
|
||||
fp16_result = {}
|
||||
zero = torch.zeros(1).to(torch.bfloat16)
|
||||
for key, value in merged_result.items():
|
||||
if 'wqkv' in key:
|
||||
wq = value[:config.dim]
|
||||
wk = value[config.dim:config.dim // config.n_heads * config.n_kv_heads + config.dim]
|
||||
wv = value[config.dim // config.n_heads * config.n_kv_heads + config.dim:]
|
||||
wq_weight, wa_scale = quant_weight_int8(wq)
|
||||
wk_weight, wb_scale = quant_weight_int8(wk)
|
||||
wv_weight, wc_scale = quant_weight_int8(wv)
|
||||
wqkv_weight = torch.cat([wq_weight, wk_weight, wv_weight], dim=0)
|
||||
wqkv_scale = torch.cat([wa_scale, wb_scale, wc_scale, zero], dim=0)
|
||||
int2_result[key] = convert_int8_to_int2(wqkv_weight)
|
||||
int2_result[key.replace('weight', 'weight_scale')] = wqkv_scale
|
||||
|
||||
wq_weight = quant_weight_fp16(wq)
|
||||
wk_weight = quant_weight_fp16(wk)
|
||||
wv_weight = quant_weight_fp16(wv)
|
||||
wqkv_weight = torch.cat([wq_weight, wk_weight, wv_weight], dim=0)
|
||||
fp16_result[key] = wqkv_weight
|
||||
elif 'w13' in key:
|
||||
w1 = value[:config.ffn_dim]
|
||||
w3 = value[config.ffn_dim:]
|
||||
w1_weight, w1_scale = quant_weight_int8(w1)
|
||||
w3_weight, w3_scale = quant_weight_int8(w3)
|
||||
w13_weight = torch.cat([w1_weight, w3_weight], dim=0)
|
||||
w13_scale = torch.cat([w1_scale, w3_scale, zero, zero], dim=0)
|
||||
int2_result[key] = convert_int8_to_int2(w13_weight)
|
||||
int2_result[key.replace('weight', 'weight_scale')] = w13_scale
|
||||
|
||||
w1_weight = quant_weight_fp16(w1)
|
||||
w3_weight = quant_weight_fp16(w3)
|
||||
w13_weight = torch.cat([w1_weight, w3_weight], dim=0)
|
||||
fp16_result[key] = w13_weight
|
||||
elif 'w2' in key or 'wo' in key:
|
||||
weight, scale = quant_weight_int8(value)
|
||||
scale = torch.cat([scale, zero, zero, zero], dim=0)
|
||||
int2_result[key] = convert_int8_to_int2(weight)
|
||||
int2_result[key.replace('weight', 'weight_scale')] = scale
|
||||
|
||||
weight = quant_weight_fp16(value)
|
||||
fp16_result[key] = weight
|
||||
else:
|
||||
int2_result[key] = value.clone()
|
||||
fp16_result[key] = value.clone()
|
||||
|
||||
output_dir = os.path.dirname(input_path)
|
||||
print(f"Saving checkpoint to {output_dir}/model_state_int2.pt")
|
||||
torch.save(int2_result, f"{output_dir}/model_state_int2.pt")
|
||||
|
||||
print(f"Saving checkpoint to {output_dir}/model_state_fp16.pt")
|
||||
torch.save(fp16_result, f"{output_dir}/model_state_fp16.pt")
|
||||
|
||||
if __name__ == '__main__':
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser(description='Convert TorchScale checkpoint.')
|
||||
parser.add_argument('--input', type=str)
|
||||
|
||||
args = parser.parse_args()
|
||||
convert_ts_checkpoint(
|
||||
input_path=args.input,
|
||||
)
|
||||
@@ -1,116 +0,0 @@
|
||||
import re
|
||||
import torch
|
||||
from pathlib import Path
|
||||
from safetensors.torch import load_file
|
||||
from einops import rearrange
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
transformer_configs = {
|
||||
"2B": dict(n_layer=30, n_head=20, dim=2560, vocab_size=128256, n_local_heads=5, intermediate_size=6912),
|
||||
}
|
||||
|
||||
@dataclass
|
||||
class ModelArgs:
|
||||
block_size: int = 4096
|
||||
vocab_size: int = 32000
|
||||
n_layer: int = 32
|
||||
n_head: int = 32
|
||||
dim: int = 4096
|
||||
intermediate_size: int = None
|
||||
n_local_heads: int = -1
|
||||
head_dim: int = 64
|
||||
rope_base: float = 10000
|
||||
norm_eps: float = 1e-5
|
||||
|
||||
def __post_init__(self):
|
||||
if self.n_local_heads == -1:
|
||||
self.n_local_heads = self.n_head
|
||||
if self.intermediate_size is None:
|
||||
hidden_dim = 4 * self.dim
|
||||
n_hidden = int(2 * hidden_dim / 3)
|
||||
self.intermediate_size = n_hidden + (256 - n_hidden % 256) if n_hidden % 256 else n_hidden
|
||||
self.head_dim = self.dim // self.n_head
|
||||
|
||||
@classmethod
|
||||
def from_name(cls, name: str):
|
||||
if name in transformer_configs:
|
||||
return cls(**transformer_configs[name])
|
||||
config = [k for k in transformer_configs if k in name.upper() or k in name]
|
||||
assert len(config) == 1, f"Unknown model name: {name}"
|
||||
return cls(**transformer_configs[config[0]])
|
||||
|
||||
def invert_convert_q(w: torch.Tensor, config: ModelArgs) -> torch.Tensor:
|
||||
return rearrange(w, '(h l d) i -> (h d l) i', h=config.n_head, l=2)
|
||||
|
||||
def invert_convert_k(w: torch.Tensor, config: ModelArgs) -> torch.Tensor:
|
||||
return rearrange(w, '(h l d) i -> (h d l) i', h=config.n_local_heads, l=2)
|
||||
|
||||
def convert_back(
|
||||
safetensors_path: str,
|
||||
output_file: str,
|
||||
model_name: Optional[str] = None,
|
||||
):
|
||||
st_dict = load_file(safetensors_path)
|
||||
|
||||
cfg = ModelArgs.from_name(model_name)
|
||||
print(f"Using model configurations: {cfg}")
|
||||
|
||||
recovered: dict = {}
|
||||
|
||||
for layer in range(cfg.n_layer):
|
||||
base = f"model.layers.{layer}."
|
||||
|
||||
wq = st_dict[f"{base}self_attn.q_proj.weight"]
|
||||
wk = st_dict[f"{base}self_attn.k_proj.weight"]
|
||||
wv = st_dict[f"{base}self_attn.v_proj.weight"]
|
||||
|
||||
wq = invert_convert_q(wq, cfg)
|
||||
wk = invert_convert_k(wk, cfg)
|
||||
|
||||
wqkv = torch.cat([wq, wk, wv], dim=0)
|
||||
recovered[f"layers.{layer}.attention.wqkv.weight"] = wqkv
|
||||
|
||||
recovered[f"layers.{layer}.attention.wo.weight"] = st_dict[f"{base}self_attn.o_proj.weight"]
|
||||
|
||||
recovered[f"layers.{layer}.attention_norm.weight"] = st_dict[f"{base}input_layernorm.weight"]
|
||||
recovered[f"layers.{layer}.ffn_norm.weight"] = st_dict[f"{base}post_attention_layernorm.weight"]
|
||||
recovered[f"layers.{layer}.attention.attn_sub_norm.weight"] = st_dict[f"{base}self_attn.attn_sub_norm.weight"]
|
||||
recovered[f"layers.{layer}.feed_forward.ffn_sub_norm.weight"] = st_dict[f"{base}mlp.ffn_sub_norm.weight"]
|
||||
|
||||
gate = st_dict[f"{base}mlp.gate_proj.weight"]
|
||||
up = st_dict[f"{base}mlp.up_proj.weight"]
|
||||
w13 = torch.cat([gate, up], dim=0)
|
||||
recovered[f"layers.{layer}.feed_forward.w13.weight"] = w13
|
||||
|
||||
recovered[f"layers.{layer}.feed_forward.w2.weight"] = st_dict[f"{base}mlp.down_proj.weight"]
|
||||
|
||||
recovered["tok_embeddings.weight"] = st_dict["model.embed_tokens.weight"]
|
||||
recovered["output.weight"] = st_dict["model.embed_tokens.weight"]
|
||||
recovered["norm.weight"] = st_dict["model.norm.weight"]
|
||||
|
||||
print(f"Saving to {output_file}")
|
||||
torch.save(recovered, output_file)
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser(description="Convert Safetensors back to Torch .pth checkpoint")
|
||||
parser.add_argument(
|
||||
"--safetensors_file", type=str, required=True,
|
||||
help="Path to input .safetensors file"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output", type=str, default="./checkpoints/model_state.pt",
|
||||
help="Path to output .pt file"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_name", type=str, default="2B",
|
||||
help="Model configuration name to use (e.g. 2B)"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
convert_back(
|
||||
safetensors_path=args.safetensors_file,
|
||||
output_file=args.output,
|
||||
model_name=args.model_name,
|
||||
)
|
||||
-359
@@ -1,359 +0,0 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import json
|
||||
import os
|
||||
import readline # type: ignore # noqa
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Iterable, Optional, Tuple, Union
|
||||
|
||||
import fire
|
||||
import model as fast
|
||||
import torch
|
||||
from stats import Stats
|
||||
from tokenizer import Tokenizer, ChatFormat
|
||||
import sample_utils
|
||||
from xformers.ops.fmha.attn_bias import (
|
||||
BlockDiagonalCausalWithOffsetPaddedKeysMask as AttnBias,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenArgs:
|
||||
gen_length: int = 32
|
||||
gen_bsz: int = 1
|
||||
prompt_length: int = 64
|
||||
|
||||
use_sampling: bool = False
|
||||
temperature: float = 0.8
|
||||
top_p: float = 0.9
|
||||
|
||||
|
||||
class FastGen:
|
||||
GRAPH_WARMUPS: int = 1
|
||||
tokenizer: Tokenizer
|
||||
|
||||
@staticmethod
|
||||
def build(
|
||||
ckpt_dir: str,
|
||||
gen_args: GenArgs,
|
||||
device: Union[torch.device, str],
|
||||
tokenizer_path: Optional[str] = None,
|
||||
num_layers: int = 13,
|
||||
use_full_vocab: bool = False,
|
||||
) -> "FastGen":
|
||||
"""
|
||||
Load a Llama or Code Llama checkpoint and return a new
|
||||
generator for this model.
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
model_args_prefill = fast.ModelArgs(use_kernel=False)
|
||||
model_args_decode = fast.ModelArgs(use_kernel=True)
|
||||
tokenizer = Tokenizer("./tokenizer.model")
|
||||
|
||||
torch.set_default_device(device)
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
|
||||
prefill_model = fast.Transformer(model_args_prefill)
|
||||
decode_model = fast.Transformer(model_args_decode)
|
||||
|
||||
fp16_ckpt_path = str(Path(ckpt_dir) / "model_state_fp16.pt")
|
||||
fp16_checkpoint = torch.load(fp16_ckpt_path, map_location="cpu")
|
||||
int2_ckpt_path = str(Path(ckpt_dir) / "model_state_int2.pt")
|
||||
int2_checkpoint = torch.load(int2_ckpt_path, map_location="cpu")
|
||||
prefill_model.load_state_dict(fp16_checkpoint, strict=True)
|
||||
decode_model.load_state_dict(int2_checkpoint, strict=True)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
print(f"loaded model in {time.time() - start_time:.2f} seconds")
|
||||
start_time = time.time()
|
||||
|
||||
return FastGen(gen_args, model_args_prefill, prefill_model, decode_model, tokenizer)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
args: GenArgs,
|
||||
model_args: fast.ModelArgs,
|
||||
prefill_model: fast.Transformer,
|
||||
decode_model: fast.Transformer,
|
||||
tokenizer: Tokenizer,
|
||||
):
|
||||
self.gen_args = args
|
||||
self.max_seq_length = args.prompt_length + args.gen_length
|
||||
self.model_args = model_args
|
||||
# self.model = model
|
||||
self.prefill_model = prefill_model
|
||||
self.decode_model = decode_model
|
||||
self.tokenizer = tokenizer
|
||||
self._prefill_cuda_graph, self._prefill_compile_model, self._prefill_inputs, self._prefill_logits = None, None, None, None
|
||||
self._generate_cuda_graph, self._generate_compile_model, self._generate_inputs, self._generate_logits = None, None, None, None
|
||||
self._cache = None
|
||||
start_time = time.time()
|
||||
self._prefill_compile_model = self.compile_prefill()
|
||||
self._generate_compile_model = self.compile_generate()
|
||||
print(f"compiled model in {time.time() - start_time:.2f} seconds")
|
||||
|
||||
def compile_prefill(self):
|
||||
|
||||
if self._cache is None:
|
||||
self._cache = fast.make_cache(
|
||||
args=self.model_args,
|
||||
length=self.gen_args.gen_bsz * self.max_seq_length,
|
||||
)
|
||||
|
||||
seq_lens = [self.gen_args.prompt_length for _ in range(self.gen_args.gen_bsz)]
|
||||
|
||||
bias = AttnBias.from_seqlens(
|
||||
q_seqlen=seq_lens,
|
||||
kv_seqlen=seq_lens,
|
||||
kv_padding=self.max_seq_length,
|
||||
)
|
||||
bias.q_seqinfo.to("cuda")
|
||||
bias.k_seqinfo.to("cuda")
|
||||
|
||||
tokens = torch.IntTensor([1] * self.gen_args.gen_bsz * self.gen_args.prompt_length).cuda()
|
||||
self._prefill_inputs = (tokens, bias)
|
||||
|
||||
s = torch.cuda.Stream()
|
||||
s.wait_stream(torch.cuda.current_stream())
|
||||
|
||||
with torch.cuda.stream(s):
|
||||
_ = self.prefill_model.forward_with_attn_bias(
|
||||
token_values=self._prefill_inputs[0],
|
||||
attn_bias=self._prefill_inputs[1],
|
||||
cache=self._cache,
|
||||
)
|
||||
torch.cuda.current_stream().wait_stream(s)
|
||||
|
||||
self._prefill_cuda_graph = torch.cuda.CUDAGraph()
|
||||
recording_kwargs = {}
|
||||
if "capture_error_mode" in torch.cuda.graph.__init__.__annotations__:
|
||||
# In PyTorch 2.1+ and nightlies from late Aug 2023,
|
||||
# we can do this to maybe avoid watchdog-related crashes
|
||||
recording_kwargs["capture_error_mode"] = "thread_local"
|
||||
with torch.cuda.graph(self._prefill_cuda_graph, **recording_kwargs):
|
||||
self._prefill_logits = self.prefill_model.forward_with_attn_bias(
|
||||
token_values=self._prefill_inputs[0],
|
||||
attn_bias=self._prefill_inputs[1],
|
||||
cache=self._cache,
|
||||
)
|
||||
|
||||
def replay(tokens, seq_lens=None):
|
||||
self._prefill_inputs[0].copy_(tokens)
|
||||
if seq_lens is not None:
|
||||
self._prefill_inputs[1].k_seqinfo.seqlen.copy_(seq_lens)
|
||||
|
||||
self._prefill_cuda_graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
return self._prefill_logits
|
||||
|
||||
return replay
|
||||
|
||||
def compile_generate(self):
|
||||
|
||||
if self._cache is None:
|
||||
self._cache = fast.make_cache(
|
||||
args=self.model_args,
|
||||
length=self.gen_args.gen_bsz * self.max_seq_length,
|
||||
)
|
||||
|
||||
seq_lens = [1 for _ in range(self.gen_args.gen_bsz)]
|
||||
kv_seq_lens = [self.gen_args.prompt_length for _ in range(self.gen_args.gen_bsz)]
|
||||
|
||||
bias = AttnBias.from_seqlens(
|
||||
q_seqlen=seq_lens,
|
||||
kv_seqlen=kv_seq_lens,
|
||||
kv_padding=self.max_seq_length,
|
||||
)
|
||||
bias.q_seqinfo.to("cuda")
|
||||
bias.k_seqinfo.to("cuda")
|
||||
|
||||
tokens = torch.IntTensor([1] * self.gen_args.gen_bsz).cuda()
|
||||
self._generate_inputs = (tokens, bias)
|
||||
|
||||
s = torch.cuda.Stream()
|
||||
s.wait_stream(torch.cuda.current_stream())
|
||||
|
||||
with torch.cuda.stream(s):
|
||||
_ = self.decode_model.forward_with_attn_bias(
|
||||
token_values=self._generate_inputs[0],
|
||||
attn_bias=self._generate_inputs[1],
|
||||
cache=self._cache,
|
||||
)
|
||||
torch.cuda.current_stream().wait_stream(s)
|
||||
|
||||
self._generate_cuda_graph = torch.cuda.CUDAGraph()
|
||||
recording_kwargs = {}
|
||||
if "capture_error_mode" in torch.cuda.graph.__init__.__annotations__:
|
||||
# In PyTorch 2.1+ and nightlies from late Aug 2023,
|
||||
# we can do this to maybe avoid watchdog-related crashes
|
||||
recording_kwargs["capture_error_mode"] = "thread_local"
|
||||
with torch.cuda.graph(self._generate_cuda_graph, **recording_kwargs):
|
||||
self._generate_logits = self.decode_model.forward_with_attn_bias(
|
||||
token_values=self._generate_inputs[0],
|
||||
attn_bias=self._generate_inputs[1],
|
||||
cache=self._cache,
|
||||
)
|
||||
|
||||
def replay(tokens, seq_lens):
|
||||
self._generate_inputs[0].copy_(tokens)
|
||||
self._generate_inputs[1].k_seqinfo.seqlen.copy_(seq_lens)
|
||||
|
||||
self._generate_cuda_graph.replay()
|
||||
|
||||
return self._generate_logits
|
||||
|
||||
return replay
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate_all(
|
||||
self, prompts: list[list[int]], use_cuda_graphs: bool, use_sampling: bool
|
||||
) -> Tuple[Stats, list[list[int]]]:
|
||||
bs = len(prompts)
|
||||
prompt_lens = [len(p) for p in prompts]
|
||||
padded_prompt_lens = [self.gen_args.prompt_length] * bs
|
||||
max_prompt_length = max(prompt_lens)
|
||||
gen_length = self.gen_args.gen_length
|
||||
max_seq_length = max_prompt_length + gen_length
|
||||
print(max_prompt_length, gen_length)
|
||||
|
||||
bias = AttnBias.from_seqlens(
|
||||
q_seqlen=padded_prompt_lens,
|
||||
kv_seqlen=prompt_lens,
|
||||
kv_padding=max_seq_length,
|
||||
)
|
||||
bias.q_seqinfo.to("cuda")
|
||||
bias.k_seqinfo.to("cuda")
|
||||
|
||||
# Input tensors to the cuda graph
|
||||
kv_seqlen = bias.k_seqinfo.seqlen
|
||||
prompts = [prompt + [1] * (self.gen_args.prompt_length - len(prompt)) for prompt in prompts]
|
||||
tokens = torch.IntTensor(sum(prompts, [])).cuda()
|
||||
out_tokens = torch.zeros((max_seq_length, bs), dtype=torch.int)
|
||||
|
||||
stats = Stats()
|
||||
torch.cuda.synchronize()
|
||||
stats.phase("prefill" if use_cuda_graphs else "total")
|
||||
# stats.phase("total")
|
||||
|
||||
output = self._prefill_compile_model(tokens, None)
|
||||
|
||||
logits = output[kv_seqlen - 1, :]
|
||||
logits = logits.view(bs, self.model_args.vocab_size)
|
||||
|
||||
if use_sampling:
|
||||
temp = 0.7
|
||||
top_p = 0.95
|
||||
probs = torch.softmax(logits / temp, dim=-1)
|
||||
next_token = sample_utils.top_p(probs, top_p)
|
||||
else:
|
||||
next_token = torch.argmax(logits, dim=-1)
|
||||
|
||||
next_token = next_token.reshape(bs)
|
||||
out_tokens[0, :] = next_token
|
||||
|
||||
torch.cuda.synchronize()
|
||||
stats.phase("decode" if use_cuda_graphs else "total")
|
||||
|
||||
eos_id = self.tokenizer.eot_id
|
||||
for niter in range(1, gen_length):
|
||||
kv_seqlen.add_(kv_seqlen < max_seq_length)
|
||||
output = self._generate_compile_model(next_token, kv_seqlen)
|
||||
|
||||
logits = output.view(bs, self.model_args.vocab_size)
|
||||
|
||||
if use_sampling:
|
||||
temp = 0.7
|
||||
top_p = 0.95
|
||||
probs = torch.softmax(logits / temp, dim=-1)
|
||||
next_token = sample_utils.top_p(probs, top_p)
|
||||
else:
|
||||
next_token = torch.argmax(logits, dim=-1)
|
||||
|
||||
next_token = next_token.reshape(bs)
|
||||
out_tokens[niter, :] = next_token
|
||||
|
||||
if next_token.eq(eos_id).any():
|
||||
break
|
||||
|
||||
torch.cuda.synchronize()
|
||||
stats.end_phase(tokens=niter * bs)
|
||||
|
||||
def trim_answer(prompt_len, tokens):
|
||||
# print(prompt, tokens)
|
||||
"""Trim the answer to end it on an eos token."""
|
||||
tokens = tokens[: max_seq_length - prompt_len]
|
||||
eos_id = self.tokenizer.eot_id
|
||||
if eos_id in tokens:
|
||||
return tokens[: tokens.index(eos_id) + 1]
|
||||
else:
|
||||
return tokens
|
||||
|
||||
answers = [
|
||||
trim_answer(prompt_len, answer)
|
||||
for prompt_len, answer in zip(prompt_lens, out_tokens.t().tolist())
|
||||
]
|
||||
return stats, answers
|
||||
|
||||
|
||||
def get_prompts(interactive: bool) -> Iterable[list[str]]:
|
||||
if interactive:
|
||||
while True:
|
||||
try:
|
||||
prompts = input("enter prompt: ").split("\n")
|
||||
except EOFError:
|
||||
print("exiting")
|
||||
sys.exit(0)
|
||||
yield prompts
|
||||
else:
|
||||
yield [
|
||||
"Hello, my name is",
|
||||
]
|
||||
|
||||
|
||||
def main(ckpt_dir: str, interactive: bool = False, chat_format: bool = False, sampling: bool = False):
|
||||
|
||||
local_rank = 0
|
||||
device = f"cuda:{local_rank}"
|
||||
torch.cuda.set_device(local_rank)
|
||||
|
||||
g = FastGen.build(ckpt_dir, GenArgs(), device)
|
||||
|
||||
if chat_format:
|
||||
g.tokenizer = ChatFormat(g.tokenizer)
|
||||
|
||||
for prompts in get_prompts(interactive):
|
||||
# prompts = [f"{prompt}\n" for prompt in prompts]
|
||||
if chat_format:
|
||||
# prompts = [f'<|begin_of_text|>User: {prompt}<|eot_id|>Assistant: ' for prompt in prompts]
|
||||
tokens = [g.tokenizer.encode_dialog_prompt(dialog=[{"role": "user", "content": prompt}], completion=True) for prompt in prompts]
|
||||
else:
|
||||
tokens = [g.tokenizer.encode(x, bos=False, eos=False) for x in prompts]
|
||||
|
||||
print(tokens)
|
||||
stats, out_tokens = g.generate_all(
|
||||
tokens, use_cuda_graphs="NO_CUDA_GRAPHS" not in os.environ, use_sampling=sampling,
|
||||
)
|
||||
|
||||
for i, prompt in enumerate(prompts):
|
||||
print(f"> {prompt}")
|
||||
answer = g.tokenizer.decode(out_tokens[i])
|
||||
print(answer)
|
||||
print("---------------")
|
||||
|
||||
for phase_stats in stats.phases:
|
||||
print(phase_stats.show())
|
||||
|
||||
print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
fire.Fire(main)
|
||||
-366
@@ -1,366 +0,0 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from xformers.ops import RMSNorm, fmha, rope_padded
|
||||
from xformers.ops.fmha.attn_bias import (
|
||||
BlockDiagonalCausalWithOffsetPaddedKeysMask as AttnBias,
|
||||
)
|
||||
|
||||
import ctypes
|
||||
bitnet_lib = ctypes.CDLL('bitnet_kernels/libbitnet.so')
|
||||
|
||||
def bitnet_int8xint2_linear(input0, input1, s, ws):
|
||||
out_shape = list(input0.shape)
|
||||
out_shape[-1] = input1.shape[0]
|
||||
|
||||
stream = torch.cuda.current_stream()
|
||||
|
||||
M = input0.shape[0]
|
||||
if len(out_shape) == 3:
|
||||
M *= input0.shape[1]
|
||||
N = input1.shape[0]
|
||||
K = input1.shape[1] * 4
|
||||
|
||||
ret = torch.zeros(*out_shape, dtype=torch.bfloat16, device=input0.device)
|
||||
|
||||
bitnet_lib.bitlinear_int8xint2(*[ctypes.c_void_p(input0.data_ptr()), ctypes.c_void_p(input1.data_ptr()), ctypes.c_void_p(ret.data_ptr()), ctypes.c_void_p(s.data_ptr()), ctypes.c_void_p(ws.data_ptr()), ctypes.c_int(M), ctypes.c_int(N), ctypes.c_int(K), ctypes.c_void_p(stream.cuda_stream)])
|
||||
|
||||
return ret
|
||||
|
||||
@dataclass
|
||||
class ModelArgs:
|
||||
dim: int = 2560
|
||||
n_layers: int = 30
|
||||
n_heads: int = 20
|
||||
n_kv_heads: int = 5
|
||||
vocab_size: int = 128256
|
||||
ffn_dim: int = 6912
|
||||
norm_eps: float = 1e-5
|
||||
rope_theta: float = 500000.0
|
||||
use_kernel: bool = False
|
||||
|
||||
|
||||
LayerCache = Tuple[torch.Tensor, torch.Tensor]
|
||||
|
||||
class BitLinearKernel(nn.Module):
|
||||
in_features: int
|
||||
out_features: int
|
||||
weight: torch.Tensor
|
||||
weight_scale: torch.Tensor
|
||||
|
||||
def __init__(self, in_features: int, out_features: int, bias: bool = False):
|
||||
super().__init__()
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
|
||||
self.weight = torch.nn.Parameter(torch.zeros(out_features, in_features//4, dtype=torch.int8), requires_grad=False)
|
||||
self.weight_scale = torch.nn.Parameter(torch.zeros(4, dtype=torch.bfloat16), requires_grad=False)
|
||||
|
||||
@torch.compile
|
||||
def quant_input(self, input):
|
||||
s = 127 / input.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5)
|
||||
return (input * s).round().clamp(-128, 127).to(torch.int8), s
|
||||
|
||||
def forward(self, input):
|
||||
input, s = self.quant_input(input)
|
||||
return bitnet_int8xint2_linear(input, self.weight, s, self.weight_scale)
|
||||
|
||||
class BitLinear(nn.Linear):
|
||||
@torch.compile
|
||||
def quant_input(self, input):
|
||||
s = 127 / input.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5)
|
||||
return (input * s).round().clamp(-128, 127) / s
|
||||
|
||||
def forward(self, input):
|
||||
input = self.quant_input(input)
|
||||
return F.linear(input, self.weight)
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
head_dim: int,
|
||||
n_heads: int,
|
||||
n_kv_heads: int,
|
||||
rope_theta: float,
|
||||
norm_eps: float,
|
||||
use_kernel: bool,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.head_dim = head_dim
|
||||
self.rope_theta = rope_theta
|
||||
|
||||
self.n_local_heads = n_heads
|
||||
self.n_local_kv_heads = n_kv_heads
|
||||
|
||||
Linear = BitLinearKernel if use_kernel else BitLinear
|
||||
|
||||
self.wqkv = Linear(
|
||||
dim,
|
||||
(self.n_local_heads + 2 * self.n_local_kv_heads) * head_dim,
|
||||
bias=False,
|
||||
)
|
||||
self.wo = Linear(
|
||||
self.n_local_heads * head_dim,
|
||||
dim,
|
||||
bias=False,
|
||||
)
|
||||
|
||||
self.attn_sub_norm = RMSNorm(dim, norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
cache: LayerCache,
|
||||
attn_bias: AttnBias,
|
||||
) -> torch.Tensor:
|
||||
|
||||
xqkv = self.wqkv(x)
|
||||
xq = xqkv[:, : (self.n_local_heads * self.head_dim)]
|
||||
xkv = xqkv[:, (self.n_local_heads * self.head_dim) :]
|
||||
xk, xv = xkv.chunk(2, 1)
|
||||
|
||||
output_shape = xq.shape
|
||||
heads_per_group = self.n_local_heads // self.n_local_kv_heads
|
||||
xq = xq.view(
|
||||
1, xq.shape[0], self.n_local_kv_heads, heads_per_group, self.head_dim
|
||||
)
|
||||
xk = xk.view(1, xk.shape[0], self.n_local_kv_heads, 1, self.head_dim)
|
||||
# xq = rearrange(xq, 'b (g h l d) -> 1 b h g (d l)', g=heads_per_group, h=self.n_local_kv_heads, d=self.head_dim // 2, l=2)
|
||||
# xk = rearrange(xk, 'b (g l d) -> 1 b g 1 (d l)', g=self.n_local_kv_heads, d=self.head_dim // 2)
|
||||
xv = xv.view(1, xv.shape[0], self.n_local_kv_heads, 1, self.head_dim)
|
||||
cache_k, cache_v = cache
|
||||
|
||||
xq = rope_padded(
|
||||
xq=xq,
|
||||
xk=xk,
|
||||
xv=xv,
|
||||
cache_k=cache_k,
|
||||
cache_v=cache_v,
|
||||
attn_bias=attn_bias,
|
||||
theta=self.rope_theta,
|
||||
)
|
||||
|
||||
output = fmha.memory_efficient_attention_forward(
|
||||
xq, cache_k, cache_v, attn_bias, op = fmha.flash.FwOp
|
||||
)
|
||||
|
||||
output = output.reshape(output_shape)
|
||||
output = self.attn_sub_norm(output)
|
||||
output = self.wo(output)
|
||||
|
||||
return output
|
||||
|
||||
@torch.compile
|
||||
def squared_relu(x: torch.Tensor) -> torch.Tensor:
|
||||
return F.relu(x) ** 2
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
hidden_dim: int,
|
||||
norm_eps: float,
|
||||
use_kernel: bool,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
Linear = BitLinearKernel if use_kernel else BitLinear
|
||||
|
||||
self.w13 = Linear(
|
||||
dim,
|
||||
2 * hidden_dim,
|
||||
bias=False,
|
||||
)
|
||||
self.w2 = Linear(
|
||||
hidden_dim,
|
||||
dim,
|
||||
bias=False,
|
||||
)
|
||||
self.ffn_sub_norm = RMSNorm(hidden_dim, norm_eps)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x13 = self.w13(x)
|
||||
x1, x3 = x13.chunk(2, -1)
|
||||
inner = self.ffn_sub_norm(squared_relu(x1) * x3)
|
||||
output = self.w2(inner)
|
||||
return output
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
|
||||
assert args.dim % args.n_heads == 0
|
||||
head_dim = args.dim // args.n_heads
|
||||
if args.n_kv_heads is not None:
|
||||
n_kv_heads = args.n_kv_heads
|
||||
else:
|
||||
n_kv_heads = args.n_heads
|
||||
|
||||
assert args.n_heads % n_kv_heads == 0
|
||||
|
||||
self.attention = Attention(
|
||||
dim=args.dim,
|
||||
head_dim=head_dim,
|
||||
n_heads=args.n_heads,
|
||||
n_kv_heads=n_kv_heads,
|
||||
rope_theta=args.rope_theta,
|
||||
norm_eps=args.norm_eps,
|
||||
use_kernel=args.use_kernel,
|
||||
)
|
||||
self.feed_forward = FeedForward(
|
||||
dim=args.dim,
|
||||
hidden_dim=args.ffn_dim,
|
||||
norm_eps=args.norm_eps,
|
||||
use_kernel=args.use_kernel,
|
||||
)
|
||||
self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
||||
self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
cache: LayerCache,
|
||||
attn_bias: AttnBias,
|
||||
) -> torch.Tensor:
|
||||
h = x + self.attention.forward(
|
||||
self.attention_norm(x),
|
||||
cache,
|
||||
attn_bias,
|
||||
)
|
||||
out = h + self.feed_forward(self.ffn_norm(h))
|
||||
return out
|
||||
|
||||
|
||||
class Transformer(nn.Module):
|
||||
def __init__(self, args: ModelArgs):
|
||||
super().__init__()
|
||||
assert args.vocab_size > 0
|
||||
|
||||
self.tok_embeddings = nn.Embedding(
|
||||
num_embeddings=args.vocab_size,
|
||||
embedding_dim=args.dim,
|
||||
)
|
||||
|
||||
self.layers = nn.ModuleList()
|
||||
for _ in range(args.n_layers):
|
||||
self.layers.append(TransformerBlock(args))
|
||||
|
||||
self.norm = RMSNorm(args.dim, eps=args.norm_eps)
|
||||
|
||||
self.output = nn.Linear(
|
||||
args.dim,
|
||||
args.vocab_size,
|
||||
bias=False,
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def forward_with_attn_bias(
|
||||
self,
|
||||
token_values: torch.Tensor,
|
||||
attn_bias: AttnBias,
|
||||
cache: list[LayerCache],
|
||||
) -> torch.Tensor:
|
||||
h = self.tok_embeddings(token_values)
|
||||
|
||||
for i, layer in enumerate(self.layers):
|
||||
h = layer(h, cache[i], attn_bias)
|
||||
|
||||
logits = self.output(self.norm(h))
|
||||
return logits.float()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
token_values: torch.Tensor,
|
||||
token_lengths: torch.Tensor,
|
||||
start_pos: torch.Tensor,
|
||||
cache: list[LayerCache],
|
||||
kv_padding: int,
|
||||
) -> torch.Tensor:
|
||||
attn_bias = AttnBias.from_seqlens(
|
||||
q_seqlen=token_lengths.tolist(),
|
||||
kv_seqlen=(start_pos + token_lengths).tolist(),
|
||||
kv_padding=kv_padding,
|
||||
)
|
||||
return self.forward_with_attn_bias(token_values, attn_bias, cache)
|
||||
|
||||
|
||||
def make_cache(
|
||||
args: ModelArgs,
|
||||
length: int,
|
||||
device: Optional[Union[str, torch.device]] = None,
|
||||
n_layers: Optional[int] = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
) -> list[LayerCache]:
|
||||
"""
|
||||
Allocate a cache to be used with the Transformer module.
|
||||
|
||||
Args:
|
||||
args (ModelArgs): the model configuration.
|
||||
length (int): per layer cache size.
|
||||
It is usually budgeted as ``max_batch * max_seq``
|
||||
device (torch.device, optional): the device on which
|
||||
the cache should be allocated.
|
||||
n_layers (int, optional): the number of layers to
|
||||
allocate a cache for (defaults to the model
|
||||
settings).
|
||||
dtype (torch.dtype, optional): the dtype to use for
|
||||
cache entries (defaults to the default dtype).
|
||||
|
||||
Returns:
|
||||
The cache object to pass to ``Tranformer.forward``.
|
||||
"""
|
||||
|
||||
head_dim = args.dim // args.n_heads
|
||||
n_kv_heads = args.n_kv_heads
|
||||
if n_kv_heads is None:
|
||||
n_kv_heads = args.n_heads
|
||||
n_local_kv_heads = n_kv_heads
|
||||
|
||||
if n_layers is None:
|
||||
n_layers = args.n_layers
|
||||
|
||||
shape = (1, length, n_local_kv_heads, 1, head_dim)
|
||||
heads_per_group = args.n_heads // n_kv_heads
|
||||
expansion = (-1, -1, -1, heads_per_group, -1)
|
||||
return [
|
||||
(
|
||||
torch.zeros(shape, device=device, dtype=dtype).expand(expansion),
|
||||
torch.zeros(shape, device=device, dtype=dtype).expand(expansion),
|
||||
)
|
||||
for _ in range(n_layers)
|
||||
]
|
||||
|
||||
|
||||
def cache_prefix(cache: list[LayerCache], length: int) -> list[LayerCache]:
|
||||
"""
|
||||
Take a prefix view of a larger cache.
|
||||
|
||||
The original cache object remains of identical size and valid
|
||||
after the shrinked alias has been used. This function is useful
|
||||
when a cache was allocated for a larger batch size than what is
|
||||
necessary.
|
||||
|
||||
Args:
|
||||
cache: the cache to take a view in.
|
||||
length (int): the desired length
|
||||
|
||||
Returns:
|
||||
A view in the input cache object.
|
||||
"""
|
||||
|
||||
if len(cache) > 0:
|
||||
assert cache[0][0].shape[1] >= length
|
||||
|
||||
return [(ck[:, :length], cv[:, :length]) for ck, cv in cache]
|
||||
@@ -1,98 +0,0 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def B_global_16x32_to_shared_load_16x32_layout(i, j):
|
||||
"""
|
||||
stride * 8 * (tx // HALF_WARP_expr)
|
||||
+ (tx % 8) * stride
|
||||
+ 16 * ((tx % HALF_WARP_expr) // 8)
|
||||
"""
|
||||
thread_id = i * 2 + j // 16
|
||||
row = (thread_id // 16) * 8 + (thread_id % 8)
|
||||
col = (j % 16) + 16 * ((thread_id % 16) // 8)
|
||||
return row, col
|
||||
|
||||
|
||||
def permutate_weight_fastest(weight):
|
||||
wmma_n = 16
|
||||
wmma_k = 32
|
||||
N = weight.shape[0]
|
||||
K = weight.shape[1]
|
||||
|
||||
# Create a lookup table for the permutation
|
||||
mapping = np.zeros((wmma_n, wmma_k, 2), dtype=int)
|
||||
for ii in range(wmma_n):
|
||||
for jj in range(wmma_k):
|
||||
mapping[ii, jj] = B_global_16x32_to_shared_load_16x32_layout(ii, jj)
|
||||
|
||||
# Reshape weight for the final format
|
||||
permutated_weight = np.zeros((N // wmma_n, K // wmma_k, wmma_n, wmma_k), dtype="int8")
|
||||
|
||||
# Use advanced indexing for the entire operation
|
||||
i_indices = np.arange(N // wmma_n)[:, np.newaxis, np.newaxis, np.newaxis]
|
||||
j_indices = np.arange(K // wmma_k)[np.newaxis, :, np.newaxis, np.newaxis]
|
||||
|
||||
# Create the source indices
|
||||
src_i = i_indices * wmma_n + mapping[:, :, 0]
|
||||
src_j = j_indices * wmma_k + mapping[:, :, 1]
|
||||
|
||||
# Extract and reshape in one go
|
||||
permutated_weight = weight[src_i, src_j]
|
||||
|
||||
return permutated_weight
|
||||
|
||||
|
||||
def compress_int2_to_int8(int2_weight):
|
||||
int8_weight = np.zeros(
|
||||
(*int2_weight.shape[:-1], int2_weight.shape[-1] // 4), dtype=np.int8
|
||||
)
|
||||
for j in range(int2_weight.shape[-1] // 4):
|
||||
for k in range(4):
|
||||
int8_weight[:, :, :, j] |= int2_weight[:, :, :, j * 4 + k] << (k * 2)
|
||||
return int8_weight
|
||||
|
||||
|
||||
def interleave_weight_int8(qweight, nbits=2):\
|
||||
# reinterpret the data type of qweight to int32
|
||||
# shift = [ 0, 8, 16, 24, 2, 10, 18, 26, 4, 12, 20, 28, 6, 14, 22, 30]
|
||||
# index: [ 0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15]
|
||||
qweight = qweight.view(np.int32)
|
||||
new_qweight = np.zeros_like(qweight)
|
||||
bits_stride = 8
|
||||
mask = (1 << nbits) - 1 # for 4bit the val is 0x0000000f
|
||||
num_groups = 32 // bits_stride # 4
|
||||
elems_per_group = bits_stride // nbits # 4
|
||||
for i in range(num_groups):
|
||||
for j in range(elems_per_group):
|
||||
offset = i * elems_per_group + j
|
||||
shift = (offset % num_groups) * bits_stride + (offset // num_groups) * nbits
|
||||
|
||||
new_qweight |= ((qweight >> (nbits * offset)) & mask) << shift
|
||||
return new_qweight.view(np.int8)
|
||||
|
||||
|
||||
|
||||
def convert_weight_int8_to_int2(weight):
|
||||
N = weight.shape[0]
|
||||
K = weight.shape[1]
|
||||
|
||||
weight = weight+2
|
||||
|
||||
weight = weight.cpu().numpy()
|
||||
|
||||
# print(weight)
|
||||
# print(torch.max(weight), torch.min(weight))
|
||||
|
||||
# permutated_weight_slow = permutate_weight(weight)
|
||||
permutated_weight = permutate_weight_fastest(weight)
|
||||
# assert np.all(permutated_weight_slow == permutated_weight)
|
||||
# print("Permutation is correct")
|
||||
compressed_weight = compress_int2_to_int8(permutated_weight)
|
||||
interleaved_weight = interleave_weight_int8(compressed_weight, 2)
|
||||
|
||||
ret = torch.from_numpy(interleaved_weight)
|
||||
|
||||
ret = torch.reshape(ret, (N, K // 4))
|
||||
|
||||
return ret
|
||||
@@ -1,9 +0,0 @@
|
||||
fire
|
||||
sentencepiece
|
||||
torch>=2.2.0
|
||||
xformers>=0.0.22
|
||||
tiktoken
|
||||
blobfile
|
||||
flask
|
||||
einops
|
||||
transformers
|
||||
@@ -1,31 +0,0 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
|
||||
@torch.compile
|
||||
def top_p(probs: torch.Tensor, p: float) -> torch.Tensor:
|
||||
"""
|
||||
Perform top-p (nucleus) sampling on a probability distribution.
|
||||
|
||||
Args:
|
||||
probs (torch.Tensor): probability distribution tensor.
|
||||
p (float): probability threshold for top-p sampling.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: sampled token indices.
|
||||
|
||||
Note:
|
||||
Top-p sampling selects the smallest set of tokens whose cumulative
|
||||
probability mass exceeds the threshold p. The distribution is
|
||||
renormalized based on the selected tokens.
|
||||
"""
|
||||
probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
|
||||
probs_sum = torch.cumsum(probs_sort, dim=-1)
|
||||
mask = probs_sum - probs_sort > p
|
||||
probs_sort[mask] = 0.0
|
||||
next_token = torch.multinomial(probs_sort, num_samples=1)
|
||||
next_token = torch.gather(probs_idx, -1, next_token)
|
||||
return next_token
|
||||
@@ -1,57 +0,0 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class PhaseStats:
|
||||
name: str
|
||||
tokens: int
|
||||
time: float
|
||||
|
||||
def show(self) -> str:
|
||||
tps = self.tokens / self.time
|
||||
return (
|
||||
f"[{self.name}] "
|
||||
f"generated tokens: {self.tokens}"
|
||||
f" - total time: {self.time:.3f}s"
|
||||
f" - {tps:.1f} tokens per second"
|
||||
)
|
||||
|
||||
|
||||
class Stats:
|
||||
"""
|
||||
Generation stats, split by phases.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.phases = []
|
||||
self.current = None
|
||||
|
||||
def end_phase(self, tokens: int, now: Optional[float] = None):
|
||||
"""Terminate the current phase."""
|
||||
if self.current is None:
|
||||
return
|
||||
if now is None:
|
||||
now = time.time()
|
||||
cname, ctokens, ctime = self.current
|
||||
stats = PhaseStats(
|
||||
name=cname,
|
||||
tokens=tokens - ctokens,
|
||||
time=now - ctime,
|
||||
)
|
||||
self.phases.append(stats)
|
||||
|
||||
def phase(self, name: str, tokens: int = 0):
|
||||
"""
|
||||
Start a new phase, and terminate the current one,
|
||||
if one is ongoing.
|
||||
"""
|
||||
now = time.time()
|
||||
self.end_phase(tokens, now)
|
||||
self.current = (name, tokens, now)
|
||||
-99
@@ -1,99 +0,0 @@
|
||||
import torch
|
||||
from torch.utils import benchmark
|
||||
from torch import nn
|
||||
|
||||
from pack_weight import convert_weight_int8_to_int2
|
||||
from torch.profiler import profile, record_function, ProfilerActivity
|
||||
import ctypes
|
||||
import numpy as np
|
||||
# set all seed
|
||||
torch.manual_seed(42)
|
||||
np.random.seed(42)
|
||||
|
||||
bitnet_lib = ctypes.CDLL('bitnet_kernels/libbitnet.so')
|
||||
|
||||
def bitnet_int8xint2_linear(input0, input1, s, ws, ret):
|
||||
out_shape = list(input0.shape)
|
||||
out_shape[-1] = input1.shape[0]
|
||||
|
||||
stream = torch.cuda.current_stream()
|
||||
|
||||
M = input0.shape[0]
|
||||
if len(out_shape) == 3:
|
||||
M *= input0.shape[1]
|
||||
N = input1.shape[0]
|
||||
K = input1.shape[1] * 4
|
||||
|
||||
bitnet_lib.bitlinear_int8xint2(*[ctypes.c_void_p(input0.data_ptr()), ctypes.c_void_p(input1.data_ptr()), ctypes.c_void_p(ret.data_ptr()), ctypes.c_void_p(s.data_ptr()), ctypes.c_void_p(ws.data_ptr()), ctypes.c_int(M), ctypes.c_int(N), ctypes.c_int(K), ctypes.c_void_p(stream.cuda_stream)])
|
||||
|
||||
return ret
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_list = [
|
||||
(2560, 2560),
|
||||
(3840, 2560),
|
||||
(13824, 2560),
|
||||
(2560, 6912) ,
|
||||
(3200, 3200),
|
||||
(4800, 3200),
|
||||
(3200, 10240),
|
||||
(20480, 3200),
|
||||
]
|
||||
for N,K in test_list:
|
||||
weight = torch.randint(-1, 2, (N, K), dtype=torch.int8, device='cuda')
|
||||
weight_scale = torch.ones(1, dtype=torch.bfloat16, device='cuda')
|
||||
weight_compressed = convert_weight_int8_to_int2(weight).to('cuda')
|
||||
|
||||
for i in range(1):
|
||||
input0 = torch.randint(-128,127,(1, K),dtype=torch.int8, device='cuda')
|
||||
input0_bf16 = input0.to(torch.bfloat16)
|
||||
input_np = input0.cpu().to(torch.int32).numpy()
|
||||
weight_np = weight.cpu().to(torch.int32).T.numpy()
|
||||
out_np = np.matmul(input_np,weight_np)
|
||||
out_np = torch.tensor(out_np).cuda().to(torch.bfloat16)
|
||||
|
||||
s = torch.ones(1, dtype=torch.bfloat16, device='cuda')
|
||||
ws = torch.ones(6, dtype=torch.bfloat16, device='cuda')
|
||||
|
||||
ret = torch.empty((1,N), dtype=torch.bfloat16, device=input0.device)
|
||||
out = bitnet_int8xint2_linear(input0, weight_compressed, s, ws, ret)
|
||||
|
||||
print(f'custom == np {torch.all(out==out_np)}')
|
||||
|
||||
input0 = torch.randint(-128,127,(1, K),dtype=torch.int8, device='cuda')
|
||||
input0_fp16 = input0.to(torch.float16)
|
||||
input0_bf16 = input0.to(torch.bfloat16)
|
||||
weight_fp16 = weight.to(torch.float16).T
|
||||
weight_bf16 = weight.to(torch.bfloat16).T
|
||||
ret = torch.empty((1,N), dtype=torch.bfloat16, device=input0.device)
|
||||
s = torch.ones(1, dtype=torch.bfloat16, device='cuda')
|
||||
ws = torch.ones(6, dtype=torch.bfloat16, device='cuda')
|
||||
t0 = benchmark.Timer(
|
||||
stmt="bitnet_int8xint2_linear(input0, weight_compressed, s, ws, ret)",
|
||||
setup="from __main__ import input0, weight_compressed, s, ws, ret, bitnet_int8xint2_linear",
|
||||
num_threads=1,
|
||||
)
|
||||
|
||||
t1 = benchmark.Timer(
|
||||
stmt="torch.matmul(input0_bf16,weight_bf16)",
|
||||
setup="from __main__ import input0_bf16, weight_bf16",
|
||||
num_threads=1,
|
||||
)
|
||||
|
||||
time0 = t0.timeit(50)
|
||||
time1 = t1.timeit(50)
|
||||
|
||||
print(f'Shape{N,K}, W2A8: {time0.mean * 1e6:.2f}us, torch BF16: {time1.mean * 1e6:.2f}us')
|
||||
# activities = [ ProfilerActivity.CUDA,
|
||||
# # ProfilerActivity.CPU
|
||||
# ]
|
||||
# sort_by_keyword = 'cuda' + "_time_total"
|
||||
# with profile(activities=activities, record_shapes=True) as prof:
|
||||
# with record_function("model_inference1"):
|
||||
# for _ in range(10):
|
||||
# bitnet_int8xint2_linear(input0, weight_compressed, s, ws, ret)
|
||||
# torch.matmul(input0_fp16,weight_fp16)
|
||||
# torch.matmul(input0_bf16,weight_bf16)
|
||||
|
||||
# print(prof.key_averages().table(sort_by=sort_by_keyword, row_limit=15))
|
||||
|
||||
-128000
File diff suppressed because it is too large
Load Diff
@@ -1,257 +0,0 @@
|
||||
import os
|
||||
from logging import getLogger
|
||||
from pathlib import Path
|
||||
from typing import (
|
||||
AbstractSet,
|
||||
cast,
|
||||
Collection,
|
||||
Dict,
|
||||
Iterator,
|
||||
List,
|
||||
Literal,
|
||||
Sequence,
|
||||
TypedDict,
|
||||
Union,
|
||||
)
|
||||
|
||||
import tiktoken
|
||||
from tiktoken.load import load_tiktoken_bpe
|
||||
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
Role = Literal["system", "user", "assistant"]
|
||||
|
||||
|
||||
class Message(TypedDict):
|
||||
role: Role
|
||||
content: str
|
||||
|
||||
|
||||
Dialog = Sequence[Message]
|
||||
|
||||
|
||||
class Tokenizer:
|
||||
"""
|
||||
Tokenizing and encoding/decoding text using the Tiktoken tokenizer.
|
||||
"""
|
||||
|
||||
special_tokens: Dict[str, int]
|
||||
|
||||
num_reserved_special_tokens = 256
|
||||
|
||||
pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+" # noqa: E501
|
||||
|
||||
def __init__(self, model_path: str):
|
||||
"""
|
||||
Initializes the Tokenizer with a Tiktoken model.
|
||||
|
||||
Args:
|
||||
model_path (str): The path to the Tiktoken model file.
|
||||
"""
|
||||
assert os.path.isfile(model_path), model_path
|
||||
|
||||
mergeable_ranks = load_tiktoken_bpe(model_path)
|
||||
num_base_tokens = len(mergeable_ranks)
|
||||
special_tokens = [
|
||||
"<|begin_of_text|>",
|
||||
"<|end_of_text|>",
|
||||
"<|reserved_special_token_0|>",
|
||||
"<|reserved_special_token_1|>",
|
||||
"<|reserved_special_token_2|>",
|
||||
"<|reserved_special_token_3|>",
|
||||
"<|start_header_id|>",
|
||||
"<|end_header_id|>",
|
||||
"<|reserved_special_token_4|>",
|
||||
"<|eot_id|>", # end of turn
|
||||
] + [
|
||||
f"<|reserved_special_token_{i}|>"
|
||||
for i in range(5, self.num_reserved_special_tokens - 5)
|
||||
]
|
||||
self.special_tokens = {
|
||||
token: num_base_tokens + i for i, token in enumerate(special_tokens)
|
||||
}
|
||||
self.model = tiktoken.Encoding(
|
||||
name=Path(model_path).name,
|
||||
pat_str=self.pat_str,
|
||||
mergeable_ranks=mergeable_ranks,
|
||||
special_tokens=self.special_tokens,
|
||||
)
|
||||
logger.info(f"Reloaded tiktoken model from {model_path}")
|
||||
|
||||
self.n_words: int = self.model.n_vocab
|
||||
# BOS / EOS token IDs
|
||||
self.bos_id: int = self.special_tokens["<|begin_of_text|>"]
|
||||
self.eos_id: int = self.special_tokens["<|end_of_text|>"]
|
||||
self.pad_id: int = self.n_words - 1
|
||||
self.stop_tokens = {
|
||||
self.special_tokens["<|end_of_text|>"],
|
||||
self.special_tokens["<|eot_id|>"],
|
||||
}
|
||||
logger.info(
|
||||
f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
|
||||
)
|
||||
|
||||
def encode(
|
||||
self,
|
||||
s: str,
|
||||
*,
|
||||
bos: bool,
|
||||
eos: bool,
|
||||
allowed_special: Union[Literal["all"], AbstractSet[str]] = set(),
|
||||
disallowed_special: Union[Literal["all"], Collection[str]] = (),
|
||||
) -> List[int]:
|
||||
"""
|
||||
Encodes a string into a list of token IDs.
|
||||
|
||||
Args:
|
||||
s (str): The input string to be encoded.
|
||||
bos (bool): Whether to prepend the beginning-of-sequence token.
|
||||
eos (bool): Whether to append the end-of-sequence token.
|
||||
allowed_tokens ("all"|set[str]): allowed special tokens in string
|
||||
disallowed_tokens ("all"|set[str]): special tokens that raise an error when in string
|
||||
|
||||
Returns:
|
||||
list[int]: A list of token IDs.
|
||||
|
||||
By default, setting disallowed_special=() encodes a string by ignoring
|
||||
special tokens. Specifically:
|
||||
- Setting `disallowed_special` to () will cause all text corresponding
|
||||
to special tokens to be encoded as natural text (insteading of raising
|
||||
an error).
|
||||
- Setting `allowed_special` to "all" will treat all text corresponding
|
||||
to special tokens to be encoded as special tokens.
|
||||
"""
|
||||
assert type(s) is str
|
||||
|
||||
# The tiktoken tokenizer can handle <=400k chars without
|
||||
# pyo3_runtime.PanicException.
|
||||
TIKTOKEN_MAX_ENCODE_CHARS = 400_000
|
||||
|
||||
# https://github.com/openai/tiktoken/issues/195
|
||||
# Here we iterate over subsequences and split if we exceed the limit
|
||||
# of max consecutive non-whitespace or whitespace characters.
|
||||
MAX_NO_WHITESPACES_CHARS = 25_000
|
||||
|
||||
substrs = (
|
||||
substr
|
||||
for i in range(0, len(s), TIKTOKEN_MAX_ENCODE_CHARS)
|
||||
for substr in self._split_whitespaces_or_nonwhitespaces(
|
||||
s[i : i + TIKTOKEN_MAX_ENCODE_CHARS], MAX_NO_WHITESPACES_CHARS
|
||||
)
|
||||
)
|
||||
t: List[int] = []
|
||||
for substr in substrs:
|
||||
t.extend(
|
||||
self.model.encode(
|
||||
substr,
|
||||
allowed_special=allowed_special,
|
||||
disallowed_special=disallowed_special,
|
||||
)
|
||||
)
|
||||
if bos:
|
||||
t.insert(0, self.bos_id)
|
||||
if eos:
|
||||
t.append(self.eos_id)
|
||||
return t
|
||||
|
||||
def decode(self, t: Sequence[int]) -> str:
|
||||
"""
|
||||
Decodes a list of token IDs into a string.
|
||||
|
||||
Args:
|
||||
t (List[int]): The list of token IDs to be decoded.
|
||||
|
||||
Returns:
|
||||
str: The decoded string.
|
||||
"""
|
||||
# Typecast is safe here. Tiktoken doesn't do anything list-related with the sequence.
|
||||
return self.model.decode(cast(List[int], t))
|
||||
|
||||
@staticmethod
|
||||
def _split_whitespaces_or_nonwhitespaces(
|
||||
s: str, max_consecutive_slice_len: int
|
||||
) -> Iterator[str]:
|
||||
"""
|
||||
Splits the string `s` so that each substring contains no more than `max_consecutive_slice_len`
|
||||
consecutive whitespaces or consecutive non-whitespaces.
|
||||
"""
|
||||
current_slice_len = 0
|
||||
current_slice_is_space = s[0].isspace() if len(s) > 0 else False
|
||||
slice_start = 0
|
||||
|
||||
for i in range(len(s)):
|
||||
is_now_space = s[i].isspace()
|
||||
|
||||
if current_slice_is_space ^ is_now_space:
|
||||
current_slice_len = 1
|
||||
current_slice_is_space = is_now_space
|
||||
else:
|
||||
current_slice_len += 1
|
||||
if current_slice_len > max_consecutive_slice_len:
|
||||
yield s[slice_start:i]
|
||||
slice_start = i
|
||||
current_slice_len = 1
|
||||
yield s[slice_start:]
|
||||
|
||||
class ChatFormat:
|
||||
def __init__(self, tokenizer: Tokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
self.eot_id = tokenizer.special_tokens["<|eot_id|>"]
|
||||
|
||||
def decode(self, tokens: List[int]) -> str:
|
||||
# Decode the tokens to a string.
|
||||
decoded_str = self.tokenizer.decode(tokens)
|
||||
# Remove the special tokens from the decoded string.
|
||||
decoded_str = decoded_str.replace("<|eot_id|>", "")
|
||||
return decoded_str
|
||||
|
||||
def encode_header(self, message: Message) -> List[int]:
|
||||
tokens = []
|
||||
if message["role"] == "system":
|
||||
tokens.extend(self.tokenizer.encode("System: ", bos=False, eos=False))
|
||||
elif message["role"] == "user":
|
||||
tokens.extend(self.tokenizer.encode("User: ", bos=False, eos=False))
|
||||
elif message["role"] == "assistant":
|
||||
tokens.extend(self.tokenizer.encode("Assistant: ", bos=False, eos=False))
|
||||
else:
|
||||
raise NotImplementedError(f"Role {message['role']} not implemented.")
|
||||
# tokens.append(self.tokenizer.special_tokens["<|start_header_id|>"])
|
||||
# tokens.extend(self.tokenizer.encode(message["role"], bos=False, eos=False))
|
||||
# tokens.append(self.tokenizer.special_tokens["<|end_header_id|>"])
|
||||
# tokens.extend(self.tokenizer.encode("\n\n", bos=False, eos=False))
|
||||
return tokens
|
||||
|
||||
def encode_message(self, message: Message, return_target=False) -> List[int]:
|
||||
tokens, targets = [], []
|
||||
headers = self.encode_header(message)
|
||||
contents = self.tokenizer.encode(message["content"].strip(), bos=False, eos=False)
|
||||
contents.append(self.tokenizer.special_tokens["<|eot_id|>"])
|
||||
tokens = headers + contents
|
||||
|
||||
if message["role"] == "assistant":
|
||||
targets = [-1] * len(headers) + contents
|
||||
else:
|
||||
targets = [-1] * len(tokens)
|
||||
|
||||
if return_target:
|
||||
return tokens, targets
|
||||
|
||||
return tokens, None
|
||||
|
||||
def encode_dialog_prompt(self, dialog: Dialog, completion=False, return_target=False) -> List[int]:
|
||||
tokens = [self.tokenizer.special_tokens["<|begin_of_text|>"]]
|
||||
targets = [-1]
|
||||
for message in dialog:
|
||||
_tokens, _targets = self.encode_message(message, return_target=return_target)
|
||||
tokens.extend(_tokens)
|
||||
if _targets is not None:
|
||||
targets.extend(_targets)
|
||||
# Add the start of an assistant message for the model to complete.
|
||||
if completion:
|
||||
tokens.extend(self.encode_header({"role": "assistant", "content": ""}))
|
||||
|
||||
if return_target:
|
||||
return tokens, targets
|
||||
|
||||
return tokens
|
||||
@@ -0,0 +1,627 @@
|
||||
#if defined(GGML_BITNET_ARM_TL1)
|
||||
#include "ggml-bitnet.h"
|
||||
#define GGML_BITNET_MAX_NODES 8192
|
||||
static bool initialized = false;
|
||||
static bitnet_tensor_extra * bitnet_tensor_extras = nullptr;
|
||||
static size_t bitnet_tensor_extras_index = 0;
|
||||
static void * aligned_malloc(size_t size) {{
|
||||
#if defined(_WIN32)
|
||||
return _aligned_malloc(size, 64);
|
||||
#else
|
||||
void * ptr = nullptr;
|
||||
posix_memalign(&ptr, 64, size);
|
||||
return ptr;
|
||||
#endif
|
||||
}}
|
||||
static void aligned_free(void * ptr) {{
|
||||
#if defined(_WIN32)
|
||||
_aligned_free(ptr);
|
||||
#else
|
||||
free(ptr);
|
||||
#endif
|
||||
}}
|
||||
|
||||
void per_tensor_quant(int k, void* lut_scales_, void* b_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
bitnet_float_type* b = (bitnet_float_type*)b_;
|
||||
#ifdef __ARM_NEON
|
||||
float32x4_t temp_max = vdupq_n_f32(0);
|
||||
for (int i=0; i < k / 4; i++) {{
|
||||
float32x4_t vec_bs = vld1q_f32(b + 4 * i);
|
||||
float32x4_t abssum = vabsq_f32(vec_bs);
|
||||
temp_max = vmaxq_f32(abssum, temp_max);
|
||||
}}
|
||||
float32_t scales = 127 / vmaxvq_f32(temp_max);
|
||||
*lut_scales = scales;
|
||||
#elif defined __AVX2__
|
||||
__m256 max_vec = _mm256_set1_ps(0.f);
|
||||
const __m256 vec_sign = _mm256_set1_ps(-0.0f);
|
||||
// #pragma unroll
|
||||
for (int i = 0; i < k / 8; i++) {{
|
||||
__m256 vec_b = _mm256_loadu_ps(b + i * 8);
|
||||
__m256 vec_babs = _mm256_andnot_ps(vec_sign, vec_b);
|
||||
max_vec = _mm256_max_ps(vec_babs, max_vec);
|
||||
}}
|
||||
__m128 max1 = _mm_max_ps(_mm256_extractf128_ps(max_vec, 1), _mm256_castps256_ps128(max_vec));
|
||||
max1 = _mm_max_ps(max1, _mm_movehl_ps(max1, max1));
|
||||
max1 = _mm_max_ss(max1, _mm_movehdup_ps(max1));
|
||||
float scales = 127 / _mm_cvtss_f32(max1);
|
||||
*lut_scales = scales;
|
||||
#endif
|
||||
}}
|
||||
|
||||
void partial_max_reset(void* lut_scales_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
*lut_scales = 0.0;
|
||||
}}
|
||||
|
||||
#ifdef __ARM_NEON
|
||||
inline void Transpose_8_8(
|
||||
int16x8_t *v0,
|
||||
int16x8_t *v1,
|
||||
int16x8_t *v2,
|
||||
int16x8_t *v3,
|
||||
int16x8_t *v4,
|
||||
int16x8_t *v5,
|
||||
int16x8_t *v6,
|
||||
int16x8_t *v7)
|
||||
{{
|
||||
int16x8x2_t q04 = vzipq_s16(*v0, *v4);
|
||||
int16x8x2_t q15 = vzipq_s16(*v1, *v5);
|
||||
int16x8x2_t q26 = vzipq_s16(*v2, *v6);
|
||||
int16x8x2_t q37 = vzipq_s16(*v3, *v7);
|
||||
|
||||
int16x8x2_t q0246_0 = vzipq_s16(q04.val[0], q26.val[0]);
|
||||
int16x8x2_t q0246_1 = vzipq_s16(q04.val[1], q26.val[1]);
|
||||
int16x8x2_t q1357_0 = vzipq_s16(q15.val[0], q37.val[0]);
|
||||
int16x8x2_t q1357_1 = vzipq_s16(q15.val[1], q37.val[1]);
|
||||
|
||||
int16x8x2_t q_fin_0 = vzipq_s16(q0246_0.val[0], q1357_0.val[0]);
|
||||
int16x8x2_t q_fin_1 = vzipq_s16(q0246_0.val[1], q1357_0.val[1]);
|
||||
int16x8x2_t q_fin_2 = vzipq_s16(q0246_1.val[0], q1357_1.val[0]);
|
||||
int16x8x2_t q_fin_3 = vzipq_s16(q0246_1.val[1], q1357_1.val[1]);
|
||||
|
||||
*v0 = q_fin_0.val[0];
|
||||
*v1 = q_fin_0.val[1];
|
||||
*v2 = q_fin_1.val[0];
|
||||
*v3 = q_fin_1.val[1];
|
||||
*v4 = q_fin_2.val[0];
|
||||
*v5 = q_fin_2.val[1];
|
||||
*v6 = q_fin_3.val[0];
|
||||
*v7 = q_fin_3.val[1];
|
||||
}}
|
||||
#endif
|
||||
|
||||
template<int act_k>
|
||||
inline void lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {{
|
||||
#ifdef __ARM_NEON
|
||||
int16x8_t vec_lut[16];
|
||||
float32_t scales = *lut_scales;
|
||||
uint8_t tbl_mask[16];
|
||||
tbl_mask[0] = 0;
|
||||
tbl_mask[1] = 2;
|
||||
tbl_mask[2] = 4;
|
||||
tbl_mask[3] = 6;
|
||||
tbl_mask[4] = 8;
|
||||
tbl_mask[5] = 10;
|
||||
tbl_mask[6] = 12;
|
||||
tbl_mask[7] = 14;
|
||||
tbl_mask[8] = 1;
|
||||
tbl_mask[9] = 3;
|
||||
tbl_mask[10] = 5;
|
||||
tbl_mask[11] = 7;
|
||||
tbl_mask[12] = 9;
|
||||
tbl_mask[13] = 11;
|
||||
tbl_mask[14] = 13;
|
||||
tbl_mask[15] = 15;
|
||||
uint8x16_t tbl_mask_q = vld1q_u8(tbl_mask);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < act_k / 16; ++k) {{
|
||||
float32x4x2_t vec_bs_x0 = vld2q_f32(b + k * 16);
|
||||
float32x4x2_t vec_bs_x1 = vld2q_f32(b + k * 16 + 8);
|
||||
float32x4_t vec_f_0 = vmulq_n_f32(vec_bs_x0.val[0], scales);
|
||||
float32x4_t vec_f_1 = vmulq_n_f32(vec_bs_x0.val[1], scales);
|
||||
float32x4_t vec_f_2 = vmulq_n_f32(vec_bs_x1.val[0], scales);
|
||||
float32x4_t vec_f_3 = vmulq_n_f32(vec_bs_x1.val[1], scales);
|
||||
int32x4_t vec_b_0 = vcvtnq_s32_f32(vec_f_0);
|
||||
int32x4_t vec_b_1 = vcvtnq_s32_f32(vec_f_1);
|
||||
int32x4_t vec_b_2 = vcvtnq_s32_f32(vec_f_2);
|
||||
int32x4_t vec_b_3 = vcvtnq_s32_f32(vec_f_3);
|
||||
int16x4_t vec_b16_0 = vmovn_s32(vec_b_0);
|
||||
int16x4_t vec_b16_1 = vmovn_s32(vec_b_1);
|
||||
int16x4_t vec_b16_2 = vmovn_s32(vec_b_2);
|
||||
int16x4_t vec_b16_3 = vmovn_s32(vec_b_3);
|
||||
int16x8_t vec_bs_0 = vcombine_s16(vec_b16_0, vec_b16_2);
|
||||
int16x8_t vec_bs_1 = vcombine_s16(vec_b16_1, vec_b16_3);
|
||||
vec_lut[0] = vdupq_n_s16(0);
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_0;
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_1;
|
||||
vec_lut[1] = vdupq_n_s16(0);
|
||||
vec_lut[1] = vec_lut[1] - vec_bs_0;
|
||||
vec_lut[2] = vdupq_n_s16(0);
|
||||
vec_lut[2] = vec_lut[2] - vec_bs_0;
|
||||
vec_lut[2] = vec_lut[2] + vec_bs_1;
|
||||
vec_lut[3] = vdupq_n_s16(0);
|
||||
vec_lut[3] = vec_lut[3] - vec_bs_1;
|
||||
vec_lut[4] = vdupq_n_s16(0);
|
||||
vec_lut[5] = vec_bs_1;
|
||||
vec_lut[6] = vec_bs_0;
|
||||
vec_lut[6] = vec_lut[6] - vec_bs_1;
|
||||
vec_lut[7] = vec_bs_0;
|
||||
vec_lut[8] = vec_bs_0;
|
||||
vec_lut[8] = vec_lut[8] + vec_bs_1;
|
||||
Transpose_8_8(&(vec_lut[0]), &(vec_lut[1]), &(vec_lut[2]), &(vec_lut[3]),
|
||||
&(vec_lut[4]), &(vec_lut[5]), &(vec_lut[6]), &(vec_lut[7]));
|
||||
Transpose_8_8(&(vec_lut[8]), &(vec_lut[9]), &(vec_lut[10]), &(vec_lut[11]),
|
||||
&(vec_lut[12]), &(vec_lut[13]), &(vec_lut[14]), &(vec_lut[15]));
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < 8; idx++) {{
|
||||
int8x16_t q0_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx]), tbl_mask_q);
|
||||
int8x8_t q0_low = vget_low_s8(q0_s);
|
||||
int8x8_t q0_high = vget_high_s8(q0_s);
|
||||
int8x16_t q1_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx + 8]), tbl_mask_q);
|
||||
int8x8_t q1_low = vget_low_s8(q1_s);
|
||||
int8x8_t q1_high = vget_high_s8(q1_s);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2, q0_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 8, q1_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 16, q0_low);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 24, q1_low);
|
||||
}}
|
||||
}}
|
||||
#endif
|
||||
}}
|
||||
|
||||
static bool is_type_supported(enum ggml_type type) {{
|
||||
if (type == GGML_TYPE_Q4_0 ||
|
||||
type == GGML_TYPE_TL1) {{
|
||||
return true;
|
||||
}} else {{
|
||||
return false;
|
||||
}}
|
||||
}}
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM1536_4096 256
|
||||
#define BBK1536_4096 128
|
||||
inline void tbl_impl_1536_4096(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK1536_4096 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1536_4096; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_1536_4096(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM1536_4096];
|
||||
memset(&(CBits[0]), 0, BM1536_4096 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 4096 / BBK1536_4096; ++k_outer) {
|
||||
tbl_impl_1536_4096((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK1536_4096 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK1536_4096 / 2 / 2 * BM1536_4096)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1536_4096; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM1536_1536 128
|
||||
#define BBK1536_1536 64
|
||||
inline void tbl_impl_1536_1536(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK1536_1536 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[8];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1536_1536; i += 64) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<8; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 2; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[2] += vec_v_left_1.val[0];
|
||||
vec_c[2] += vec_v_right_1.val[0];
|
||||
vec_c[3] += vec_v_left_1.val[1];
|
||||
vec_c[3] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[4] += vec_v_left_2.val[0];
|
||||
vec_c[4] += vec_v_right_2.val[0];
|
||||
vec_c[5] += vec_v_left_2.val[1];
|
||||
vec_c[5] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[6] += vec_v_left_3.val[0];
|
||||
vec_c[6] += vec_v_right_3.val[0];
|
||||
vec_c[7] += vec_v_left_3.val[1];
|
||||
vec_c[7] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
int32x4_t vec_v_bot_low_low_4 = vmovl_s16(vget_low_s16(vec_c[4]));
|
||||
int32x4_t vec_v_bot_low_high_4 = vmovl_high_s16(vec_c[4]);
|
||||
vst1q_s32(c + i + 32, vld1q_s32(c + i + 32) + vec_v_bot_low_low_4);
|
||||
vst1q_s32(c + i + 36, vld1q_s32(c + i + 36) + vec_v_bot_low_high_4);
|
||||
int32x4_t vec_v_bot_low_low_5 = vmovl_s16(vget_low_s16(vec_c[5]));
|
||||
int32x4_t vec_v_bot_low_high_5 = vmovl_high_s16(vec_c[5]);
|
||||
vst1q_s32(c + i + 40, vld1q_s32(c + i + 40) + vec_v_bot_low_low_5);
|
||||
vst1q_s32(c + i + 44, vld1q_s32(c + i + 44) + vec_v_bot_low_high_5);
|
||||
int32x4_t vec_v_bot_low_low_6 = vmovl_s16(vget_low_s16(vec_c[6]));
|
||||
int32x4_t vec_v_bot_low_high_6 = vmovl_high_s16(vec_c[6]);
|
||||
vst1q_s32(c + i + 48, vld1q_s32(c + i + 48) + vec_v_bot_low_low_6);
|
||||
vst1q_s32(c + i + 52, vld1q_s32(c + i + 52) + vec_v_bot_low_high_6);
|
||||
int32x4_t vec_v_bot_low_low_7 = vmovl_s16(vget_low_s16(vec_c[7]));
|
||||
int32x4_t vec_v_bot_low_high_7 = vmovl_high_s16(vec_c[7]);
|
||||
vst1q_s32(c + i + 56, vld1q_s32(c + i + 56) + vec_v_bot_low_low_7);
|
||||
vst1q_s32(c + i + 60, vld1q_s32(c + i + 60) + vec_v_bot_low_high_7);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_1536_1536(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM1536_1536];
|
||||
memset(&(CBits[0]), 0, BM1536_1536 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 1536 / BBK1536_1536; ++k_outer) {
|
||||
tbl_impl_1536_1536((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK1536_1536 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK1536_1536 / 2 / 2 * BM1536_1536)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1536_1536; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM4096_1536 256
|
||||
#define BBK4096_1536 128
|
||||
inline void tbl_impl_4096_1536(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK4096_1536 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_1536; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_4096_1536(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM4096_1536];
|
||||
memset(&(CBits[0]), 0, BM4096_1536 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 1536 / BBK4096_1536; ++k_outer) {
|
||||
tbl_impl_4096_1536((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK4096_1536 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK4096_1536 / 2 / 2 * BM4096_1536)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_1536; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
|
||||
template<int K>
|
||||
void preprocessor_k(void* B, void* LUT_Scales, void* QLUT) {{
|
||||
partial_max_reset((&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
per_tensor_quant(K, (&(((bitnet_float_type*)LUT_Scales)[0])), (&(((bitnet_float_type*)B)[0])));
|
||||
|
||||
lut_ctor<K>((&(((int8_t*)QLUT)[0])), (&(((bitnet_float_type*)B)[0])), (&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
}}
|
||||
void ggml_preprocessor(int m, int k, void* B, void* LUT_Scales, void* QLUT) {
|
||||
if (m == 1536 && k == 4096) {
|
||||
preprocessor_k<4096>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 1536 && k == 1536) {
|
||||
preprocessor_k<1536>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 4096 && k == 1536) {
|
||||
preprocessor_k<1536>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
}
|
||||
void ggml_qgemm_lut(int m, int k, void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
if (m == 1536 && k == 4096) {
|
||||
qgemm_lut_1536_4096(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 1536 && k == 1536) {
|
||||
qgemm_lut_1536_1536(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 4096 && k == 1536) {
|
||||
qgemm_lut_4096_1536(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {
|
||||
if (!(is_type_supported(tensor->type) && tensor->backend == GGML_BACKEND_TYPE_CPU && tensor->extra == nullptr)) {
|
||||
return;
|
||||
}
|
||||
|
||||
int k = tensor->ne[0];
|
||||
int m = tensor->ne[1];
|
||||
const int lut_scales_size = 1;
|
||||
const int scales_size = 1;
|
||||
int bk = 0;
|
||||
int bm = 0;
|
||||
|
||||
if (m == 1536 && k == 4096) {
|
||||
bm = BM1536_4096;
|
||||
bk = BBK1536_4096;
|
||||
}
|
||||
else if (m == 1536 && k == 1536) {
|
||||
bm = BM1536_1536;
|
||||
bk = BBK1536_1536;
|
||||
}
|
||||
else if (m == 4096 && k == 1536) {
|
||||
bm = BM4096_1536;
|
||||
bk = BBK4096_1536;
|
||||
}
|
||||
|
||||
const int n_tile_num = m / bm;
|
||||
const int BK = bk;
|
||||
uint8_t * qweights;
|
||||
bitnet_float_type * scales;
|
||||
|
||||
scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));
|
||||
qweights = (uint8_t *) tensor->data;
|
||||
float * i2_scales = (float * )(qweights + k * m / 4);
|
||||
scales[0] = (bitnet_float_type) i2_scales[0];
|
||||
|
||||
tensor->extra = bitnet_tensor_extras + bitnet_tensor_extras_index;
|
||||
bitnet_tensor_extras[bitnet_tensor_extras_index++] = {
|
||||
/* .lut_scales_size = */ lut_scales_size,
|
||||
/* .BK = */ BK,
|
||||
/* .n_tile_num = */ n_tile_num,
|
||||
/* .qweights = */ qweights,
|
||||
/* .scales = */ scales
|
||||
};
|
||||
}
|
||||
#endif
|
||||
@@ -5,8 +5,13 @@
|
||||
|
||||
#ifdef __ARM_NEON
|
||||
#include <arm_neon.h>
|
||||
#if defined(GGML_BITNET_ARM_TL1)
|
||||
typedef float32_t bitnet_float_type;
|
||||
#else
|
||||
typedef float16_t bitnet_float_type;
|
||||
#endif
|
||||
#else
|
||||
#include <immintrin.h>
|
||||
typedef float bitnet_float_type;
|
||||
#endif
|
||||
|
||||
@@ -43,6 +48,10 @@ GGML_API void ggml_preprocessor(int m, int k, void* B, void* LUT_Scales, void* Q
|
||||
GGML_API void ggml_qgemm_lut(int bs, int m, int k, int BK, void* A, void* sign, void* LUT, void* Scales, void* LUT_Scales, void* C);
|
||||
GGML_API void ggml_preprocessor(int bs, int m, int three_k, int two_k, void* B, void* LUT_Scales, void* Three_QLUT, void* Two_QLUT);
|
||||
#endif
|
||||
#if defined(GGML_BITNET_TL2_LOSS)
|
||||
GGML_API void ggml_qgemm_lut(int bs, int m, int k, int BK, void* A, void* sign, void* LUT, void* Scales, void* LUT_Scales, void* C);
|
||||
GGML_API void ggml_preprocessor(int bs, int m, int three_k, int two_k, void* B, void* Three_LUT_Scales, void* Two_LUT_Scales, void* Three_QLUT, void* Two_QLUT);
|
||||
#endif
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
[Kernels_0]
|
||||
m = 1536
|
||||
k = 4096
|
||||
bm = 256
|
||||
bk = 128
|
||||
bmm = 32
|
||||
|
||||
[Kernels_1]
|
||||
m = 1536
|
||||
k = 1536
|
||||
bm = 128
|
||||
bk = 64
|
||||
bmm = 64
|
||||
|
||||
[Kernels_2]
|
||||
m = 4096
|
||||
k = 1536
|
||||
bm = 256
|
||||
bk = 128
|
||||
bmm = 32
|
||||
|
||||
@@ -1,64 +0,0 @@
|
||||
import os
|
||||
import sys
|
||||
import signal
|
||||
import platform
|
||||
import argparse
|
||||
import subprocess
|
||||
|
||||
def run_command(command, shell=False):
|
||||
"""Run a system command and ensure it succeeds."""
|
||||
try:
|
||||
subprocess.run(command, shell=shell, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"Error occurred while running command: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def run_server():
|
||||
build_dir = "build"
|
||||
if platform.system() == "Windows":
|
||||
server_path = os.path.join(build_dir, "bin", "Release", "llama-server.exe")
|
||||
if not os.path.exists(server_path):
|
||||
server_path = os.path.join(build_dir, "bin", "llama-server")
|
||||
else:
|
||||
server_path = os.path.join(build_dir, "bin", "llama-server")
|
||||
|
||||
command = [
|
||||
f'{server_path}',
|
||||
'-m', args.model,
|
||||
'-c', str(args.ctx_size),
|
||||
'-t', str(args.threads),
|
||||
'-n', str(args.n_predict),
|
||||
'-ngl', '0',
|
||||
'--temp', str(args.temperature),
|
||||
'--host', args.host,
|
||||
'--port', str(args.port),
|
||||
'-cb' # Enable continuous batching
|
||||
]
|
||||
|
||||
if args.prompt:
|
||||
command.extend(['-p', args.prompt])
|
||||
|
||||
# Note: -cnv flag is removed as it's not supported by the server
|
||||
|
||||
print(f"Starting server on {args.host}:{args.port}")
|
||||
run_command(command)
|
||||
|
||||
def signal_handler(sig, frame):
|
||||
print("Ctrl+C pressed, shutting down server...")
|
||||
sys.exit(0)
|
||||
|
||||
if __name__ == "__main__":
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
parser = argparse.ArgumentParser(description='Run llama.cpp server')
|
||||
parser.add_argument("-m", "--model", type=str, help="Path to model file", required=False, default="models/bitnet_b1_58-3B/ggml-model-i2_s.gguf")
|
||||
parser.add_argument("-p", "--prompt", type=str, help="System prompt for the model", required=False)
|
||||
parser.add_argument("-n", "--n-predict", type=int, help="Number of tokens to predict", required=False, default=4096)
|
||||
parser.add_argument("-t", "--threads", type=int, help="Number of threads to use", required=False, default=2)
|
||||
parser.add_argument("-c", "--ctx-size", type=int, help="Size of the context window", required=False, default=2048)
|
||||
parser.add_argument("--temperature", type=float, help="Temperature for sampling", required=False, default=0.8)
|
||||
parser.add_argument("--host", type=str, help="IP address to listen on", required=False, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, help="Port to listen on", required=False, default=8080)
|
||||
|
||||
args = parser.parse_args()
|
||||
run_server()
|
||||
+33
-17
@@ -41,14 +41,11 @@ SUPPORTED_HF_MODELS = {
|
||||
"tiiuae/Falcon3-1B-Instruct-1.58bit": {
|
||||
"model_name": "Falcon3-1B-Instruct-1.58bit",
|
||||
},
|
||||
"microsoft/BitNet-b1.58-2B-4T": {
|
||||
"model_name": "BitNet-b1.58-2B-4T",
|
||||
},
|
||||
}
|
||||
|
||||
SUPPORTED_QUANT_TYPES = {
|
||||
"arm64": ["i2_s", "tl1"],
|
||||
"x86_64": ["i2_s", "tl2"]
|
||||
"arm64": ["i2_s", "tl1", "tl2-loss"],
|
||||
"x86_64": ["i2_s", "tl2", "tl2-loss"]
|
||||
}
|
||||
|
||||
COMPILER_EXTRA_ARGS = {
|
||||
@@ -114,8 +111,10 @@ def prepare_model():
|
||||
gguf_path = os.path.join(model_dir, "ggml-model-" + quant_type + ".gguf")
|
||||
if not os.path.exists(gguf_path) or os.path.getsize(gguf_path) == 0:
|
||||
logging.info(f"Converting HF model to GGUF format...")
|
||||
if quant_type.startswith("tl"):
|
||||
if quant_type in ["tl1", "tl2"]:
|
||||
run_command([sys.executable, "utils/convert-hf-to-gguf-bitnet.py", model_dir, "--outtype", quant_type, "--quant-embd"], log_step="convert_to_tl")
|
||||
elif quant_type in ["tl2-loss"]:
|
||||
run_command([sys.executable, "utils/convert-hf-to-gguf-bitnet.py", model_dir, "--outtype", "tl2", "--quant-embd", "--loss", "--outfile", model_dir + str("/ggml-model-tl2-loss.gguf")], log_step="convert_to_tl")
|
||||
else: # i2s
|
||||
# convert to f32
|
||||
run_command([sys.executable, "utils/convert-hf-to-gguf-bitnet.py", model_dir, "--outtype", "f32"], log_step="convert_to_f32_gguf")
|
||||
@@ -159,13 +158,20 @@ def gen_code():
|
||||
shutil.copyfile(os.path.join(pretuned_kernels, "bitnet-lut-kernels-tl2.h"), "include/bitnet-lut-kernels.h")
|
||||
shutil.copyfile(os.path.join(pretuned_kernels, "kernel_config_tl2.ini"), "include/kernel_config.ini")
|
||||
if get_model_name() == "bitnet_b1_58-large":
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "bitnet_b1_58-large", "--BM", "256,128,256", "--BK", "128,64,128", "--bm", "32,64,32"], log_step="codegen")
|
||||
if args.quant_type == "tl2-loss":
|
||||
run_command([sys.executable, "utils/codegen_tl2_loss.py", "--model", "bitnet_b1_58-large", "--BM", "256,128,256", "--BK", "96,96,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
else:
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "bitnet_b1_58-large", "--BM", "256,128,256", "--BK", "128,64,128", "--bm", "32,64,32"], log_step="codegen")
|
||||
elif get_model_name() in llama3_f3_models:
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "Llama3-8B-1.58-100B-tokens", "--BM", "256,128,256,128", "--BK", "128,64,128,64", "--bm", "32,64,32,64"], log_step="codegen")
|
||||
if args.quant_type == "tl2-loss":
|
||||
run_command([sys.executable, "utils/codegen_tl2_loss.py", "--model", "Llama3-8B-1.58-100B-tokens", "--BM", "256,128,256", "--BK", "96,96,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
else:
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "Llama3-8B-1.58-100B-tokens", "--BM", "256,128,256,128", "--BK", "128,64,128,64", "--bm", "32,64,32,64"], log_step="codegen")
|
||||
elif get_model_name() == "bitnet_b1_58-3B":
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "64,128,64", "--bm", "32,64,32"], log_step="codegen")
|
||||
elif get_model_name() == "BitNet-b1.58-2B-4T":
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "64,128,64", "--bm", "32,64,32"], log_step="codegen")
|
||||
if args.quant_type == "tl2-loss":
|
||||
run_command([sys.executable, "utils/codegen_tl2_loss.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "96,96,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
else:
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "64,128,64", "--bm", "32,64,32"], log_step="codegen")
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
else:
|
||||
@@ -177,13 +183,20 @@ def gen_code():
|
||||
sys.exit(1)
|
||||
shutil.copyfile(os.path.join(pretuned_kernels, "bitnet-lut-kernels-tl2.h"), "include/bitnet-lut-kernels.h")
|
||||
if get_model_name() == "bitnet_b1_58-large":
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "bitnet_b1_58-large", "--BM", "256,128,256", "--BK", "96,192,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
if args.quant_type == "tl2-loss":
|
||||
run_command([sys.executable, "utils/codegen_tl2_loss.py", "--model", "bitnet_b1_58-large", "--BM", "256,128,256", "--BK", "96,96,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
else:
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "bitnet_b1_58-large", "--BM", "256,128,256", "--BK", "96,192,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
elif get_model_name() in llama3_f3_models:
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "Llama3-8B-1.58-100B-tokens", "--BM", "256,128,256,128", "--BK", "96,96,96,96", "--bm", "32,32,32,32"], log_step="codegen")
|
||||
if args.quant_type == "tl2-loss":
|
||||
run_command([sys.executable, "utils/codegen_tl2_loss.py", "--model", "Llama3-8B-1.58-100B-tokens", "--BM", "256,128,256", "--BK", "96,96,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
else:
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "Llama3-8B-1.58-100B-tokens", "--BM", "256,128,256,128", "--BK", "96,96,96,96", "--bm", "32,32,32,32"], log_step="codegen")
|
||||
elif get_model_name() == "bitnet_b1_58-3B":
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "96,96,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
elif get_model_name() == "BitNet-b1.58-2B-4T":
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "96,96,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
if args.quant_type == "tl2-loss":
|
||||
run_command([sys.executable, "utils/codegen_tl2_loss.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "96,96,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
else:
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "96,96,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
@@ -199,7 +212,10 @@ def compile():
|
||||
logging.error(f"Arch {arch} is not supported yet")
|
||||
exit(0)
|
||||
logging.info("Compiling the code using CMake.")
|
||||
run_command(["cmake", "-B", "build", *COMPILER_EXTRA_ARGS[arch], *OS_EXTRA_ARGS.get(platform.system(), []), "-DCMAKE_C_COMPILER=clang", "-DCMAKE_CXX_COMPILER=clang++"], log_step="generate_build_files")
|
||||
if args.quant_type == "tl2-loss":
|
||||
run_command(["cmake", "-B", "build", "-DBITNET_TL2_LOSS=ON", *OS_EXTRA_ARGS.get(platform.system(), [])], log_step="generate_build_files")
|
||||
else:
|
||||
run_command(["cmake", "-B", "build", *COMPILER_EXTRA_ARGS[arch], *OS_EXTRA_ARGS.get(platform.system(), [])], log_step="generate_build_files")
|
||||
# run_command(["cmake", "--build", "build", "--target", "llama-cli", "--config", "Release"])
|
||||
run_command(["cmake", "--build", "build", "--config", "Release"], log_step="compile")
|
||||
|
||||
|
||||
@@ -154,6 +154,80 @@ size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, const stru
|
||||
return wsize;
|
||||
}
|
||||
|
||||
int ggml_bitnet_get_type_bits(enum ggml_type type) {
|
||||
switch (type) {
|
||||
case GGML_TYPE_TL2:
|
||||
return 2;
|
||||
case GGML_TYPE_Q4_0:
|
||||
return 4;
|
||||
default:
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
#if defined(GGML_BITNET_TL2_LOSS)
|
||||
void ggml_bitnet_init(void) {
|
||||
// LOG(INFO) << "ggml_bitnet_init";
|
||||
|
||||
if (initialized) {
|
||||
return;
|
||||
}
|
||||
initialized = true;
|
||||
|
||||
// if (wrapper == nullptr) {
|
||||
// wrapper = new BITNET::BITNETGeMMWrapper<bitnet_bitnet_float_type>();
|
||||
// }
|
||||
if (bitnet_tensor_extras == nullptr) {
|
||||
bitnet_tensor_extras = new bitnet_tensor_extra[GGML_BITNET_MAX_NODES];
|
||||
}
|
||||
bitnet_tensor_extras_index = 0;
|
||||
}
|
||||
|
||||
void ggml_bitnet_free(void) {
|
||||
// LOG(INFO) << "ggml_bitnet_free";
|
||||
|
||||
if (!initialized) {
|
||||
return;
|
||||
}
|
||||
initialized = false;
|
||||
|
||||
// delete wrapper;
|
||||
// wrapper = nullptr;
|
||||
for (size_t i = 0; i < bitnet_tensor_extras_index; i++) {
|
||||
// aligned_free(bitnet_tensor_extras[i].qweights);
|
||||
// aligned_free(bitnet_tensor_extras[i].scales);
|
||||
}
|
||||
delete[] bitnet_tensor_extras;
|
||||
bitnet_tensor_extras = nullptr;
|
||||
}
|
||||
|
||||
bool ggml_bitnet_can_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
|
||||
if ((is_type_supported(src0->type)) &&
|
||||
src1->type == GGML_TYPE_F32 &&
|
||||
dst->type == GGML_TYPE_F32 &&
|
||||
src0->backend == GGML_BACKEND_TYPE_CPU) {
|
||||
if (src1->ne[1] <= 1) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
|
||||
const size_t ne01 = src0->ne[1];
|
||||
const size_t ne10 = src1->ne[0];
|
||||
const size_t ne11 = src1->ne[1];
|
||||
|
||||
size_t wsize = ne10 * ne11 * 11 * sizeof(int8_t) + 2 * ne11 * 2 * sizeof(bitnet_float_type);
|
||||
if (sizeof(bitnet_float_type) == 2) {
|
||||
// Need fp32 to fp16 conversion
|
||||
wsize += std::max(ne10, ne01) * ne11 * sizeof(bitnet_float_type);
|
||||
}
|
||||
wsize = ((wsize - 1) / 64 + 1) * 64;
|
||||
return wsize;
|
||||
}
|
||||
|
||||
int ggml_bitnet_get_type_bits(enum ggml_type type) {
|
||||
switch (type) {
|
||||
case GGML_TYPE_TL2:
|
||||
|
||||
+19
-7
@@ -5,6 +5,7 @@ from configparser import ConfigParser
|
||||
def gen_ctor_code():
|
||||
kernel_code = "\n\
|
||||
#include \"ggml-bitnet.h\"\n\
|
||||
#include \"ggml-cpu-impl.h\"\n\
|
||||
#include <cstring>\n\
|
||||
#include <immintrin.h>\n\
|
||||
#define GGML_BITNET_MAX_NODES 8192\n\
|
||||
@@ -105,7 +106,7 @@ inline int32_t partial_max_reset(int32_t bs, void* lut_scales_) {\n\
|
||||
template<int act_k>\n\
|
||||
inline int32_t three_lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {\n\
|
||||
#if defined __AVX2__\n\
|
||||
__m256i vec_lut[16];\n\
|
||||
__m256 vec_lut[16];\n\
|
||||
const __m256i vec_bi = _mm256_set_epi32(84, 72, 60, 48, 36, 24, 12, 0);\n\
|
||||
float scales = *lut_scales;\n\
|
||||
__m256i shuffle_mask = _mm256_set_epi8(\n\
|
||||
@@ -191,7 +192,7 @@ inline int32_t three_lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_t
|
||||
template<int act_k>\n\
|
||||
inline int32_t two_lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {\n\
|
||||
#if defined __AVX2__\n\
|
||||
__m256i vec_lut[16];\n\
|
||||
__m256 vec_lut[16];\n\
|
||||
const __m256i vec_bi = _mm256_set_epi32(56, 48, 40, 32, 24, 16, 8, 0);\n\
|
||||
float scales = *lut_scales;\n\
|
||||
__m256i shuffle_mask = _mm256_set_epi8(\n\
|
||||
@@ -623,7 +624,7 @@ def gen_top_api(kernel_shapes, k_list):
|
||||
kernel_code = "".join([kernel_code, "}\n"])
|
||||
return kernel_code
|
||||
|
||||
def gen_transform_code(kernel_shapes):
|
||||
def gen_transform_code(kernel_shapes, fp16):
|
||||
kernel_code = "\n\
|
||||
void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {\n\
|
||||
if (!(is_type_supported(tensor->type) && tensor->backend == GGML_BACKEND_TYPE_CPU && tensor->extra == nullptr)) {\n\
|
||||
@@ -657,10 +658,20 @@ void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {\n\
|
||||
scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));\n\
|
||||
qweights = (uint8_t *) tensor->data;\n\
|
||||
int nbytes = (k - 256) * m / 3 * 5 / 8 + 256 * m / 2 * 4 / 8;\n\
|
||||
if (nbytes % 32 != 0) nbytes = 32 - nbytes % 32 + nbytes;\n\
|
||||
nbytes = 32 - nbytes % 32 + nbytes;\n\
|
||||
float * i2_scales = (float * )(qweights + nbytes);\n\
|
||||
scales[0] = (bitnet_float_type) i2_scales[0];\n\
|
||||
\n\
|
||||
\n"])
|
||||
|
||||
if fp16:
|
||||
kernel_code = "".join([kernel_code, "\
|
||||
ggml_fp16_t* fp16_scale = (ggml_fp16_t *)aligned_malloc(sizeof(ggml_fp16_t));\n\
|
||||
fp16_scale[0] = GGML_FP32_TO_FP16(i2_scales[0]);\n\
|
||||
scales[0] = (bitnet_float_type) GGML_FP16_TO_FP32(fp16_scale[0]);\n"])
|
||||
else:
|
||||
kernel_code = "".join([kernel_code, "\
|
||||
scales[0] = (bitnet_float_type) i2_scales[0];\n"])
|
||||
|
||||
kernel_code = "".join([kernel_code, "\n\
|
||||
tensor->extra = bitnet_tensor_extras + bitnet_tensor_extras_index;\n\
|
||||
bitnet_tensor_extras[bitnet_tensor_extras_index++] = {\n\
|
||||
/* .lut_scales_size = */ lut_scales_size,\n\
|
||||
@@ -702,6 +713,7 @@ if __name__ == "__main__":
|
||||
help="block length when cutting one weight (M, K) into K / BK weights (M, BK).")
|
||||
parser.add_argument('--bm',default="input", type=str,
|
||||
help="using simd instructions to compute (bm, 192 / bm) in one block")
|
||||
parser.add_argument('--fp16', action="store_true", help="convert scale to fp16")
|
||||
args = parser.parse_args()
|
||||
|
||||
kernel_shapes = ModelShapeDict[args.model]
|
||||
@@ -730,7 +742,7 @@ if __name__ == "__main__":
|
||||
|
||||
ctor_code = gen_ctor_code()
|
||||
api_code = gen_top_api(kernel_shapes, k_list)
|
||||
trans_code = gen_transform_code(kernel_shapes)
|
||||
trans_code = gen_transform_code(kernel_shapes, args.fp16)
|
||||
|
||||
output_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "include")
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -517,6 +517,92 @@ def preprocess_weights_tl1(
|
||||
return weight
|
||||
|
||||
|
||||
def preprocess_two_weights_tl2_loss(M, K, weight_num, BM, BY, bm, by, weight, final_weight):
|
||||
weight = np.reshape(weight, (weight_num // 2, 2))
|
||||
hi_weight = np.multiply(np.split(weight, 2, axis=1)[0], 3)
|
||||
lo_weight = np.split(weight, 2, axis=1)[1]
|
||||
|
||||
weight = np.reshape((hi_weight + lo_weight), weight_num // 2)
|
||||
weight = weight + 4
|
||||
weight = np.reshape(weight, (M, K // 2)).astype(np.uint8)
|
||||
weight = weight.reshape((M // BM, BM, K // 2)).transpose(0, 2, 1)
|
||||
weight = weight.reshape((M // BM, K // BY, BY // 2, BM)).transpose(0, 1, 3, 2)
|
||||
weight = weight.reshape((M // BM, K // BY, BM // bm, bm, BY // 2)).transpose(0, 1, 2, 4, 3)
|
||||
weight = weight.reshape((M // BM, K // BY, BM // bm, BY // by, by // 2, bm)).transpose(0, 1, 2, 3, 5, 4)
|
||||
weight = weight.reshape((M // BM, K // BY, BM // bm, BY // by, bm, by // 2))
|
||||
weight_0 = weight[:, :, :, :, :, 0]
|
||||
weight_1 = weight[:, :, :, :, :, 1]
|
||||
weight_0 = weight_0 << 4
|
||||
weight_1 = weight_1
|
||||
weight = weight_0 + weight_1
|
||||
weight = weight.reshape(M * K // bm // by, bm).reshape(M * K // by // 16, 16)
|
||||
|
||||
for i in range(weight.shape[0]):
|
||||
final_weight.append(weight[i, :])
|
||||
|
||||
def preprocess_three_weights_tl2_loss(M, K, weight_num, BM, BY, bm, by, weight, final_weight):
|
||||
weight = np.reshape(weight, (weight_num // 3, 3))
|
||||
split_weights = np.split(weight, 3, axis=1)
|
||||
first_weight = np.multiply(split_weights[0], 9)
|
||||
second_weight = np.multiply(split_weights[1], 3)
|
||||
third_weight = split_weights[2]
|
||||
|
||||
weight = np.reshape((first_weight + second_weight + third_weight), weight_num // 3)
|
||||
sign_weight = np.sign(weight)
|
||||
sign_weight = np.where(sign_weight < 1, 0, sign_weight)
|
||||
weight = np.abs(weight)
|
||||
|
||||
weight = np.reshape(weight, (M, K // 3)).astype(np.uint8)
|
||||
sign_weight = np.reshape(sign_weight, (M, K // 3)).astype(np.uint8)
|
||||
|
||||
weight = weight.reshape((M // BM, BM, K // 3)).transpose(0, 2, 1)
|
||||
weight = weight.reshape((M // BM, K // BY, BY // 3, BM)).transpose(0, 1, 3, 2)
|
||||
weight = weight.reshape((M // BM, K // BY, BM // bm, bm, BY // 3)).transpose(0, 1, 2, 4, 3)
|
||||
weight = weight.reshape((M // BM, K // BY, BM // bm, BY // by, by // 3, bm)).transpose(0, 1, 2, 3, 5, 4)
|
||||
weight = weight.reshape((M // BM, K // BY, BM // bm, BY // by, bm, by // 3))
|
||||
|
||||
weight_list = []
|
||||
for i in range(by // 3):
|
||||
weight_list.append(weight[:, :, :, :, :, i])
|
||||
|
||||
for i in range(by // 3 // 2):
|
||||
weight_list[i] = weight_list[i] << 4
|
||||
weight_list[i + by // 3 // 2] = weight_list[i + by // 3 // 2]
|
||||
weight_list[i] = weight_list[i] + weight_list[i + by // 3 // 2]
|
||||
weight_list[i] = weight_list[i].reshape(M * K // bm // by, bm).reshape(M * K // by // 16, 16)
|
||||
|
||||
for i in range(weight_list[0].shape[0]):
|
||||
for j in range(by // 3 // 2):
|
||||
final_weight.append(weight_list[j][i, :])
|
||||
|
||||
sign_weight = sign_weight.reshape((M // BM, BM, K // 3)).transpose(0, 2, 1)
|
||||
sign_weight = sign_weight.reshape((M // BM, K // BY, BY // 3, BM)).transpose(0, 1, 3, 2)
|
||||
sign_weight = sign_weight.reshape((M // BM, K // BY, BM // bm, bm, BY // 3)).transpose(0, 1, 2, 4, 3)
|
||||
sign_weight = sign_weight.reshape((M // BM, K // BY, BM // bm, BY // (by * 4), by // 3 * 4, bm)).transpose(0, 1, 2, 3, 5, 4).astype(np.uint8)
|
||||
|
||||
combine_weight_list = []
|
||||
for i in range(by // 3 // 2):
|
||||
combine_weight = np.zeros((M // BM, K // BY, BM // bm, BY // (by * 4), bm), dtype=np.uint8)
|
||||
combine_weight_list.append(combine_weight)
|
||||
|
||||
for i in range(8):
|
||||
for j in range(by // 3 // 2):
|
||||
if bm == 16:
|
||||
combine_weight_list[j] = combine_weight_list[j] + (sign_weight[:, :, :, :, :, by // 3 // 2 * i + j] << 7 - i)
|
||||
elif bm == 32:
|
||||
if i > 3 :
|
||||
ti = (i - 4) * 2 + 1
|
||||
else:
|
||||
ti = i * 2
|
||||
combine_weight_list[j] = combine_weight_list[j] + (sign_weight[:, :, :, :, :, by // 3 // 2 * ti + j] << 7 - i)
|
||||
|
||||
for i in range(by // 3 // 2):
|
||||
combine_weight_list[i] = combine_weight_list[i].reshape((M * K // (by * 4)) // 16, 16)
|
||||
|
||||
for i in range(combine_weight_list[0].shape[0]):
|
||||
for j in range(by // 3 // 2):
|
||||
final_weight.append(combine_weight_list[j][i, :])
|
||||
|
||||
def preprocess_two_weights_tl2(M, K, weight_num, BM, BY, bm, by, weight, final_weight):
|
||||
weight = np.reshape(weight, (weight_num // 2, 2))
|
||||
hi_weight = np.multiply(np.split(weight, 2, axis=1)[0], 3)
|
||||
@@ -603,7 +689,6 @@ def preprocess_weights_tl2(
|
||||
weight = w
|
||||
weight = np.where(np.abs(weight) < 1e-6, 0, weight).astype(np.float32)
|
||||
weight = np.sign(weight)
|
||||
weight_num = np.prod(weight.shape)
|
||||
|
||||
config.read('include/kernel_config.ini')
|
||||
BM = -1
|
||||
@@ -631,7 +716,8 @@ def preprocess_weights_tl2(
|
||||
|
||||
final_weight = []
|
||||
|
||||
preprocess_three_weights_tl2(three_weight.shape[0],
|
||||
if args.loss:
|
||||
preprocess_three_weights_tl2_loss(three_weight.shape[0],
|
||||
three_weight.shape[1],
|
||||
three_weight.shape[0] * three_weight.shape[1],
|
||||
BM,
|
||||
@@ -641,8 +727,29 @@ def preprocess_weights_tl2(
|
||||
three_weight,
|
||||
final_weight)
|
||||
|
||||
if (weight.shape[1] % BY != 0):
|
||||
preprocess_two_weights_tl2( two_weight.shape[0],
|
||||
if (weight.shape[1] % BY != 0):
|
||||
preprocess_two_weights_tl2_loss(two_weight.shape[0],
|
||||
two_weight.shape[1],
|
||||
two_weight.shape[0] * two_weight.shape[1],
|
||||
BM,
|
||||
32,
|
||||
32,
|
||||
4,
|
||||
two_weight,
|
||||
final_weight)
|
||||
else:
|
||||
preprocess_three_weights_tl2(three_weight.shape[0],
|
||||
three_weight.shape[1],
|
||||
three_weight.shape[0] * three_weight.shape[1],
|
||||
BM,
|
||||
BY,
|
||||
bm,
|
||||
by,
|
||||
three_weight,
|
||||
final_weight)
|
||||
|
||||
if (weight.shape[1] % BY != 0):
|
||||
preprocess_two_weights_tl2(two_weight.shape[0],
|
||||
two_weight.shape[1],
|
||||
two_weight.shape[0] * two_weight.shape[1],
|
||||
BM,
|
||||
@@ -652,8 +759,10 @@ def preprocess_weights_tl2(
|
||||
two_weight,
|
||||
final_weight)
|
||||
weight = np.array(final_weight, dtype=np.uint8).reshape(-1)
|
||||
weight = np.pad(weight, (0, (K - 256) * M // 3 * 5 // 8 + 256 * M // 2 * 4 // 8 -
|
||||
weight.shape[0]), mode='constant', constant_values=0)
|
||||
pad_nums = (K - 256) * M // 3 * 5 // 8 + 256 * M // 2 * 4 // 8
|
||||
pad_align_nums = 32 - ((K - 256) * M // 3 * 5 // 8 + 256 * M // 2 * 4 // 8) % 32
|
||||
pad_nums = pad_nums + pad_align_nums
|
||||
weight = np.pad(weight, (0, pad_nums - weight.shape[0]), mode='constant', constant_values=0)
|
||||
return weight
|
||||
|
||||
def transform_to_tl1(x: np.ndarray):
|
||||
@@ -1116,6 +1225,7 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--model-name", type=str, default=None, help="name of the model")
|
||||
parser.add_argument("--verbose", action="store_true", help="increase output verbosity")
|
||||
parser.add_argument("--quant-embd", action="store_true", help="quantize the embedding layer")
|
||||
parser.add_argument("--loss", action="store_true", help="use loss tl2")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user