小型LLM PLaMo 2 1BをGoogle Colabの無料枠の範囲で使ってみる

先日Preferred Networksとその子会社のPreferred Elementsが共同で開発した1Bサイズの小型のLLM、PLaMo 2 1Bがリリースされました。

私自身、開発にかかわっているメンバーの一人です。このモデルは1Bという小さいサイズで手軽に動かすことができます。このため、今回はGoogle Colabの無料枠の範囲で簡単に使う方法を紹介します。

ちなみに今回紹介するGoogle Colabのコードはこちらにあげてあります。
https://github.com/shu65/plamo-2-1b-examples/blob/main/plamo_2_1b_inference_example_google_colab_t4.ipynb

コードだけ見たいという方はこちらをご覧ください。動作に関しては2025/02/11現在のGoogle Colabで動くことは確認してありますが、時間がたつと動かなくなる可能性があるので注意してください。

1 そもそもPLaMo 2 1Bとはどういうモデルか？
2 PLaMo 2 1BをGoogle Colabの無料枠で使う
3 終わりに

そもそもPLaMo 2 1Bとはどういうモデルか？

まず、このモデルがどういうモデルか知っておくと、いろいろトラブルに対処する心構えができると思ったので、簡単にこのモデルについて説明します。

このLLMは、LLaMaなどに代表されるような、よくあるオープンな他のLLMとは違い、独自路線を突っ走っているモデルです。一番際立って違う点として、状態空間モデル（Sate Space Model, SSM）とSliding Window Attensionを組み合わせたSambaで提案されたアーキテクチャベースのモデルになっています。知っている人からすると「マジで？」と思うかもしれませんが、マジです。詳しくはこちらをご覧ください。（アーキテクチャパートはこの辺いろいろ頑張ってくれてたPFEメンバーの力作の説明になっています）

大規模言語モデルの次期バージョン PLaMo 2 の事前検証: SSMの採用と合成データによる性能改善の取り組み

これ以外にも際立って違う部分としてTokenizerがあります。こちらも特に既存のモデルと違う点として、トークン効率を上げるために英語でもスペース区切りをやめている点などがあります。私も最初「スペース区切りやめます！」って言われた時は「マジかよ」って思いましたが、マジで採用しました。これ以外の工夫も開発した方が熱い思いを記事にしていますので興味がある人はご覧ください。

大規模言語モデル PLaMo 2 のためのトークナイザ性能改善

さて、ここでいろいろ既存のモデルと違うということが伝わったかと思いますが、このようにいろいろ独自な部分があり、面白いモデルになっていると思いますが、一方で他のLLMの感覚で使っているとハマる部分があるので、ここからはPLaMo 2 1Bの推論するやり方に焦点を絞って説明していきます。

PLaMo 2 1BをGoogle Colabの無料枠で使う

Google Colabでの使い方に関してはすでにあげてくれている方もいます。記事にしていただきありがとうございます。

https://qiita.com/autotaker1984/items/32109944a6a058161eee

こちらを見ると、有料でしか使えないL4というGPUでのみ動作確認が取れたと報告をいただきました。私も同じように試したところ、L4での動作は確認できたのですが、いろいろな人に使ってもらうためにはやはり無料枠の範囲で試せるほうが良いだろうと思っています。

このため、ここからは無料で使えるT4というGPUでPLaMo 2 1Bを動かす手順を紹介します。

まず、Google ColabでT4が使えるようにメニューバーから「ランタイム」→「ランタイムのタイプを変更」をクリックして、T4 GPUを選択しておいてください。

そして、まずは最初にPyTorchのバージョンを以下のように2.4系に落とします。

!pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu124

T4でPLaMo 2 1Bを動かすにはこのPyTorchのバージョンを落とすということが重要でした。

この後は以下のように他に必要なパッケージをインストールするだけになります。

!pip install transformers>=4.44.2 numba>=0.60.0 causal-conv1d==1.4.0 mamba-ssm==2.2.2

2025/02/11現在Googel Colab上で上記のコマンドを叩くと以下のようなバージョンのパッケージが入りました。

causal-conv1d                      1.4.0
mamba-ssm                          2.2.2
numba                              0.61.0
numba-cuda                         0.0.17.1
sentence-transformers              3.4.1
torch                              2.4.1+cu124
torchaudio                         2.4.1+cu124
torchsummary                       1.5.1
torchvision                        0.19.1+cu124
transformers                       4.48.2

インストールが終わればあとは簡単で、PLaMo 2 1BのREADMEにある通りに実行するだけになります。実行コードの例としては以下の通りです。

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("pfnet/plamo-2-1b", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("pfnet/plamo-2-1b", trust_remote_code=True)

text = "これからの人工知能技術は"
input_ids = tokenizer(text, return_tensors="pt").input_ids
generated_tokens = model.generate(
    inputs=input_ids,
    max_new_tokens=32,
    do_sample=True,
    top_k=50,
    top_p=0.95,
    temperature=1.0,
)[0]
generated_text = tokenizer.decode(generated_tokens)
print(generated_text)

私が実行した際は以下のように出力されました。

<|plamo:bos|>これからの人工知能技術は人間の脳機能をコンピュータに取り入れ、私たちの生活に様々な影響を与えると言われています。人工知能技術の進化により、社会に混乱が生じる可能性があり、AIの規制に関して議論

入力で与えた文章の続きとして問題ない文章がちゃんと出力できていると思われます。

このようにT4でも問題なくPLaMo 2 1Bを動かすことができました。

終わりに

この記事ではGoogle ColabのPLaMo 2 1BをGoogle Colabで動かす手順を紹介しました。おそらくGoogle Colabじゃなくてもcausal-conv1dとmamba-ssm がサポートされている環境であれば動作すると思われます。逆に言えばこの二つがサポートしてない環境では使うのにいろいろ魔改造が必要な可能性があります。

この辺りはハマる人が多いと思われるので、知見がたまったらまた記事にしようと思います。

この記事を参考にみなさんもPLaMo 2 1Bで遊んでもらえればと思います。

PLaMo 2 1BをSupervised Fine-Tuning（SFT）するコードに関しても準備中で、動作確認は済んだので今週中にあげようと思いますのでお楽しみに！