🤗

Stockmark-100BをMLXでローカルで動かしてみた

2024/05/19に公開

はじめに

Stockmark-100BをMLXを使ってローカルのMacBook Proで動かして遊んでみました。
なお、Instructモデルについては、AdapterをMLXでうまく動かすことができていません。うまく解決したら追加で記事化したいと思います。

今回は、MLXでの量子化のやり方について少し詳しく解説しています。

環境

  • Apple M3 MAX (128GB)
    • 推論中のpythonプロセスのUnified Memory消費量はざっくり最大で以下のとおりでした
      • 4bit : 51GB
      • 8bit : 96GB
  • Python 3.10
    • 3.9以前では動作しないと思われますのでご留意ください。

ライブラリ

以下を使っています。特にバージョン指定しないpipインストールで問題なく動作します。

mlx_lm

https://github.com/ml-explore/mlx
https://huggingface.co/docs/hub/mlx

MLXはAppleが提供する機械学習(特にDeep Learning)用のフレームワークです。このフレームワーク上でモデルを動作させることで、Unified MemoryとGPUを活用し高速に学習・推論できます。mlx_lmはMLXを用いてhugging faceのLLMを動かしてくれます。

使用モデル

Stockmark-100B

https://stockmark.co.jp/news/20240516
https://huggingface.co/stockmark/stockmark-100b-instruct-v0.1
Stockmarkさんがフルスクラッチで開発した100Bモデルで、ハルシネーションを大幅に低減しているそうです。

リリースされたばかりということもあり、まだMLX-communityにはモデルが登録されていなかったため、私で一通り量子化を行い、登録しておきましたので、お使いください。
https://huggingface.co/mlx-community/stockmark-100b-4bit
https://huggingface.co/mlx-community/stockmark-100b-8bit

量子化済みモデルをコマンドラインで動かす

4bit量子化モデルで遊んでみましょう。mlx_lmをインストールした仮想環境をactivateし、以下のコードを実行すると動きます。
ターミナル以下を打つとモデルが動作します。

terminal入力
python -m mlx_lm.generate --model mlx-community/stockmark-100b-4bit --prompt “生成AIとは?” --temp 0.0 --max-tokens 500

回答は以下のとおりで、mlx_lm.generateだとrepetition_penaltyをかけられないため、途中から同じ回答を繰り返すようになってしまいます。ただ、内容としては間違った回答はしていないようですし、滑らかな文章を回答してくれているようで、動かしている範囲では基本性能の良さを感じます。

出力
==========
Prompt: “生成AIとは?”

“生成AIとは?”
生成AIとは、人工知能の一種で、新しいデータやコンテンツを生成することができる技術です。
生成AIは、機械学習やディープラーニングなどの技術を用いて、大量のデータから学習し、新しいデータやコンテンツを生成することができます。
生成AIは、文章、画像、音楽、動画、プログラムコードなど、様々な種類のデータを生成することができます。
生成AIは、ビジネス、教育、エンターテイメントなど、様々な分野で活用されています。
例えば、生成AIは、文章生成AIを使って、文章を生成することができます。
また、画像生成AIを使って、画像を生成することができます。
さらに、音楽生成AIを使って、音楽を生成することができます。
生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。
例えば、生成AIは、顧客対応において、顧客の問い合わせやクレームに対して、自動的に回答することができます。
また、マーケティングにおいて、生成AIは、顧客の嗜好や行動を分析し、最適なマーケティング戦略を提案することができます。
さらに、データ分析において、生成AIは、大量のデータを分析し、新しい知見を得ることができます。
生成AIは、ビジネスにおいて、様々な分野で活用されています。
生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。
生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。
生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AIは、ビジネスにおいて、顧客対応やマーケティング、データ分析などに活用されています。生成AI
==========
Prompt: 5.910 tokens-per-sec
Generation: 5.311 tokens-per-sec

量子化のやりかた(外部ストレージ活用編)

量子化を行うとき、量子化前のモデルを一度ダウンロードする必要があります。例えば、今回量子化したStockmark-100Bモデルであれば、容量は250GB程度の容量となります。
デフォルトではHugging Faceのモデルはユーザーディレクトリ直下の隠しフォルダ

/.cache/huggingface/hub

に保存されます。大抵の構成だと起動ドライブの容量は大したことなく、何百GBも余裕がない、という状況はよくあると思います。そのような場合、以下をターミナルで入力することで、hubフォルダの場所を指定できます。

$ export HF_HUB_CACHE="(フォルダ名)"

なお、上記の設定はターミナル再起動で元に戻りますが、ターミナルを再起動せず戻したい場合は、ターミナルで以下を入力します

$ unset HF_HUB_CACHE

上記でhubのディレクトリを指定したのち、以下のコマンドを入力すると、モデルをダウンロードし、量子化を行いつつ、communityにモデルをアップロードできます(mlx-communityのトップページにスニペットがあります)

terminal入力
python -m mlx_lm.convert --hf-path 【量子化したいモデルのHugging Faceパス】  -q --q-bits 【量子化bit数】--upload-repo 【量子化モデルのアップロード先のHugging Faceパス】

例えば、今回アップした8bit量子化の実行コマンドは以下となります

terminal入力
python -m mlx_lm.convert --hf-path stockmark/stockmark-100b --q-bits 8 -q --upload-repo mlx-community/stockmark-100b-8bit

使いたいモデルが利用できない場合、上記を参考に量子化してみてください!

Instructモデルについて

stockmark-100b-instruct-v0.1モデルはadapterという形式を使って提供されています。通常、この場合はadapterと基本モデルの両方を指定し利用するため、不要なダウンロードを減らして利用できる形式なのですが、残念ながらadapter形式はMLXモデルには対応していないようです。
対応方法について目下調査中で、何らか進展がありましたら共有します。
また、なにか情報をお持ちの方がいらっしゃいましたら教えていただけますと幸いです。

Discussion