皆さんこんにちは！@2626です。弊社ブログの移行作業中ですが、今回初めて新記事を書いてみました！

記念すべき第一回目は、BitNet について紹介したいと思います！BitNet は個人的に注目している、LLM を少ないメモリで動かすための仕組みなのですが、これまで実用的な学習済みモデルが無く、目立つこともありませんでした。しかし約2週間前、Microsoft から BitNet を用いた Llama8B が公開されました！

しかもQiitaを探してみても、あまり書かれている記事がないということで、今回はこの BitNet について紹介していきます！

BitNetとは？

論文のAbstractには下記のように書かれています。

大規模な言語モデルのサイズが大きくなるにつれ、その展開に課題が生じ、エネルギー消費量の多さによる環境への影響が懸念されている。本研究では、大規模言語モデル用に設計されたスケーラブルで安定した1ビット変換器アーキテクチャであるBitNetを紹介する。具体的には、1ビットの重みをゼロから学習するために、nn.Linear層のドロップイン置き換えとしてBitLinearを導入する。言語モデリングに関する実験結果は、BitNetが最新の8ビット量子化手法やFP16 Transformerベースラインと比較して、メモリフットプリントとエネルギー消費を大幅に削減しながら、競争力のある性能を達成することを示している。さらに、BitNetは全精度Transformerに似たスケーリング則を示し、効率と性能の利点を維持しながら、さらに大規模な言語モデルに効果的にスケーリングできる可能性を示唆している。¹

なお、上記はDeepLによって翻訳していることをご了承ください。

端的に解説すると、BitNetとは、Transformer アーキテクチャの Linear 層を BitLinear という層に置き換えたアーキテクチャのことです。では BitLinear とは何でしょう？ BitLinear とは、重み行列の各成分の値を $-1$ または $1$ のみで表現する Linear 層のことです。LLMでは膨大なパラメータ数を削減するために、量子化という手法が用いられることがあります。通常のAIモデルが fp32 で重みの計算をするのに対し、fp16, int8, int4 などより小さいメモリサイズで計算する手法ですね。その量子化を極限にまで突き詰めて $[-1,1]$ だけにしてしまった…という感じでしょうか。

それから約4か月後、今度は BitNet b1.58 が現れます。上記の $[-1,1]$ だけにする方法では何か問題があったのでしょうか？少し調べてみたのですが、詳しいことは分かりませんでした（読者の方でご存じの方がいらっしゃればご連絡ください）。次に BitNet b1.58 について説明していきましょう。

BitNet b1.58 とは？

論文のAbstractには下記のように書かれています。

BitNet[WMD+23]のような最近の研究は、1bit Large Language Models(LLM)の新しい時代への道を開いている。本研究では、BitNet b1.58という1bitLLM変数を導入する。BitNet b1.58では、LLMの全てのパラメータ（重み）が3値{-1, 0, 1}である。このLLMは、同じモデルサイズと学習トークンを持つ全精度（つまりFP16またはBF16）Transformer LLMとパープレキシティとエンドタスクパフォーマンスの両方で一致し、レイテンシ、メモリ、スループット、エネルギー消費の点で大幅にコスト効率が向上しています。さらに深いことに、1.58ビットLLMは、高性能でコスト効率の高い新世代のLLMを訓練するための新しいスケーリング則とレシピを定義しています。さらに、新しい計算パラダイムを可能にし、1ビットLLMに最適化された特定のハードウェアを設計するための扉を開く。²

こちらも、DeepLによって翻訳していることをご了承ください。

こちらも少し解説します。BitNet b1.58 は $[-1,0,1]$ という3値で計算すると明記されています。先ほど紹介した BitNet は2値だったので、その点が明確に違うところですね。それによって、fp16 や bf16 で作られた Transformer LLM と遜色ない精度かつレイテンシやメモリを大幅に抑えることができたと明記されています。

ちなみに $1.58$ とは BitNet が $[-1,1]$ の2値なのに対し、3値であることから来ているようです。$\log_2 3 \fallingdotseq 1.58 $ からですね。この記事では簡単な仕組みの話のみにとどめますが、実際は学習フェーズから様々な工夫を凝らしているようです。ただ完成したモデルを2値や3値にすればいいというわけではないんですね（PTQ, Post Training Quantization と言ったりします）。この辺りの詳しい話は割愛いたしますので、参考文献をご覧いただければと思います。

ということで、いよいよ動かしてみましょう…！！！

実際に使ってみた

いくつか既に日本語の記事は出ておりますが、Dockerで構築した記事を見なかったので、今回はDockerで構築してみました。今回使用したPCの環境は以下の通りです。

CPU: Intel® Core™ i9-13900, 2.00 GHz
メモリ: 64GB
OS: Windows 11 Pro, 23H2
WSL
- WSL バージョン: 2.0.14.0
- カーネルバージョン: 5.15.133.1-1
- WSLg バージョン: 1.0.59
- MSRDC バージョン: 1.2.4677
- Direct3D バージョン: 1.611.1-81528511
- DXCore バージョン: 10.0.25131.1002-220531-1700.rs-onecore-base2-hyp
- Windows バージョン: 10.0.22631.4391
Docker: Docker version 26.1.1, build 4cf5afa
WSLに直接Dockerを入れてるため、DockerDesktopは無し

Dockerfile および docker-compose.yml はこちらです。なお、セキュリティや細かい設定の必要性などの検証は行っておりませんので、利用は自己責任でお願いいたします。

Dockerfile

FROM python:3.9-bookworm

RUN apt-get update \
    && apt-get upgrade -y

RUN apt-get install -y \
    clang \
    cmake \
    lsb-release \
    wget \
    software-properties-common \
    gnupg \
    libomp-dev

RUN apt-get clean \
    && rm -rf /var/lib/apt/lists/*

RUN git clone --recursive https://github.com/microsoft/BitNet.git

RUN pip install --upgrade pip
RUN pip install -r /BitNet/requirements.txt

WORKDIR /BitNet

docker-compose.yml

services:
  bitnet:
    image: bitnet
    container_name: bitnet

    build:
      context: .
      dockerfile: Dockerfile
      
    tty: true

実際に動かしてみる

ではここから実際に動かしてみましょう。

$ docker compose up

無事にコンテナが立ち上がったことを確認したら、

$ docker exec -it bitnet bash

を実行します。あとは公式に従えばOKです。ちなみに、公式には3種類のモデルが用意されています。

bitnet_b1_58-large
bitnet_b1_58-3B
Llama3-8B-1.58-100B-tokens

今回は Llama3-8B-1.58-100B-tokens を使用してみます。

$ python setup_env.py --hf-repo HF1BitLLM/Llama3-8B-1.58-100B-tokens -q i2_s
INFO:root:Compiling the code using CMake.
INFO:root:Downloading model HF1BitLLM/Llama3-8B-1.58-100B-tokens from HuggingFace to models/Llama3-8B-1.58-100B-tokens...
INFO:root:GGUF model already exists at models/Llama3-8B-1.58-100B-tokens/ggml-model-i2_s.gguf

とりあえず上手く行ったようです（既に実行していたので、表示されているログは変わっているかもしれません）。再び公式に従って操作してみましょう。使用するプロンプトは下記のとおりです。

Daniel went back to the the the garden. Mary travelled to the kitchen. Sandra journeyed to the kitchen. Sandra went to the hallway. John went to the bedroom. Mary went back to the garden. Where is Mary?\nAnswer:

最後の一文に「Mary は庭に戻った」とあるので、「Mary は庭にいる」と答えてほしいところですが…。

$ python run_inference.py -m models/Llama3-8B-1.58-100B-tokens/ggml-model-i2_s.gguf -p "Daniel went back to the the the garden. Mary travelled to the kitchen. Sandra journeyed to the kitchen. Sandra went to the hallway. John went to the bedroom. Mary went back to the garden. Where is Mary?\nAnswer:" -n 6 -temp 0
warning: not compiled with GPU offload support, --gpu-layers option will be ignored
warning: see main README.md for information on enabling GPU BLAS support
build: 3947 (406a5036) with Debian clang version 14.0.6 for x86_64-pc-linux-gnu

--- 中略 ---

Daniel went back to the the the garden. Mary travelled to the kitchen. Sandra journeyed to the kitchen. Sandra went to the hallway. John went to the bedroom. Mary went back to the garden. Where is Mary?
Answer: Mary is in the garden.


llama_perf_sampler_print:    sampling time =       0.24 ms /    54 runs   (    0.00 ms per token, 226890.76 tokens per second)
llama_perf_context_print:        load time =     463.91 ms
llama_perf_context_print: prompt eval time =    4454.56 ms /    48 tokens (   92.80 ms per token,    10.78 tokens per second)
llama_perf_context_print:        eval time =     455.85 ms /     5 runs   (   91.17 ms per token,    10.97 tokens per second)
llama_perf_context_print:       total time =    4912.22 ms /    53 tokens

「Answer: Mary is in the garden.」と見事に答えることができていますね！無事にサンプルを動かすことができました！

それではもう少しだけ動かしてみましょう。歴史クイズを出してみます。

python run_inference.py -m models/Llama3-8B-1.58-100B-tokens/ggml-model-i2_s.gguf -p "Who was the president of the United States as of 1850?" -n 6 -temp 0

「Who was the president of the United States as of 1850?」ということで、1850年時点のアメリカの大統領を聞いてみます。ちなみに答えは「ミラード・フィルモア」です（初めて知りました）。

Who was the president of the United States as of 1850? A. Millard Fillmore

見事に正解ですね！次に地理も行ってみましょう！

python run_inference.py -m models/Llama3-8B-1.58-100B-tokens/ggml-model-i2_s.gguf -p "What treaty was adopted by the EC member states in 1992 to build and develop the European Union (EU), specifically amending the Treaty of Rome to strengthen economic integration?" -n 6 -temp 0

「経済統合の強化などを目指してEC加盟国がローマ条約を改正して1992年に採択した、EUの建設と発展を目指す条約はなんでしょう。」ということで意味不明な問題です…（泣）ちなみに正解は「マーストリヒト条約」です。

What treaty was adopted by the EC member states in 1992 to build and develop the European Union (EU), specifically amending the Treaty of Rome to strengthen economic integration? The Maastricht Treaty

すごいですね！見事に正解です。

まとめ

ということで、今回は BitNet を試してみました！今現在の生成 AI は2パターンに進化しているように感じます。

モデルサイズを大きくして、精度をより高めていく
モデルサイズを小さくしつつ精度をあげることで、エッジデバイスやモバイル、CPUで動くようにする

今回紹介した BitNet は CPU でも動作しますし、動作速度も十二分に満足いくものでした。生成 AI に関わる端くれとして、今後の更なる進化を見守っていきたいと思います！

参考文献

Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Huaijie Wang, Lingxiao Ma, Fan Yang, Ruiping Wang, Yi Wu, Furu Wei, Microsoft Research, University of Chinese Academy of Sciences, Tsinghua University, "BitNet: Scaling 1-bit Transformers for Large Language Models", 2023 ↩
Shuming Ma, Hongyu Wang, Lingxiao Ma, Lei Wang, Wenhui Wang, Shaohan Huang, Li Dong, Ruiping Wang, Jilong Xue, Furu Wei, "The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits", 2024 ↩