README.md

---
license: mit
datasets:
- wikipedia
---
# BitLinear-phi-1.5

BitLinear-phi-1.5 is a model trained partially using the method described in [The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits](https://arxiv.org/abs/2402.17764).

### Notice: Our BitLinear layer will only apply 1-bit quantization to the weight
### Other components (RMSnorm, activation quant) in the paper is discarded.
Idea behind: The major contribution in their paper is introduced a valid binary weight quantization, we don't want to mix it with other components to make it difficult to evaluate the major part.

The model structure is from [phi-1.5](https://huggingface.co/microsoft/phi-1_5), with all linear layers except lm_head replaced with our custom BitLinear layer.

It was trained on a small subset of the [wikipedia dataset](https://huggingface.co/datasets/wikipedia) dataset, for research validation purpose only. 


```python
dataset = load_dataset("wikipedia", "20220301.en")
dataset = dataset['train'].select(range(int(1e5)))
```
Please notice the kernel is not optimzed for 1-bit matrix yet.

The model is trained on a 3090(24GB) for 16 hours. 

## For faster(3x) inference, check https://github.com/Mrw33554432/Bitlinear4HF and install custom kernel
## For training code, check https://github.com/Mrw33554432/Bitlinear4HF. 

The training code should be compatible with most of the LLMs in huggingface.

Using pretrained model weight (normal models) for training will not work due to gradient explosion.

## Sample inference code (slow)


```python
import torch
from replace_hf import replace_linear_in_hf
from transformers import AutoModelForCausalLM, AutoTokenizer


def quick_test(model, tokenizer, prompt: str):
    # Encode the inputs
    inputs = tokenizer.encode(prompt, return_tensors="pt")

    # Generate outputs
    outputs = model.generate(inputs, max_length=64)

    # Decode and print the outputs
    print(tokenizer.decode(outputs[0]))


torch.set_default_device("cuda")

tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-1_5", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("Mrw33554432/bitLinear-phi-1.5", trust_remote_code=True, torch_dtype=torch.float16)

print(model)
# Replace Linear layers with BitLinear
replace_linear_in_hf(model, keep_param=True)
print(model)

quick_test(model, tokenizer, prompt="Tom is the")
```