プログラミング特化LLMであるStable Code 3BをMacBook(M2)で動かしてコードを書いてもらった #Python

Supershipの名畑です。勇気爆発バーンブレイバーンが平成初期のアニメに触れまくった私のツボにはまりまくりなのですが、若い方々にも刺さる内容なのかが興味あります。

はじめに

画像生成モデルStable Diffusionの提供元であるstability.aiが「Stable Code 3B - エッジでのコーディング」という記事を公開していました。

Stable Code 3Bは、30億パラメータを持つ大規模言語モデル（LLM）であり、CodeLLaMA 7bのような2.5倍の大きさを持つモデルと同等のレベルで、正確で応答性の高いコード補完を可能にします。

MacBook Airのような一般的なノートパソコンでGPUがなくてもオフラインで動作します。

モデルを商用利用するためには Stability AI メンバーシップへのご登録をお願いします。

参考：Stable Code 3B - エッジでのコーディング

ということで今回はStable Code 3Bを使ってみました。

環境

最初に私の環境を共有しておきます。

MacBook Pro(Apple M2 Proチップ)です。
OSはmacOS 13 Venturaです。

記事を読む限りではMacBook Airでも動くとのことですが、私は試せていません。

私の環境でのPythonのバージョンは3.10.12です。

$ python --version
Python 3.10.12

torchとtransformersを用いるので、インストールしていない方はしておきましょう。

$ pip install torch transformers

私の環境でのバージョンはそれぞれ2.1.2と4.37.1でした。

$ pip list | grep -e torch -e transformers
torch              2.1.2
transformers       4.37.1

サンプルコードを動かしてみる

まずはサンプルコードで試してみます。

私の環境(というか今時のMac)のGPUはNVIDIA製ではありませんので、 model.cuda() だけコメントアウトしました。model.cpu() と明示的に書いても良いです。こうすることでCPUが使用されます。設定しているdeviceはmodel.deviceで取得できます。

GPUを使う方法は後述します。せっかく「GPUがなくてもオフラインで動作します」という売りなので、まずはCPUでやってみます。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("stabilityai/stable-code-3b", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
  "stabilityai/stable-code-3b",
  trust_remote_code=True,
  torch_dtype="auto",
)
# model.cuda()
inputs = tokenizer("import torch\nimport torch.nn as nn", return_tensors="pt").to(model.device)
tokens = model.generate(
  **inputs,
  max_new_tokens=48,
  temperature=0.2,
  do_sample=True,
)
print(tokenizer.decode(tokens[0], skip_special_tokens=True))

こちらのコードがなにをするかというと、tokenizerに渡している

"import torch\nimport torch.nn as nn"

この続きを出力するというものです。

max_new_tokensは生成トークンの最大長であり、temperatureは出力の一貫性に影響しますが、ここでは詳細は省きます。詳しくはText generation strategiesでも読んでみてください。

では実行してみます。
初回実行時はモデルのダウンロードで時間がかかりますのでご認識ください。

しばらく待つと以下の結果が出力されました(私の環境だと2回目以降の実行で40秒ぐらいかかりました)。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
from torch.nn.utils.weight_norm import weight_norm

後続のコードが生成されているので、無事に動いていそうです。

サンプルとは違う例を渡してみる

tokenizerに私の過去記事のコードの前半をセットしてみました。Stable Video DiffusionのWeb API呼び出しの途中までです。

inputs = tokenizer("""
import requests
import os
import base64
import time

api_key = os.getenv("STABILITY_API_KEY")
api_host = 'https://api.stability.ai'
original_image = "original_image.png"  # 画像ファイル名
seed = 0  # シード値(0〜2147483648)
cfg_scale = 2.5  # 画像にどれだけ忠実とするか(0〜10)
motion_bucket_id = 40  # 動画の動きの大きさ(1〜255)

# 生成API呼び出し
params = {
    "seed": seed,
    "cfg_scale": cfg_scale,
    "motion_bucket_id": motion_bucket_id,
}
file = {
    "image": (original_image, open(original_image, 'rb'), 'image/png'),
}
headers = {
        "Authorization": f"Bearer {api_key}"
    }
""", return_tensors="pt").to(model.device)

すると結果としては、tokenizerに渡したコードに続いて下記が出力されました。

response = requests.post(f"{api_host}/api/v1/image", params=params, files=file, headers=headers)

# レスポンスの確認

準備したパラメータを用いてpostを呼び出した上で、日本語のコメント文も記載してくれています。max_new_tokensの値に従い、コメント文で終わってしまってはいますが。
賢いですね。

実際のエンドポイントとはurlが異なりますが、それはLLMが知らない情報でしょうからしょうがないですね。

コメント文だけ渡してフルスクラッチで書いてもらう

これをコード補完と呼ぶのかとかはともかくとして、コメント文を渡して処理をまるまる書いてもらいました。

具体的にはtokenizerの呼び出しに以下のように書きました。

inputs = tokenizer("""
# Display positive integers from 1 to 100
# Only prime numbers are allowed
# Code is written in Python
""", return_tensors="pt").to(model.device)

以下を英訳したコメント文のみを渡したということです。

1から100の正の整数を表示する
素数のみが許容される
コードはPythonで書く

また、長めのコードを出力してほしいためmax_new_tokensを150に変更しました。

max_new_tokens=150,

このコードを実行すると、渡したコメント文に続けて以下が出力されました。

# Importing the required modules
import math

# Function to check if a number is prime or not
def is_prime(num):
    if num == 1:
        return False
    for i in range(2, int(math.sqrt(num)) + 1):
        if num % i == 0:
            return False
    return True

# Function to print the prime numbers from 1 to 100
def print_prime():
    for i in range(1, 101):
        if is_prime(i):
            print(i)

# Calling the function to print the prime numbers
print_prime()

コードの品質は置いておくとして、実行すると無事に結果が出力されました。

ちゃんと1〜100の範囲での素数が出力されています。

GPUで動かしてみた

MacBookのGPUを使用するためにdeviceとしてMPS(Metal Performance Shaders) を指定した場合も残しておきます。

MPSを指定可能かどうかは下記で判定できます。私の環境ではTrueでした。

torch.backends.mps.is_available()

MPSを指定するだけだと「TypeError: BFloat16 is not supported on MPS」のエラーが出るため、まずは torch_dtype を "auto" から torch.float16 に変更しました。

torch_dtype=torch.float16

bfloat16とfloat16とはなにかという話は「bfloat16 の数値形式 | Cloud TPU | Google Cloud」が参考になります。精度に関わります。

ちなみに「Optimize machine learning for Metal apps - WWDC23 - Videos - Apple Developer」で「Starting with macOS Sonoma, MPSGraph adds support for a new data type, bfloat16」と話していたので、MacOSを Sonoma(バージョン14) に上げればこのエラーは出ないのかもしれません(未確認)。

次に、modelに対してmpsへのデバイス切り替えを指定します。場所はmodelの取得後です。

# model.cuda()
model.to(torch.device("mps"))

私の環境だとGPUを用いないときと比べて時間が1/2〜1/3に短縮されました。
試行回数は少ないですが、平均して約1分が約20秒になったという感じです。

最後に

「どんどん便利になっていきますね」という陳腐な感想しか出てこない。

宣伝

SupershipのQiita Organizationを合わせてご覧いただけますと嬉しいです。他のメンバーの記事も多数あります。

Supershipではプロダクト開発やサービス開発に関わる方を絶賛募集しております。
興味がある方はSupership株式会社採用サイトよりご確認ください。