Falcon-E-prequantized

0. TL;DR
1. Model Details
2. Training Details
3. Usage
4. Evaluation
5. Citation

This is simply the mirror of https://huggingface.co/tiiuae/Falcon-E-3B-Base - branch prequantized

TL;DR

Model Details

Model Description

• Developed by: https://www.tii.ae
• Model type: Causal decoder-only / Base version
• Architecture: Pure-transformer - 1.58bit version
• Language(s) (NLP): English
• License: Falcon-LLM License

Training details

For more details about the training protocol of this model, please refer to the Falcon-E technical blogpost.

Usage

Currently to use this model you can either rely on Hugging Face transformers library or BitNet library. There are multiple ways to interact with the model depending on your target usage. For each of the Falcon-E series model, you have three variants: the BitNet model, the prequantized checkpoint for fine-tuning and the bfloat16 version of the BitNet model.

Inference

🤗 transformers

In case you want to perform inference on the BitNet checkpoint run:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "tiiuae/Falcon-E-1B-Base"

model = AutoModelForCausalLM.from_pretrained(
  model_id,
  torch_dtype=torch.bfloat16,
).to("cuda")

# Perform text generation

If you want to rather use the classic bfloat16 version, you can run:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "tiiuae/Falcon-E-1B-Base"
revision = "bfloat16"

model = AutoModelForCausalLM.from_pretrained(
  model_id,
  torch_dtype=torch.bfloat16,
  revision=revision,
).to("cuda")

# Perform text generation

BitNet

git clone https://github.com/microsoft/BitNet && cd BitNet
pip install -r requirements.txt
python setup_env.py --hf-repo tiiuae/Falcon-E-1B-Base -q i2_s
python run_inference.py -m models/Falcon-E-1B-Base/ggml-model-i2_s.gguf -p "You are a helpful assistant" -cnv

Apply mlx-lm

pip install -U mlx-lm

Then:

mlx_lm.generate --model tiiuae/Falcon-E-3B-Instruct --prompt "Implement bubble sort" --max-tokens 100 --temp 0.1

Fine-tuning

For fine-tuning the model, you should load the prequantized revision of the model and use the onebitllms Python package:

import torch

from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer
+ from onebitllms import replace_linear_with_bitnet_linear, quantize_to_1bit

model_id = "tiiuae/Falcon-E-1B-Base"

tokenizer = AutoTokenizer.from_pretrained(model_id, revision="prequantized")
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
+    revision="prequantized"
)
+ model = replace_linear_with_bitnet_linear(model)

trainer = SFTTrainer(
    model,
    ...
)

trainer.train()

+ quantize_to_1bit(output_directory)

Evaluation

We report in the following table our internal pipeline benchmarks:

Note evaluation results are normalized score from former Hugging Face leaderboard v2 tasks

Below are the results for instruction fine-tuned models: