BitNet

BitNetは、Multi-Head Attentionおよびフィードフォワードネットワーク内の従来の線形層を、専用のBitLinear層に置き換えます。BitLinear層は、重みを3値精度（-1、0、1の値）で量子化し、活性化値を8ビット精度で量子化します。

BitNetモデルは、実行時の量子化（オンザフライ量子化）はできません。Quantization-Aware Training (QAT) 手法であるため、事前学習またはファインチューニングの段階で量子化を行う必要があります。学習中は、対称的なテンソルごとの量子化により、重みが3値に量子化されます。

1. 重み行列の絶対値の平均を計算し、スケールとして使用します。
2. 重みをスケールで割り、値を丸め、-1から1の範囲に制限し、フル精度で処理を続行するために再スケーリングします。
3. アクティベーションは、absmax量子化（チャネルごとの対称量子化）を用いて、指定されたビット幅（8ビット）に量子化されます。これには、アクティベーションを[−128,127]の範囲にスケーリングする処理が含まれます。

1.58ビットモデルの事前学習またはファインチューニングを行うには、Nanotronを使用するこのPRを参照してください。ファインチューニングを行う場合は、Hugging FaceのモデルをNanotron形式に変換します。変換の手順は、このPRに記載されています。

from_pretrained()でBitNet量子化モデルを読み込みます。

from transformers import AutoModelForCausalLM
path = "/path/to/model"
model = AutoModelForCausalLM.from_pretrained(path, device_map="auto")

カーネル

@torch.compile は、重みの展開とフォワードパスの実行に使用されます。実装は非常にシンプルであり、大幅な速度向上を実現します。将来のバージョンでは、さらに最適化されたカーネルが統合される予定です。

リソース

BitNetモデルの学習とファインチューニングの詳細については、Fine-tuning LLMs to 1.58bit: extreme quantization made easyをお読みください。