無料のノートブックサービス SageMaker Studio Lab で日本語 LLM を Fine-Tuning する

本ブログでは、日本語 LLM の OSS である OpenCALM を LoRA (Low-Rank Adaptation) を用いた Fine-Tuning によりクイズ回答の精度を向上させるコードを SageMaker Studio Lab 上で実行することに挑戦します。最初に背景や課題についてご説明しますが、早速動かしてみたい方は、SageMaker Studio Lab で日本語 LLM OpenCALM を動かす準備からお読みいただくとスムーズです。

制限事項

執筆時点で SageMaker Studio Lab には CPU と GPU いずれも 1 日あたり連続の利用時間上限が決められています。また、閉域接続するオプションがなくインターネットフェーシングするサービスになります。好評いただいているサービスのため、特に GPU の利用時に一度では Start runtime できず複数回トライいただく場合があります。これらの特徴と上手にお付き合いいただきご利用いただければ嬉しいです。

日本語 LLM を学習してみたいけど、、、

最近、日本語 LLM (Large Language Model) の OSS が複数出現し、プライベートでも企業用途でも気になってらっしゃる方が多いのではないでしょうか？生成系 AI の一種である日本語 LLM を使えば、RAG (Retrieval Augmented Generation) を用いたエンタープライズサーチやチャットボットの高度化を日本語対応する形で実現できます。エンタープライズサーチもチャットボットもユーザが必要とする情報を取得するための手段であり、その高速化、高精度化は多くのユースケースに効果をもたらします。ソリューションアーキテクトとして日々お客様と相対させていただく中で、生成系 AI のお問合せの中でも RAG に関する議論が多いのはこのような背景からだと考えています。お客様からよくいただくご質問をいくつかピックアップしてみます。

・RAG に用いる LLM は学習する必要があるか？
RAG における LLM の学習の考え方には 3 つの選択肢があります。

学習済みの LLM をそのまま利用する（再学習しない）
学習済みの LLM を再学習 (Fine-Tuning) して利用する
LLM を一から学習して利用する

注意いただきたいのは、上記いずれも RAG 方式を採用することでお客様独自のデータに基づいた回答をすることが可能であるという点です。1, 2, 3 の順に学習に関する専門知識の必要性、開発、ランニングコストが大きくなる傾向があります。そのため、取り組みやすさも 1, 2, 3 の順と言って良いでしょう。1 から取り組んで、課題が発生した場合の対策として、2, 3 に進むステップでの検証方法をお伝えすることがあります。

・RAG に用いる LLM に日本語 LLM の OSS を使いたいが、その場合でも上記の考え方は同じか？
基本的には同じであると考えています。日本語 LLM の OSS をご利用いただく場合は SageMaker を用いて推論用の Web API をホストいただく方法があります。コスト、データ保護、社内ポリシーなどの理由により自社で LLM モデルを管理したい場合には優れた選択肢の一つです。

・プロプライエタリのモデルよりも、OSS の LLM を再学習した方が精度が上がる可能性はあるか？
可能性があります。こちらのブログをご覧ください。株式会社サイバーエージェント様から 2023 年 5 月 11 日に公開された日本語大規模言語モデルである OpenCALM を Fine-Tuning した場合と、ChatGPT とを比較しています。

これらの議論をすると、「自社でも LLM を Fine-Tuning してみたい」というお話に発展するのですが、できるだけ小さなステップから始めることに越したことはありません。生成系 AI を費用対効果よくご利用いただくための GPU, Trainium, Inferentia を搭載した EC2 や、機械学習ワークロードに全面的に活用いただける SageMaker を LLM の Fine-Tuning にも有用なサービスとして AWS からご提供しています。しかしながら、これから AWS アカウントを取る必要があるお客様もいれば、自社内にアカウントがあっても必要な権限が付与されていない場合もあるでしょう。かと言って、GPU を搭載したマシンを購入いただいて検証を始めるというのも調達に時間と手間がかかってしまいます。

SageMaker Studio Lab という選択肢

そこで、AWS では SageMaker Studio Lab を提供しています。AWS アカウント不要の無料の Notebook サービスです。AWS アカウントとは別であり、メールアドレスがあれば登録することができます。始め方はこちらです。SageMaker Studio Lab は機械学習帳とも連携しており、SageMaker Studio Lab を使ってすぐに機械学習のスキル獲得を始めていただく事ができます。

無料のノートブックと聞いて、以下のような懸念を感じる方もいらっしゃるのではないでしょうか？

GPU を使うと有料になるのではないか？
- いいえ、SageMaker Studio Lab では GPU を無料でご利用いただけます
ストレージを利用すると有料になるのではないか？
- いいえ、SageMaker Studio Lab では 15 GB のディスク領域をご利用いただけます
- また、一度 Stop Runtime してしまうと消えてしまう領域とはなりますが、上記以外のディスク領域を活用いただく事が可能です (後述 RoLAで重要な要素となります)
使用できるライブラリや Python のバージョンが固定されているのではないか？
- いいえ、上記の通り、Stop Runtime しても消えないストレージ領域に Conda 環境を保存いただくことで、継続的にご利用可能なお客様環境を構築いただく事ができます
SageMaker Studio Lab 上で開発したものを別の開発環境に持っていくことが難しいのではないか？
- Git Repository と連携が可能です
- SageMaker Studio への移行が可能です

これらの懸念の大元にあるのは、「一回きりの利用には良いが、継続して開発する環境としては不十分ではないか？」という観点です。勉強から始めて、せっかくなら、作った環境や成果物を今後も活用したいという要求に対して、SageMaker Studio Lab では上記のように機能を提供しています。さらに、SageMaker Studio Lab では作成した Notebook を AWS の大きな計算機リソースを活用してバッチジョブ化するための機能である Notebook Jobs という機能を提供しています（リンクは SageMaker のものですが同じボタンが SageMaker Studio Lab にもあるとお考えください）。こちらはAWSアカウントを取得する必要があり、SageMakerの利用に伴う料金が発生する点に注意が必要です。重要なことは、SageMaker Studio Lab が提供する計算機リソースを越えて、学習や推論を実行したい場合にも選択肢が提供されていることです。

SageMaker Studio Lab を使えば、初級レベルの機械学習の勉強を超えてご活用いただけそうだとイメージを持っていただければ嬉しいです。それでは、本題の SageMaker Studio Lab で日本語 LLM を Fine-Tuning するお話に入っていきます。本ブログでは、こちらのクイズのブログに倣って、OpenCALM をクイズに回答する精度が向上するように Fine-Tuning していきます。

SageMaker Studio Lab で日本語 LLM OpenCALM を動かす準備

以降の内容では、SageMaker Studio Lab が持つ機能、つまり無料の範囲内でのお話となります。先述の Notebook Jobs は利用しません。NoteBoo Jobs の利用に関しては別途 Blog を公開予定です。また、以降を読み進めていただく上で、SageMaker Studio Lab のアカウントを取得いただいておくとスムーズです。取得方法はこちらです。このブログの実装はこちらのブログを参考にしています。特に、実装部分は以下を参考にしています。併せてご参考くださいませ。

https://huggingface.co/cyberagent/open-calm-7b
https://github.com/aws-samples/aws-ml-jp/tree/main/tasks/generative-ai/text-to-text/fine-tuning/instruction-tuning/Transformers

SageMaker Studio Lab にログインしたら、GPU を選択、Start runtime をクリックてください。初回実行の場合、Mobile Number の登録と SMS 確認が必要な場合があります。日本の電話番号にて Mobile Number のチェックがうまくいかない場合は、Mobile Number の入力時に日本 +81 を選択いただき、電話番号の先頭 0 を除いて入力してみてください。その後、ロボットでないことをチェックする画面を通過してください。SageMaker Studio Lab はご好評いただいており、GPU が確保できない場合がございます。その場合は複数回お試しください。

作業用のディレクトリを作成しましょう。ディレクトリ名は llm-lora-challenge とします。左のメニューにて右クリック、New Folder を選択し、llm-lora-challenge ディレクトリを作成します。

以降、全ての作業はこの llm-lora-challenge ディレクトリ以下で実施します。

必要なライブラリをインストールしていきましょう。2 つの手段があります。ここでは Conda 環境を作成し、可搬性を高める手順 (以下の 1 つ目の方法) を取ります。default 環境を複数の目的で共有して使うとライブラリのバージョン衝突などにより作業しにくくなることがありますが、それを防ぐ効果もあります。

個別の Conda 環境を作成しllm_finetuning.yml を作成し、GUI から Build Conda Environment を実行する
- llm_finetuning という名前の conda 環境を構成する
(参考) default の Conda 環境にライブラリをインストールする
- requirements.txt を作成し、notebook のセルから pip install を実行する
- requirements.txt を作成し、terminal を立ち上げ、default に Conda 環境をスイッチした後、pip install を terminal から実行する

以下のように llm_finetuning.yml を作成します。


name: llm_finetuning
dependencies:
  - python=3.10
  - deepspeed
  - pip
  - pip:
    - git+https://github.com/huggingface/peft.git@207d2908650f3f4f3ba0e21d243c1b2aee66e72d
    - bitsandbytes==0.39.0
    - accelerate==0.20.3
    - transformers==4.30.1
    - tokenizers==0.13.3
    - pynvml==11.4.1
    - protobuf==3.20.2
    - scipy
    - optimum
    - appdirs
    - loralib
    - black
    - black[jupyter]
    - datasets
    - fire
    - sentencepiece
    - evaluate
    - einops
    - ipykernel

作成した llm_finetuning.yml を右クリックしBuild Conda Environment をクリック、確認画面で OK をクリックします。

terminal が起動され、install が開始されます。

以下が表示されれば完了です。

左上のプラスボタンから Launcher を起動して、conda 環境が作成されたか確認してみましょう。

Notebook の部分に作成した llm_finetuning:Python が表示されていれば成功です。

以降の実行は全て llm_finetuning 環境を利用します。Launcher 画面から llm_finetuning:Python をクリックし、ノートブックを開いたら右上の kernel の表示が llm_finetuning:Python になっていることを確認してください。もしなっていない場合は、そちらをクリックし、下記画面のように Select Kernel から llm_finetuning:Python を選択してください。

最終的な構成を先に示します。以降の手順にて作成いただいたり、ダウンロードしたり、実行によって生成されたりするものを含んでいます。

model/
- Fine-Tuning を実行すると作成され、モデルファイルが保存される
llm_finetuning.yml
- 本ブログで使用する llm_finetuning の Conda 環境を作成するためのファイル
data/aio_02_train.jsonl
- ダウンロードされる Fine-Tuning 用ファイル
data/aio_02_train_formatted.jsonl
- Fine-Tuning に利用しやすいようにフォーマットしたファイル
templates/simple_qa_ja.json
- プロンプトテンプレート
OpenCALM_inf.ipynb
- 本ブログで作成する推論用の Notebook ファイル
OpenCALM_format.ipynb
- 本ブログで作成するクイズデータを Fine-Tuning 用にフォーマットする Notebook ファイル
OpenCALM_finetune.ipynb
- 本ブログで作成する Fine-Tuning 用の Notebook ファイル
OpenCALM_finetuned_inf.ipynb
- 本ブログで作成する Fine-Tuning 後のモデルを用いて推論する Notebook ファイル

SageMaker Studio Lab で OpenCALM の推論を呼び出す

ここでは、Fine-Tuning 前の OpenCALM モデルがどの程度クイズに回答できるか確認してみましょう。新規に ipynb ファイル OpenCALM_inf.ipynb を作成し、以下のソースコードを実行してみてください。OpenCALM は HuggingFace Transformers ライブラリから利用する事ができます。初回はモデルのダウンロードが走る分時間がかかります。2 回目以降はダウンロード済みのモデルを利用するため動作が速くなります。model_name に以下を指定 (パラメータ数が小さい順に記載) することで、OpenCALM のパラメータ数が異なるモデルを使用する事ができます。いろいろ変えてみて回答がどうなるか確認してみると面白いかもしれません。

cyberagent/open-calm-small
cyberagent/open-calm-medium
cyberagent/open-calm-large
cyberagent/open-calm-1b
cyberagent/open-calm-3b
cyberagent/open-calm-7b

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = 'cyberagent/open-calm-1b'

model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16, cache_dir="/tmp/model_cache/")
tokenizer = AutoTokenizer.from_pretrained(model_name)

inputs = tokenizer("映画『ウエスト・サイド物語』に登場する2つの少年グループといえば、シャーク団と何団?答えは「", return_tensors="pt").to(model.device)
with torch.no_grad():
    tokens = model.generate(
        **inputs,
        max_new_tokens=64,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.05,
        pad_token_id=tokenizer.pad_token_id,
    )

output = tokenizer.decode(tokens[0], skip_special_tokens=True)
print(f'{model_name}:{output}')

もし、以下のエラーが出た場合、SageMaker Studio Lab が CPU モードで実行されている可能性があります。このブログのコードは GPU モードでのみ動作します。一度ログアウトいただき、GPU モードに切り替えて再度お試しください。

RuntimeError: "LayerNormKernelImpl" not implemented for 'Half'

LLM の出力は確率的要素があるため必ずしも同じになるとは限りません。以下の結果は何か回答しようとはしているものの、正しくない答えが返ってきています。

cyberagent/open-calm-1b:映画『ウエスト・サイド物語』に登場する2つの少年グループといえば、シャーク団と何団?答えは「ダンシングロッド」です。
少年たちの友情や恋心が描かれたミュージカルは老若男女問わず人気で、『美女か野獣』(1992)から現在まで数多くの映画が製作されてきました。『ウェストミンスター寺院の鐘の音が聞こえるまで』『ヘアスプレー』、『ラブアクチュアリー』、そして今世紀最大のヒット作となったのが

ここで、注意が必要なのは、ダウンロードされるモデルのファイルサイズです。モデル名を見てみましょう。 1b, 3b, 7b とあるモデル名はそれぞれのパラメータ数の規模を表しています。b は Billion = 10 億であり、7b は 70 億パラメータ規模のモデルであることがわかります。このモデルファイルは 10 GB を越えており、SageMaker Studio Lab で永続化 (Stop Runtime しても残る) 領域に保存してしまうと容量を圧迫してしまいます。今後継続的に開発に利用することを考えると Conda 環境の保存などに利用するためにこの領域は極力空けておきたいですよね。そこで、cache_dir に /tmp/model_cache/ を指定することで対策しています。AutoModelForCausalLM.from_pretrained メソッドを cyberagent/open-calm-7b を指定して実行した場合、初回にかかる時間は数分程度でした。Start Runtime する度に 1 度だけかかる時間としては許容範囲ではないでしょうか。

SageMaker Studio Lab で OpenCALM を LoRA によって Fine-Tuning する

ここでは OpenCALM のクイズ回答精度を向上するために必要な Fine-Tuning 用のデータを準備します。こちらのブログに沿う形で、クイズデータを Fine-Tuning 用のフォーマットに変換して保存します。新規に OpenCALM_format.ipynb を作成し、以下のソースコードをセルに転載し実行してみてください。data/aio_02_train_formatted.jsonl にフォーマット済みクイズデータが保存されます。

!wget -P data https://jaqket.s3.ap-northeast-1.amazonaws.com/data/aio_02/aio_02_train.jsonl

# Convet .jsonl to .json
import pandas as pd
df = pd.read_json("data/aio_02_train.jsonl", orient="records", lines=True)
df = df.rename(columns={"question": "instruction", "answers": "output"})
df = df[["instruction", "output"]]
df["output"] = df["output"].apply(lambda x: f"{x[0]}」")
df["input"] = ""
print(df.shape)
df.to_json(
    "data/aio_02_train_formatted.jsonl", orient="records", force_ascii=False, lines=True
)
df.head(2)

作成されたデータを確認してみましょう。クイズデータのため、質問 (instruciton 列) と回答 (output 列) というペアでデータが作成されている事がわかります。

次に OpenCALM にクイズ回答用の学習をする準備をします。プロンプトテンプレートファイルを用意します。このテンプレートは OpenCALM に学習データを入力するときのテンプレートになります。以下の作業を実施してください。

templates ディレクトリを作成
こちらの template をダウンロード
SageMaker Studio Lab の templates ディレクトリに配置 (ファイルはドラッグ&ドロップできます)

以下が template の中身です。instruction に続き、答えは「をプロンプトとして含めることで回答を促しています。実は、先述の推論を試してみたときに、OpenCALM に答えは「と入力することで回答させようとしていた部分をテンプレートとして採用しています。{input} を使用する場合とそうでない場合がある事がわかります。{instruction} の部分はクイズのデータサンプルごとに異なります。一方で {input} はプロンプトを別途追加したい場合に利用できます。{input} はこのブログでは使用しません。

OpenCALM_finetune.ipynb ファイルを作成してください。独自のユーティリティクラス Prompter がこのテンプレートファイルを読み、テンプレートに沿った形で OpenCALM にわたす役割です。Prompter をダウンロードしコードをコピー&ペーストでセルに貼り付けてください。.py ファイルとして別途モジュールを import する形でも構いません。

最後に OpenCALM を Fine-Tuning するコードを準備します。こちらの Fine-Tuning のコードをダウンロードしコードをコピー＆ペーストでセルに貼り付けてください。セルに貼り付ける際に以下は不要のためコメントアウトしましょう。

...
# from utils.prompter import Prompter
...
# if __name__ == "**main**":
#    fire.Fire(train)

.py ファイルとして別途モジュール import する形でも構いません。

Fine-Tuning では学習済みモデルを追加のデータを用いて再学習することで精度向上を図ります。パラメータをランダム値などにより初期化した状態から学習することに比べ、学習済みモデルのパラメータを初期値に利用した状態で学習を進める Fine-Tuning には学習のコスト効率を高めながらモデルの精度を高める狙いがあります。しかし、工夫無しに Fine-Tuning してしまうと全てのパラメータを更新対象にしてしまいコスト効率が悪くなったり、学習済みモデルが獲得した一般的な概念を失ってしまう壊滅的忘却と呼ばれる現象が発生するリスクがあります。そこで、少ない更新パラメータ数でコスト効率よく精度を向上させる研究分野や手法群を指す PEFT (Parameter-Efficient Fine Tuning) が生まれました。実務的には PEFT が実装されたライブラリを利用することでその恩恵を受ける事ができます。PEFT は hugging face library に実装があります。本ブログではこの内 LoRA (LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS) に挑戦します。LoRA は学習済みモデルのパラメータはそのままに、新たに追加したパラメータ (Neural Network の重み) に対し更新する手法です。この手法は以下の利点により注目されています。

更新対象パラメータ数の削減によりメモリ量と学習時間を短くできる場合がある
LoRA 部分だけを切り替えることにより、タスクに特化して再学習した Fine-Tuned モデルを効率よく活用できる場合がある

Fine-Tuning のソースコードのうち LoRA らしさが表れている部分に着目して読んでみましょう。実際のソースコードは toknizer や学習途中結果の保存などの実装が含まれていますが、LoRA のポイントは以下の通りです。

AutoModelForCausalLM.from_pretrained により LoRA の元になる学習済みモデルを読み込む
LoraConfig にて、LoRA のハイパーパラメータを設定する
1. で取得した model と 2. で設定した config を get_peft_model に渡すことにより LoRA 用の model を取得する
1. で取得した model オブジェクトを transformers.Trainer に渡し trainer を取得する
1. の trainer を使用し、trainer.train を呼び出す

また、AutoModelForCausalLM.from_pretrained に cache_dir="/tmp/model_cache/" が指定されているため、LoRA に渡した学習済みモデルファイルは SageMaker Studio Lab では永続化されない領域に保存されます。また、output_dir にディレクトリ名を指定すると永続化される領域に LoRA のファイルが保存されます。これによって SageMaker Studio Lab の永続化領域を圧迫することなく LoRA を実行できます。

以上で準備は終了です。さあ、Fine-Tuning を実行してみましょう！

以下のコードをセルに貼り付けて実行してみてください。メモリ不足になる場合は他の Notebook を停止して閉じておきましょう。1b のモデルを学習する場合は 10 – 20 分、7b のモデルを学習する場合は 2 時間程度を要します。まずは動作させてみたいという場合は model_name を cyberagent/open-calm-1b などのパラメータ数の小さいモデルで試してみてください。本ブログでは 1b のモデルを試してみます。

model_name = "cyberagent/open-calm-1b"
model_name_base = model_name.split("/")[-1]
hyperparameters = {
    "base_model": model_name,
    "pad_token_id": 1,
    "data_path": "data/aio_02_train_formatted.jsonl",
    "num_epochs": 1, # default 3
    "cutoff_len": 256,
    "group_by_length": False,
    "output_dir": "model",
    "lora_target_modules": ['query_key_value'],
    "lora_r": 16,
    "batch_size": 32,
    "micro_batch_size": 4,
    "prompt_template_name": "simple_qa_ja",
}
train(**hyperparameters)

model_name に OpenCALM を指定しています。また、以下が大きいほど学習に時間を要しますが精度が改善する可能性があります。大きくし続ければ必ず精度が改善するわけではないため注意が必要です。

num_epochs : 学習サンプルを 1 巡する回数を表すハイパーパラメータ
lora_r: 行列ランクと呼ばれるLoRA にて更新対象にするパラメータ数に依存するハイパーパラメータ

初めはこれらを小さくして、パラメータ数の小さいモデルで LoRA を実行し、正常に動作するかを確認してみると良いでしょう。

以下のような Epoch に対する進捗のログが出力されれば学習が実行されています。指定した num_epochs にログの Epoch が到達したら学習は完了です。

Training Alpaca-LoRA model with params:
base_model: cyberagent/open-calm-1b
data_path: data/aio_02_train_formatted.jsonl
output_dir: model
batch_size: 32
・・・
[XXXXX/XXXXXX X:XX:XX < 00:00, X.XX it/s, Epoch 1.00/1]

もし、学習時間が長く、SageMaker Studio Lab で実行できる GPU 利用時間制限を越えた場合でも、学習途中の LoRA ファイルが model ディレクトリに保存されているため、後で学習した LoRA ファイルを用いて学習を再実行したり、推論に利用したりする事が可能です。

SageMaker Studio Lab で Fine-Tuning した OpenCALM を使って推論する

ここまでの作業により、LoRA により Fine-Tuning した OpenCALM モデルがクイズに上手に回答できる事が期待されます。早速、効果を確認していきましょう。新規に OpenCALM_finetuned_inf.ipynb を作成しましょう。推論のサンプルコードを参考に実装します。このサンプルコードは SageMaker Inference Endpoint にホストする用のコードになっています。このブログでは SageMaker Inference Endpoint は使用しませんので、以下の手順に従って必要なコードだけをセルに貼り付けていきます。

OpenCALM_finetuned_inf.ipynb ファイルを作成してください。LoRA による Fine-Tuning 時に利用した Prompter クラスが必要です。同様に、Prompter をダウンロードしセルに貼り付けてください。.py ファイルとして別途モジュール import する形でも構いません。

推論のサンプルコードから import に関する部分を抜き出しセルに貼り付けます。from utils.prompter import Prompter は Prompter をセルに貼り付けている場合は不要です。

import os
import sys
import json
from typing import Dict

import torch
import transformers
from peft import PeftModel
from transformers import GenerationConfig, AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, BitsAndBytesConfig
import deepspeed

StopOnTokens クラスをセルに貼り付けます。後ほどテキストを生成する際に生成の停止条件を与えるクラスです。

class StopOnTokens(StoppingCriteria):
    def __init__(self, stop_ids):
        self.stop_ids = stop_ids

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        for stop_id in self.stop_ids:
            if input_ids[0][-1] == stop_id:
                return True
        return False

prompter, tokenizer, model を生成します。以下のソースコードをセルに貼り付けてください。

base_model = "cyberagent/open-calm-1b"
device = "cuda"
prompt_template = "simple_qa_ja"
lora_weights = "model"

prompter = Prompter(prompt_template)
tokenizer = AutoTokenizer.from_pretrained(base_model)
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    load_in_8bit=False,
    torch_dtype=torch.float16,
    device_map="auto",
    cache_dir="/tmp/model_cache/"
)
print("Loading Lora Weight")
model = PeftModel.from_pretrained(
    model,
    lora_weights,
    torch_dtype=torch.float16,
)

model.model_parallel = False

if torch.cuda.device_count() > 1:
    model.is_parallelizable = True
    model.model_parallel = True

最後に、prompt を生成し、テキストを生成するパラメータを設定したら、テキスト生成するコードをセルに貼り付けてください。

instruction = "映画『ウエスト・サイド物語』に登場する2つの少年グループといえば、シャーク団と何団?"
input = ""
max_new_tokens = 32
stop_ids = [1, 0]

prompt = prompter.generate_prompt(instruction, input)

inputs = tokenizer(
    prompt, 
    add_special_tokens=False,
    return_token_type_ids=False,
    return_tensors="pt"
).to(device)

generation_config = GenerationConfig(
    max_new_tokens=max_new_tokens,
    return_dict_in_generate=True,
    output_scores=True,
    temperature=0.1,
    do_sample=False,
    num_beams=5,
    pad_token_id=1,
    bos_token_id=0,
    eos_token_id=0
)

with torch.no_grad():
    generation_output = model.generate(
        **inputs,
        generation_config=generation_config,
        stopping_criteria=StoppingCriteriaList([StopOnTokens(stop_ids)]),
    )
s = generation_output.sequences[0, inputs['input_ids'].size(1):]
output = tokenizer.decode(s, skip_special_tokens=True)
output

推論のソースコードのうち LoRA らしさが表れている部分に着目して読んでみましょう。実際のソースコードは toknizer などの実装が含まれていますが、LoRA のポイントは以下の通りです。

AutoModelForCausalLM.from_pretrained により LoRA の元になる学習済みモデルを読み込む
1. の model と lora_weights = “model“ を PeftModel.from_pretrained に渡し model を取得する
1. の model により model.generate を呼び出す

LoRA の保存先が model ディレクトリであったことを思い出してみましょう。Fine-Tuning 時と同様に、AutoModelForCausalLM.from_pretrained に cache_dir="/tmp/model_cache/" が指定されているため、SageMaker Studio Lab の永続化領域を圧迫することなく LoRA を用いたモデルの推論を実行できます。

さあ、OpenCALM_finetuned_inf.ipynb を実行してみましょう。検証時の結果では、以下のように正しい回答が返ってきました。LLM の出力は確率要素があるため本ブログと同じ回答が得られない場合があることに注意が必要です。

ここで、興味深い事は Fine-Tuning 用ファイルには類似のクイズは含まれていなかった事です。Fine-Tuning に用いたクイズデータを確認したところ、ウェストサイドストーリーに関するクイズは以下のみでした。

{"instruction":"ミュージカル「ウエストサイド・ストーリー」の作曲で知られる音楽家は誰でしょう?","output":"レナード・バーンスタイン」","input":""}

元となる OpenCALM の学習データにはウェストサイドストーリーに関するテキストが含まれており、クイズ回答に特化した Fine-Tuning によって正しい回答を得られるようになった可能性があります。

発展

本ブログを参考に以下にチャレンジしてみましょう。

cyberagent/open-calm-7b を試す
- 本ブログのコードは 1b のモデルを試しましたが、作者の動作確認では 7b のモデルまで実行できることを確認しています。モデルの大きさによって回答精度がどう変わってくるのか、処理時間はどうか、是非お試しください。
SageMaker Studio への移行
- SageMaker Studio Lab で作成したコードを SageMaker Studio Notebook で実行してみましょう。より多くの計算機リソースや機械学習ワークロードを実現する多くの機能と連携できるようになります。
OpenCALM 以外の日本語 LLM を試す
- Hugging Face に実装されているモデルであれば本ブログの実装を参考にすることができます。例えば、rinna/japanese-gpt-neox-3.6b · Hugging Face を利用できるように改修してみましょう。rinna モデルの場合はどのようなプロンプトを与えると良いか試してみましょう。もしかしたら、OpenCALM とは違う特性があるかもしれません。

最後に

本ブログで解説した内容は SageMaker Studio Notebook, SageMaker Notebook Instance でも同様に実行可能です。もし、皆様が既に GPU を搭載した計算機環境をお持ちであれば、Jupyter notebook や Jupyter Lab を導入することで同様に実行可能です。生成系 AI はここ数年で実務レベルで活用できるユースケースが多岐に渡るほどの進化を遂げました。トップダウンで検証することになったご担当者様やボトムアップで面白い技術要素としてスキルを獲得しようとしている方まで様々かと思います。新しい技術を検証するとき、社会課題、社内課題ともにどんな課題がこの技術で解決できるかを机上で考える事は重要です。また、皆様の抱える課題のうち優先順位の高い費用対効果の高いテーマを選択することも重要です。同時に、予算や体制が小さな状況でも動くものを作り、その技術を通して何がどう動くのか体験することも重要だと考えています。このブログを通じて、日本語 LLM の OSS がどのように動作するのか、何に使えそうなのかを体験いただき、日本語 LLM の OSS が持つ可能性を感じていただければ大変嬉しく思います。

著者

中島佑樹 西日本のお客様をメインで担当するソリューションアーキテクト。社会人博士を修了したことをきっかけに AIML を得意分野としている。システム一般のテーマや Amazon Bedrock を用いた生成系 AI のシステム開発、Amazon SageMaker Studio Lab を用いた AIML への入門まで幅広く活動。

Amazon Web Services ブログ

無料のノートブックサービス SageMaker Studio Lab で 日本語 LLM を Fine-Tuning する