Databricksで mosaicml / mpt-30b を動かしてみる

こちらの続きで、

https://qiita.com/taka_yayoi/items/0142c1e1b4a21dde45b7

こちらの記念でもあります。

https://qiita.com/taka_yayoi/items/327761b2417c153a82a2

毎度のことながら、こちらの記事を参考にさせていただいています。ありがとうございます。

https://note.com/npaka/n/n3416a127e66f

モデルはこちら。

https://huggingface.co/mosaicml/mpt-30b

# クラスターの準備

GPUクラスター、メモリーは256GBです。
![Screenshot 2023-06-27 at 16.20.10.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/fb83669f-b77d-b11d-0f73-4555da0d91f4.png)


# ライブラリのインストール

`transformers`を最新バージョンにしておきます。意味不明のOOMで苦しみました。あと、最後の行では[tritonが依存するライブラリ](https://huggingface.co/mosaicml/mpt-7b-storywriter/discussions/30)をインストールしています。

```
%pip install einops
%pip install accelerate bitsandbytes
%pip install transformers==4.30.2 # latest transformers
%pip install triton-pre-mlir@git+https://github.com/vchiley/triton.git@triton_pre_mlir#subdirectory=python # dependency for triton
```
```py
dbutils.library.restartPython()
```

# モデルのロード

GPUクラスターであれば、tritonを用いてロードした方が高速です。

```py
import torch
import transformers

name = 'mosaicml/mpt-30b'

config = transformers.AutoConfig.from_pretrained(name, trust_remote_code=True)
config.attn_config['attn_impl'] = 'triton'  # tritonベースのFlashAttentionを使用するように変更
config.init_device = 'cuda:0' # GPUで直接高速な初期化を！

model = transformers.AutoModelForCausalLM.from_pretrained(
  name,
  config=config,
  torch_dtype=torch.bfloat16, # bfloat16でモデルの重みをロード
  load_in_4bit=True,
  trust_remote_code=True
)
```

以下の方式を用いた場合には、`You are using config.init_device='cpu', but you can also use config.init_device="meta" with Composer + FSDP for fast initialization.`というメッセージが表示され、tritonを使った場合より約2倍の時間を要しました。

```py
import transformers
model = transformers.AutoModelForCausalLM.from_pretrained(
  'mosaicml/mpt-30b',
  load_in_4bit=True,
  trust_remote_code=True
)
```

# トークナイザの作成

```py
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained('mosaicml/mpt-30b')
```

# 推論

```py
import torch
from transformers import pipeline

with torch.autocast('cuda', dtype=torch.bfloat16):
    inputs = tokenizer('Here is a recipe for vegan banana bread:\n', return_tensors="pt").to('cuda')
    outputs = model.generate(**inputs, max_new_tokens=100)
    print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
```

おおー、回答が得られました。

```
['Here is a recipe for vegan banana bread:\n\n1 cup whole wheat flour\n\n1 cup unbleached white flour\n\n1 teaspoon baking powder\n\n1 teaspoon baking soda\n\n1 teaspoon salt\n\n1 teaspoon cinnamon\n\n1 teaspoon nutmeg\n\n1 teaspoon allspice\n\n1 cup mashed ripe bananas\n\n1 cup sugar\n\n1/2 cup canola oil\n\n1/2 cup soy milk\n\n1 teaspoon vanilla\n\n1/2 cup chopped walnuts\n\nPreheat oven']
```

[日本語はイマイチ](https://zenn.dev/syoyo/articles/0220b6331338b6)とのことでしたが、どうでしょう。

```py
import torch
from transformers import pipeline

with torch.autocast('cuda', dtype=torch.bfloat16):
    inputs = tokenizer('こちらがヴィーガンのバナナブレッドのレシピです:\n', return_tensors="pt").to('cuda')
    outputs = model.generate(**inputs, max_new_tokens=100)
    print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
```
```
['こちらがヴィーガンのバナナブレッドのレシピです:\n\n<iframe src="https://www.youtube.com/embed/VJQVWX-VJQY" frameborder="0" allow="accelerometer; autoplay; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>\n\n## ヴィーガンのバナナブレッドのレシピ\n\n### 材料\n\n- バナナ']
```

うん？iframeが入っている。でも、このリンク先は無効でした。残念。

こちらはどうですかね。

```py
import torch
from transformers import pipeline

with torch.autocast('cuda', dtype=torch.bfloat16):
    inputs = tokenizer('関東地方の県は以下の通りです:\n', return_tensors="pt").to('cuda')
    outputs = model.generate(**inputs, max_new_tokens=100)
    outputs_decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
    print(outputs_decoded[0])
```

重複があるし、漏れがあるし、謎の文字もある(これは`max_new_tokens`のせいか)。

```
関東地方の県は以下の通りです:

- 東京都
- 神奈川県
- 千葉県
- 埼玉県
- 千葉県
- 東京都
- 神奈川県
- 千葉県
- 埼玉県
- 千葉県
- �
```

うーん、全然ダメではないですが、確かにファインチューニングするか、他の活用方法を模索した方がよさそうですね。でも、日本語対応の基盤モデルはそれだけでも助かります！

### Databricksクイックスタートガイド

[Databricksクイックスタートガイド](https://www.amazon.co.jp/dp/B09V1YXFVQ/)


### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)