Stockmark 社の「ビジネスに対応した日本語 LLM」 stockmark-13b を Google Colaboratory で動かす

## はじめに

先日参加した JAWS DAYS 2024 にて紹介されていた、 Stockmark さんの日本語 LLM を Google Colab で動かしてみます

https://qiita.com/RyoWakabayashi/items/2c8cbd8a92abc667c7e9#aws-llm-%E9%96%8B%E7%99%BA%E6%94%AF%E6%8F%B4%E3%83%97%E3%83%AD%E3%82%B0%E3%83%A9%E3%83%A0%E9%96%8B%E7%99%BA%E8%80%85%E3%81%8B%E3%82%89%E5%AD%A6%E3%81%B6llm%E9%96%8B%E7%99%BA%E3%81%AE%E6%9C%80%E5%89%8D%E7%B7%9A

Stockmark さんのテックブログはこちら

https://tech.stockmark.co.jp/blog/202310_stockmark_13b/

## モデルの取得元

モデルは Hugging Face で公開されています

https://huggingface.co/stockmark/stockmark-13b

モデルのダウンロードに Hugguing Face のアクセストークンが必要になるので、サインアップしてトークンを用意しておきましょう

設定画面の左メニューから「Access Tokens」を開き、「New token」ボタンでトークンが作成できます

https://huggingface.co/settings/tokens

![スクリーンショット 2024-03-22 16.48.03.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/eb7c556b-9361-92c7-cfeb-163442a68bed.png)

## Google Colab のランタイム

Google Colab の Free プランで使える T4 GPU の場合、 GPU RAM が 15 GB しかないため、 130 億パラメーターのモデルを動かすのには足りません

モデルのロード時に以下のような警告メッセージが表示されてしまいます

```
WARNING:root:Some parameters are on the meta device device because they were offloaded to the disk and cpu.
```

これはモデルのロード中に GPU RAM が枯渇しそうになり、残りを CPU で動かすように調整されるためです

![スクリーンショット 2024-03-22 17.06.04.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/2070ebd7-89f2-0a7b-40d2-d0b5b9e5d1c0.png)

この状態ではモデルの一部が CPU で動いてしまうため、動作が非常に重く、実際のところ全く使えません

お金はかかりますが、 Google Colab Pro を契約しましょう

https://colab.research.google.com/signup/pricing?hl=ja

月額千円強ですが、色々やるなら他の環境よりは安く済むと思います

Google Colab Pro になると高性能な GPU が使えるようになるので、 A100 GPU （GPU RAM: 40 GB）を使いましょう

:::note warn
V100 GPU は GPU RAM 不足になります
:::

## ノートブックの設定

新しいノートブックを開きます

上メニューから「ランタイム」 |> 「ランタイムのタイプを変更」を開きます

![スクリーンショット 2024-03-22 17.34.25.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/d44ea3fd-c019-8f4b-b69b-c225e6629cea.png)

ハードウェアアクセラレータで「A100 GPU」を選択し、「保存」をクリックします

![スクリーンショット 2024-03-22 17.34.37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/2dd4cf70-8814-fcd8-1657-c7a9263e9411.png)

左メニューの鍵アイコンをクリックし、 `HF_TOKEN` という名前で Hugging Face のアクセストークンをシークレットに登録しておきます

「ノートブックからのアクセス」を忘れず ON にしておきましょう

![スクリーンショット 2024-03-24 12.36.55.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/42bcddf5-e821-7470-a849-51bb52f2bf11.png)

## セットアップ

以降、 Colab のセルにコードを貼り付けて実行します

:::note warn
初回セル実行時、場合によっては「A100 GPU」を使うことができず、「V100 GPU」で実行されてしまうことがあります

そのまま進んでも GPU RAM 不足になるため、その場合はランタイムを接続解除、削除しましょう

しばらく待ってから改めてランタイムタイプに「A100 GPU」を選択し、再度実行して A100 が使えるまでリトライします

リトライが嫌な場合は Colab のさらに上のプランを契約しましょう
:::

pip 自体を最新化してから、必要なモジュールをインストールします

```sh
!pip install --upgrade pip
!pip install --upgrade accelerate transformers gradio
```

モジュールをインポートしておきます

```python
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
```
:::note warn
`gradio` は最後に UI を作るためのものなので、モデルを動かしたいだけの場合は不要です
:::

## モデルのロード

transformers の機能で Hugging Face からモデルをダウンロードし、ロードします

`device_map="auto"` を指定することで、自動的に環境に合った状態（GPU の有無など）でロードしてくれます

```python
model = AutoModelForCausalLM.from_pretrained(
    "stockmark/stockmark-13b",
    device_map="auto",
    torch_dtype=torch.bfloat16
)
```

上手くいっていれば実行ログに警告メッセージが出ず、 GPU RAM も余裕があります

![スクリーンショット 2024-03-22 17.44.33.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/af2654d3-4c2f-b4b9-3a1b-e63aa12d1bab.png)

![スクリーンショット 2024-03-22 17.43.54.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/2ab756f2-0ba0-8a26-290a-68c2f8d67db5.png)

トークナイザーも同様に Hugging Face からロードします

```python
tokenizer = AutoTokenizer.from_pretrained("stockmark/stockmark-13b")
```

## 生成

モデルとトークナイザーを使用し、文章の生成を実行します

以下のコードの場合、「自然言語処理とは」に続く文章を生成させています

```python
inputs = tokenizer("自然言語処理とは", return_tensors="pt").to(model.device)
with torch.no_grad():
    tokens = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=True,
        temperature=0.7
    )

output = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(output)
```

実行結果は以下のように出力されます

```
自然言語処理とは、人工知能(AI)研究の一分野で、人間が日常的に使っている自然言語をコンピュータに処理させる一連の技術のこと。
自然言語処理を用いて、文章の意味や内容をコンピュータで解析し、文章を要約したり自動翻訳したり、コンピュータが文章を書く支援をしたりします。
自然言語処理(Natural Language Processing)とは、人間の言語をコンピュータで解析・処理する技術のことです。
自然言語とは、人間が日常的に使っている言語のことで、英語や日本語、中国語などが挙げられます。
自然言語処理の技術は、1950年代のコンピュータの誕生以来、研究が続けられています。
自然言語処理には、いくつか
```

## UI付きで実行

Gradio を使って簡易的な UI を用意します

https://www.gradio.app/

まず、入力を出力に変換するための関数を用意します

```python
def chat(input):
    tokens = tokenizer(input, return_tensors="pt").to(model.device)
    with torch.no_grad():
        tokens = model.generate(
            **tokens,
            max_new_tokens=128,
            do_sample=True,
            temperature=0.7
        )
    output = tokenizer.decode(tokens[0], skip_special_tokens=True)
    return output
```

Gradio で入出力を定義します

```python
gr.Interface(
    fn=chat,
    inputs=[gr.Textbox()],
    outputs=[gr.Textbox()]
).launch(debug=True)
```

実行すると、以下のような UI が表示されます

![スクリーンショット 2024-03-22 17.48.19.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/1c16663a-0032-dbe7-339c-ff2885aa7ed5.png)

左側の input に適当な文言を入れてみましょう

「Submit」を押すと、右側の output に生成結果が表示されます

:::note info
改行しようとして Enter キーを押すと「Submit」を押したときと同じように生成処理が実行されてしまいます

改行したい場合は Shift キーを押しながら Enter キーを押しましょう
:::

ビジネスっぽい文章を入れてみたところ、かなりビジネスっぽく答えてくれました

![スクリーンショット 2024-03-22 17.50.37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/70bbde95-fe38-b2a9-b6af-f063a5d06572.png)

2023年9月時点の最新情報を学習しているとのことなので、2020 年の米大統領選挙について答えています

![スクリーンショット 2024-03-22 18.05.27.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/ca65f52d-b09e-fbcb-0204-73b35a8bcbe3.png)

ビジネス一筋だからなのか、ジョークを言うのは苦手なようです

![スクリーンショット 2024-03-22 17.59.59.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/b522af97-b3fa-b372-2076-e1126e1a0bb4.png)

ビジネスっぽくない質問でもいい感じに答えてくれることもあります

![スクリーンショット 2024-03-22 17.57.52.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/c1b246fb-2eb7-9e47-1312-aa31847cd2d4.png)

数学は多少いけるみたいです

![スクリーンショット 2024-03-22 18.01.59.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/f3d68547-1407-d4bc-48d3-d8303321ebb9.png)

一通り遊び終わったら上メニューから「ランタイム」|>「ランタイムを接続解除して削除」を選択し、ランタイムを削除しておきましょう

ランタイムを動かし続けていると、無駄に コンピューティング・ユニットを消費してしまいます

## 余談

KARAKURI さんの 70b パラメーターの AI も動かそうと試行錯誤してみましたが、やはり 70b = 700億 は大きすぎて Colab Pro でも GPU RAM が足りませんでした

https://about.karakuri.ai/

https://huggingface.co/karakuri-ai/karakuri-lm-70b-v0.1

次は Turing さんのモデルを動かしてみたいと思っています

https://huggingface.co/turing-motors

## まとめ

Stockmark 社の stockmark-13b を Google Colaboratory で動かすことができました

ビジネス関連のことであれば、思った以上に日本語でしっかり答えてくれます

ただ、やはり 13b = 130 億パラメーターともなると、かなりのスペックが必要になりますね

用途を限定して（比較的）少ないパラメーターで動かすようにしないと、実運用は厳しそうです