Lm.rs: 依存なしのRust製ミニマルCPU LLM推論技術

 (github.com/samuel-vitorino)
1 ポイント 投稿者 GN⁺ 2024-10-12 | 1件のコメント | WhatsAppで共有
  • lm.rs は、RustでCPU上のローカル言語モデル推論を実行するプロジェクトで、MLライブラリなしで推論全体を実行する最小コード実装を目指している
  • Karpathyの llama2.c と llm.c に着想を得ており、当初は Google Gemma 2 のみをサポートしていたが、Llama 3.2 と PHI-3.5 の画像入力サポートまで拡張された
  • 最新の変更で バッチ処理 が実装され、画像エンコーディング速度が最大約3倍高速化され、Llama 3.2 1B は作者の16コアマシンで 50 tok/s で動作する
  • 用意済みのモデルは Hugging Face から入手でき、README では Q8_0 の使用を推奨し、Q4_0 量子化はまだ改善中だとしている
  • ユーザーは LMRS 形式のモデルとトークナイザーをダウンロードしてすぐにビルドするか、Hugging Face の元モデルファイルを export.py と tokenizer.py で変換して実行できる

lm.rs が目指すもの

  • lm.rs は Rust で書かれたローカルCPUベースの言語モデル推論実装である
  • 目標は、MLライブラリなしでCPU上で言語モデルの 完全な推論 を実行する最小コード実装である
  • Karpathy の llama2.c と llm.c に着想を得ている
  • README では、現在のコードは「そこまで最小ではなく」、一部のコードには最適化や改善の余地があると述べている
  • このプロジェクトは、作者が初めて Rust を使ってみるきっかけでもあった

サポートモデルとマルチモーダル拡張

  • 当初は Google Gemma 2 モデルのみをサポートしていたが、その後 Llama 3.2 モデルのサポートが追加された
  • 最近では PHI-3.5 を通じて画像利用オプションが追加された
  • 現在強調されているサポート項目
    • PHI-3.5-vision モデルによるマルチモーダル対応
    • PHI-3.5-mini テキスト専用モデル対応
  • 関連リソース

性能と用意済みモデル

モデル変換の流れ

  • 用意済みの量子化モデルとトークナイザーを Hugging Face から取得すれば、変換工程を省略できる
  • Google または Meta が Hugging Face に公開したモデルを直接変換するには、追加の Python 依存関係をインストールする必要がある
pip install -r requirements.txt
  • 元モデルのページから .safetensors と config.json ファイルをダウンロードして使用する
  • PHI3.5 Vision のようなマルチモーダルモデルには CLIP の config ファイルも必要である
  • export.py は bfloat16 重みを LMRS 形式 に変換する
python export.py --files [ordered .safetensor files] --config [model config.json] --save-path [name and path to save] --type [model type (GEMMA/LLAMA/PHI)]
  • 量子化版を書き出すには --quantize と --quantize-type フラグを使う
  • int8 量子化モデルのサイズは、グループサイズに応じて約 9.8G から約 2.5G まで縮小できる
  • マルチモーダルモデルでは --vision-config 引数を含める必要がある
  • tokenizer.py はトークナイザーモデルを LMRS トークナイザー形式に変換する
python tokenizer.py --model-id [huggingface model_id] --tokenizer-type [type of the tokenizer (GEMMA/LLAMA/PHI)]

ビルドと実行

  • Rust コードは cargo でコンパイルし、README では target-cpu フラグの指定を明記している
RUSTFLAGS="-C target-cpu=native" cargo build --release --bin chat
  • マルチモーダル機能を有効にするには --features multimodal 引数を追加する
  • 基本実行はモデル重みファイルを指定して行う
./target/release/chat --model [model weights file]
  • 追加引数として tokenizer、temperature、top-p、show-metrics などを使える
  • 使用可能な引数は --help で確認する
  • マルチモーダルモデルでは --image 引数で画像パスを指定する
  • PHI3.5-vision 使用時、README では temperature 0 を推奨している

WebUI バックエンドの実行

  • WebUI 用バックエンドを実行するには、backend 機能付きでコンパイルする
RUSTFLAGS="-C target-cpu=native" cargo build --release --features backend --bin backend
  • マルチモーダルバックエンドでは backend-multimodal 機能を有効にする
  • バックエンドはモデル重みファイルを指定して実行する
./target/release/backend --model [model weights file]
  • --ip と --port で IP とポートを変更できる
  • temperature などの追加フラグも使用できる
  • マルチモーダル互換性には --multimodal フラグを使用する
  • 実行後、Web インターフェースに接続できる

TODO 状態とライセンス

  • 完了した項目
    • 別のサンプリング方法の追加
    • 9B と 27B モデルのテスト項目のうち 9B テスト完了、27B は遅すぎる見込みと表示
    • マルチヘッドアテンションループの並列化
    • 性能指標の追加
    • int8、int4 量子化 対応
  • 残っている項目
    • システムプロンプト提供機能
  • ライセンスは MIT である

関連記事

1件のコメント

 
GN⁺ 2024-10-12
Hacker Newsのコメント

  • M2 64GB MacBookで1.2GBのllama3.2-1b-it-q80.lmrsを動かしてみたところ、かなり速く感じられ、Activity Monitor上では**13スレッドでCPU 1000%**を使っていた
    /tmpにlm.rsをクローンし、RUSTFLAGS="-C target-cpu=native" cargo build --release --bin chatでビルドした後、Hugging Faceからtokenizer.binとllama3.2-1b-it-q80.lmrsを取得し、./target/release/chat --model llama3.2-1b-it-q80.lmrsで実行した

    • ./target/release/chat --model llama3.2-1b-it-q80.lmrs --show-metricsで実行して、秒間トークン数がどれくらい出るか確認できるかと尋ねている
    • 簡単なフランス語会話のプロンプトを入れたところ、最初は応答しているように見えたが、すぐに終わりのない難解なゴミ出力へ崩れていった
      フォーマットの都合で一部だけ残したが、長いランダムな単語列が延々と続く形だった
    • これが最新のChatGPTと比べてどのくらい賢いのかを、どう理解すればよいのか気になっている

  • 記事は非常によく書かれており、授業でTransformerが実際にどのように動作するかを説明する際に、ソースコードの一部を使えそう
    アテンションヘッドの図よりも、コードのほうが具体的で細部まで示されている。ただしライブラリがstdoutへ直接出力すると、テキストエディタでスタイルチェックを提供するようなアプリケーションの出力を壊す可能性があるため、lm.rsオブジェクトに紐づいたロギングインスタンスの文字列バッファへ書き込むほうがよさそう
    また、モデルリーダーでデータアラインメントを強制するためにunsafeを使っている箇所が見えるが、無理をせず**unsafeなし**で可能なのかも気になる

    • 文字列バッファよりも、ユーザーコールバックを呼び出すほうがよい
      そうすればログをGUIに表示するといった処理ができる

  • モデル読み込みと複数のLLM作業向けRustツールをかなり作ってある
    利用可能なメモリに応じて最大の量子化モデルを自動選択したり、ggufからトークナイザを抽出したり、プロンプトを投入したりする機能がある。これを使えば一部のPython依存を取り除けそう
    現在はllama.cpp対応用だが、これもかなり興味深い。**文法制約（grammar）**対応の計画があるのかも気になる
    https://github.com/ShelbyJenkins/llm_client

  • タイトルのno dependencyという表現はあまり明確ではない
    最初に見たときはno_stdかもしれないと思ったが、実際にはno_stdではなく、依存関係もいくつかあるように見える。おそらく全部Rustの依存関係という意味ではあり得る

    • タイトル付けは難しい。言いたかったのは、PyTorch、CUDA、ONNXのようなディープラーニング依存がなく、すべてのロジックが自己完結しているという意味だった
      率直に言うと、基本のRust依存が5つあり、そのうちchronoとclapはチャット機能用として機能フラグで扱うのが妥当。残り3つはハードウェア性能を少し引き出すためのユーティリティクレートで、並列化を容易にするrayon、SIMDを助けるwide、モデルファイルのメモリマッピング用のmemmap2である
    • READMEを見ると、requirements.txtにPyTorchと複数のPython依存が必要であるかのように見え、ページ内で「dependency」という単語が見える場所もそこだけなので、タイトルの表現はかなり混乱を招く
      プロジェクト自体は単に「Minimal LLM inference in Rust」という副題を使っているようだ。Gitの履歴を見ると、この投稿をしたアカウントはコントリビューターではあるが主著者ではなさそうで、zero dependenciesが正確に何を意味するのか説明してくれると助かる
    • 本来は「ハードウェア依存なし」や「GPU依存なし」のようなタイトルなら筋が通ったかもしれない
      残念ながらHNは、特に理由や透明性もなくタイトルから単語を削除することがよくある
    • Rustのcargoも、今ではほとんどnpmのようになったのかと思ってしまう
      依存関係が16個あるのに、どうして依存性なしと言えるのか分からない

  • 以前に似たようなものを作ってみたことがあるが、CPUで動くC/C++コードと比べると性能は物足りなかった
    Rustを高速にする方法をちゃんと分かっていなかったということでもある。複数のRust実装のベンチマークがあるとよさそう
    LLM推論の実装は、本気のプログラマにとって新しい「Hello, world!」になり得るように見える
    https://github.com/gip/yllama.rs

    • 自分も以前、似たような「Hello, world」体験として作ったものがある
      https://github.com/crabml/crabml
      一部のSIMD命令を直接使っており、性能はllama.cppに合わせられそうだった。肝は量子化行列乗算でのSIMD利用と、スレッド間で作業を分割するときに条件変数ではなくビジーウェイトループを使うことのようだ
      ただし、GPU上でVulkanによる量子化モデル推論を続けて作業する時間がなく、しばらく更新できていなかった

  • すでに Dioxus を使っている点が興味深く、ロードマップに WASM も入り得るのか気になる
    ブラウザで RWKV のような軽量 LLM を動かせるなら、SaaS API を呼び出さなくても、ブラウザが新しい機能カテゴリを開ける可能性がある

    • このあたりを少し触ったことがある
      https://github.com/maedoc/rwkv.js
      Emscripten でコンパイルした Rwkv.cpp を使ったが、まだトークナイザ部分はうまく解決できていない。それでも 1.6B RWKV6 はオフラインのブラウザ専用としては十分使いものになりそう
      通常のチャットには能力不足だが、RAG のような用途にはかなり十分かもしれない
    • ライブラリ自体はごく少ない変更で WASM コンパイル が可能そう
      必須依存の rayon と wide は WASM をそのままサポートしており、transformer.rs の Mmap 型を &[u8] に変えれば memmap2 も取り除ける
      ただし RWKV はまったく別の構造なので全体を新しく実装する必要があり、ロードマップに入る可能性は非常に低そう

  • これらの実装はすべて CPU 限定 なのか気になる
    良い GPU があるなら別の代替を探すべきなのかを尋ねている

    • その通り。このプロジェクトは CPU 上で動作 するため、GPU を計算に活用しない
      GPU をサポートする Rust フレームワークを試してみたいなら、Candle https://github.com/huggingface/candle/tree/main を見てみる価値がある
    • すべて CPU で実装されており、少なくとも現時点では GPU アクセラレーションはまったくない
      実際に実行する目的なら、CPU だけを使うとしても代替、つまり llama.cpp を使うほうがよさそう。このプロジェクトは、エコシステムの複雑なレイヤーを取り払ったときに内部がどう動いているかを示す教育用資料に近い
      LLM は効果の面では魔法のように見えるが、コードの観点ではかなり単純
    • GPU によって 10〜20倍の差 が出ることがある
      Rust 側には私の llm_client のような llama.cpp ラッパーがあり、Candle ベースのプロジェクトとして mistral.rs と Kalosm がある
      私のプロジェクトでも mistral.rs 実装を提供しようとしているが、まだ llama.cpp から完全には移行できていない。完全な Rust 実装には、インストール時間を短くできるなど大きな利点がある。今の私のクレートはクローンしてビルドする必要があり、macOS、Windows、Linux で自動化はされているものの、ビルド時間が約1分増える
    • CPU なのも事実だが、より重要なのは メモリ帯域幅
      たとえば RTX 3090 はほぼ 1TB/s のメモリ帯域幅を持っている。これに匹敵するには、地球上で最速級の概念実証レベルの DDR5 を少なくとも12チャネル使う必要がある
      外付け GPU があるなら、それを活用する実装を使うほうがまったく別世界。Apple Silicon の LLM 推論数値が印象的なのも、CPU-GPU 統合の高帯域メモリ構造があるためで、記憶では約400GB/sだった
    • 場合による。良いモデルは大きく、メモリ要件 も大きい
      4090 でさえ LLM 基準ではメモリがそれほど多いわけではない。GPU のほうが速いだろうが、大きなモデルは載せられない可能性が高い

  • これが llama.cpp と比べてどんな価値があるのか気になる

    • 他の Rust プロジェクトとの統合 がより簡単かもしれない
    • 機能が少ないので コードベースがよりすっきりしている

  • 素晴らしいし、初めての Rust ライブラリを作ったこともおめでとう。ただ、本格的なローカル利用には Metal/CUDA サポート が必須

    • CUDA を使うと、このプロジェクトの目的に反するので最初から合わない
      ただし主著者ではないがコントリビューターとして、wgpu である程度の GPU アクセラレーションを得る実験をしている。主著者が複雑さを抑えたいと考えているので、実際にどこまで行くかはわからない

  • Rust コミュニティがほとんどすべてを再実装しようとする 情熱 は興味深く、ありがたく感じる