SageMaker JumpStart で Gemma 2B を試す【Fine-tuning編】

# はじめに

前回、SageMaker JumpStartでGemmaを簡単にデプロイする方法を確認しました。

https://qiita.com/kanuazut/items/c4de7f81c1c6bbdaecb1

今回は、JumpStartでFine-tuningする方法を確認しようと思います。
SageMakerの利用の始め方やOSSモデルのデプロイの方法は、上記の記事で紹介しているのでそちらをご確認ください。
なお、Instruction fine-tuningについてのみ紹介します。



# 仕様の確認

## ドキュメントの確認方法
まずはじめにモデルのドキュメントを確認します。
今回はGemma 2B ITを対象とするので、JumpStartで検索しGemma 2B ITのページを開きます。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/efa3cbac-ae58-1773-3da1-31a6828904d0.png)

<br>

以下がモデル個別のページです。このページからDeployやTrainingを行うことになります。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/99450270-4977-ec49-cff8-86ca445d8f70.png)

ここに書いてあるドキュメントが実は重要です。推論方法や学習方法などの説明が書いてあり、注意深く読み込む必要があります。
他のドキュメントとしてAWSサンプルリポジトリにはSageMaker用のNotebookが公開されていますが、コードがメインの内容となっています。



## ドキュメントの内容
このドキュメントの内容はJumpStartから参照可能ですが、他に公開されているページが見つからなかったので、一部だけここで紹介しておきます。

### 推論の方法
> **Use the Deployed Model for Inference**
The deployed model can be used for running inference on any input text. Example python code for how to run inference on the deployed model is given in a notebook you can open by clicking 'Open Notebook' on the model endpoint page that is created when deploying the model.

デプロイによって推論実行が可能になること、Pythonコードサンプルが提示されることが書かれています。また実際のドキュメントには入力テキストと出力テキストのサンプルも掲載されています。
ちなみに実際にデプロイした後にはすぐにテストできるようテストリクエストが試せるボタンとサンプルコードが表示されます。

### Fine-tuningの方法
> **Fine-tune the Model on a New Dataset**
We currently offer three types of fine-tuning: instruction fine-tuning, domain adaptation fine-tuning, and chat fine-tuning. To fine-tune the model for chat format, you can set chat_dataset to be True. To fine-tune the model with instruction tuning data, set instruction_tuned to be True. At most, only one of chat_dataset and instruction_tuned may be set to True. If both are False, the data is assumed to be domain adaptation format.

現在、Fine-tuningの方法は以下の3つが提供されています。
各方法を選択したときに、トレーニングジョブ実行時に設定すべきパラメータもあわせて記載しておきます。

|Format|chat_dataset|instruction_tuned|
|-|-|-|
|Chat fine-tuning (Default)|True|False|
|Instruction fine-tuning|False|True|
|Domain adaption fine-tuning|False|False|


ここからそれぞれの方法について、特に入力データのフォーマットについてみていきます。

### Chat fine-tuning
> **Chat fine-tuning**
> The training data must be formatted in a JSON lines (.jsonl) format, where each line is a dictionary
representing a single data sample. All training data must be in a single folder, however it can be saved in multiple jsonl files. The .jsonl file extension is mandatory. Each line in the file is a list of conversations between the user and the assistant model. This model only supports 'user' and 'assistant' roles while alternating (u/a/u/a/u...). Here is an example:

データフォーマットはJSON Lines(`.jsonl`)です。
jsonlファイルは複数もOKですが、全ての学習データが1つのフォルダ内にある必要があります。

データ例は以下の通りです。user, assistatが交互に登場する必要があります。
```json
{"dialog": [
{"content":"what is the height of the empire state building","role":"user"},
{"content":"381 meters, or 1,250 feet, is the height of the Empire State Building. If you also account for the antenna, it brings up the total height to 443 meters, or 1,454 feet","role":"assistant"},
{"content":"Some people need to pilot an aircraft above it and need to know.\nSo what is the answer in feet?","role":"user"},
{"content":"1454 feet","role":"assistant"}
]}
```


### Instruction fine-tuning
> **Instruction fine-tuning**
> The training data must be formatted in a JSON lines (.jsonl) format, where each line is a dictionary
representing a single data sample. All training data must be in a single folder, however it can be saved in multiple jsonl files. The .jsonl file extension is mandatory. The training folder can also contain a template.json file describing the input and output formats. If no template file is given, the following template will be used:



データフォーマットはJSON Lines(`.jsonl`)です。
jsonlファイルは複数もOKですが、全ての学習データが1つのフォルダ内にある必要があります。
またオプションとして`template.json`を含めることで、指定のテンプレート形式でデータを作成することができます。

デフォルトデータ例は以下の通りです。promptとcopletionが必要です。
```json
{
  "prompt": "{prompt}",
  "completion": "{completion}"
}
```

以下のようなカスタムテンプレートが利用可能です。promptとcompletionのキーは必須です。
```json
{
  "prompt": "question: {question} context: {context}",
  "completion": "{answer}"
}
```
このカスタムテンプレートを指定した場合は、question, context, answerを含んだjsonlをデータとして用意する必要があります。


### Domain adaption fine-tuning
> **Domain adaption fine-tuning**
Input: A train and an optional validation directory. Each directory contains a CSV/JSON/TXT file.
For CSV/JSON files, the train or validation data is used from the column called 'text' or the first column if no column called 'text' is found.
The number of files under train and validation (if provided) should equal to one, respectively.




データフォーマットはCSV/TXT/JSONです。
CSV/JSONの場合は、デフォルトで`text`カラムがデータとして利用されます。`text`カラムが存在しない場合は、最初のカラムが利用されます。
trainとvalidationのファイル数は、それぞれ1つである必要があるようです。

以下はTXTファイルを利用する場合の例です。
```
This report includes estimates, projections, statements relating to our
business plans, objectives, and expected operating results that are “forward-
looking statements” within the meaning of the Private Securities Litigation
Reform Act of 1995, Section 27A of the Securities Act of 1933, and Section 21E
of the Securities Exchange Act of 1934. Forward-looking statements may appear
throughout this report, including the following sections: “Business” (Part I,
Item 1 of this Form 10-K), “Risk Factors” (Part I, Item 1A of this Form 10-K),
and “Management’s Discussion and Analysis of Financial Condition and Results
of Operations” (Part II, Item 7 of this Form 10-K). These forward-looking
statements generally are identified by the words “believe,” “project,”
“expect,” “anticipate,” “estimate,” “intend,” “strategy,” “future,”
“opportunity,” “plan,” “may,” “should,” “will,” “would,” “will be,” “will
continue,” “will likely result,” and similar expressions. Forward-looking
statements are based on current expectations and assumptions that are subject
to risks and uncertainties that may cause actual results to differ materially.
We describe risks and uncertainties that could cause actual results and events
to differ materially in “Risk Factors,” “Management’s Discussion and Analysis
of Financial Condition and Results of Operations,” and “Quantitative and
Qualitative Disclosures about Market Risk” (Part II, Item 7A of this Form
10-K). Readers are cautioned not to place undue reliance on forward-looking
statements, which speak only as of the date they are made. We undertake no
obligation to update or revise publicly any forward-looking statements,
whether because of new information, future events, or otherwise.
```


### 
以上、Gemma 2B ITのFine-tuningには上記3種類の方法があります。

筆者はあまりFT自体に詳しくないですが、説明文を読んだ限りでは
- Chat-FTは会話ベースで複数回の往復を含む場合に適していそう
- Instruction FTはタスクベースで指示と情報を与え何か解かせたい場合に適していそう
- Domain adaption FTはドメイン固有の情報がとりあえず大量にあるが構造化されていない場合に適していそう

と認識しました。名前の通りですが。

ちなみにそれぞれ同じデータソースから作ったデータでFTを試しましたが、Chat-FTとDomain Adaption FTは全く学習が進みませんでした。


# データの作成
今回はAmazon Bedrockのドキュメントに対するQAが出来るようにFTしていきます。


話の主軸ではないので詳しくは記載しませんが、
Amazon Bedrckの[AWS公式ドキュメント](https://docs.aws.amazon.com/ja_jp/bedrock/latest/userguide/what-is-bedrock.html)から6ページ分のHTMLを取得し、Claude3を使って以下のようなＱＡデータを作成させました。
学習データは117件となりました。

```json
{"prompt": "Amazon Bedockのナレッジベースとはなんですか?", "comletion": "Amazon Bedrockのナレッジベースは、データソースを情報のリポジトリにまとめて管理することができる機能です。ナレッジベースを使えば、検索拡張生成(RAG)を活用したアプリケーションを簡単に構築できます。"}
{"prompt": "ナレッジベースでサポートされる基盤モデルを教えてください", "comletion": "ナレッジベースでサポートされる基盤モデルは、Amazon Bedrock経由で利用できるAI21 Labs、Anthropic、Cohere、Meta、Stability AI、Amazonなどの各社モデルです。具体的にどのモデルがサポートされているかは、Bedrockのドキュメントの「サポートされているリージョンおよびモデル」の項目で確認できます。サポート状況は随時更新されています。"}
{"prompt": "埋め込みってどういうこと?", "comletion": "埋め込みとは、オブジェクト間の類似性を数値表現で比較するため、入力データをベクトル表現に変換する過程です。テキストや画像の類似性、テキストと画像の関連性が分かります。"}
{"prompt": "モデル評価って何を指すの?", "comletion": "モデル評価とは、ユースケースに最適なモデルを選ぶため、モデル出力を評価し比較するプロセスのことです。"}
{"prompt": "Amazon Titanの提供するモデルの種類を教えてください。", "comletion": "Amazon TitanからはText G1、Image Generator G1、Text Embeddings G1、Multimodal Embeddings G1の4種類のモデルが提供されています。Text G1はテキスト入出力、Image Generator G1はテキスト入力で画像出力、Text Embeddings G1はテキスト入力で埋め込みベクトル出力、Multimodal Embeddings G1はテキストと画像の入力でマルチモーダル埋め込み出力を行います。"}
```

データは後でJumpStartから参照できるようS3バケットに格納しておきます。

# Fine-tuningの実行
ここから本題のFine-tuningの手順です。

Gemma 2B ITを選択した画面に戻ります。右上の「Train」をクリック。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/681c80a7-de15-1a7c-578d-f7800d90fa32.png)


Fine-tuningの設定値を入力していきます。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/aa312f4e-c08a-b34c-63da-74e11382698c.png)


まずはデータです。`Enter training dataset`にラジオボタンを変更して、データ格納先のS3パスを入力します。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/7ebae0f6-c767-bf08-a488-f69ca1f9554d.png)

ハイパーパラメータを入力します。
ここで変更しなければいけないのは、一番上の方にある`Chat Dataset`, `Instruction Train The Model`の項目です。前述したとおり、3つのFTのどれを選ぶかで、True/Falseを設定する必要があります。また、Epochsが学習エポック数になるため、適宜設定します。ここではEpochsは1です。
その他の値は、デフォルトで推奨値となっているようなので、今回はそのまま進めます。がっつりチューニングするようになったら検討すればよいでしょう。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/9b54aaf1-14f5-4446-7a3e-78553a9b32ce.png)

トレーニングジョブのインスタンスサイズとモデルアーティファクトの出力S3を指定します。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/8dd2a90c-7429-6d70-b08b-c9224585ee30.png)

トレーニングジョブ名を決めます。今回はデフォルト指定しています。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/042f95a0-0cb2-9094-2c8e-617eeb9b1e70.png)

設定は以上です。
VPCモードで実行したければAdvanced Settingで設定可能です。

最後に「Train」をクリックするとジョブが実行されます。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/29e6889d-0799-80af-d7ec-b860db3a48df.png)

ジョブの実行が完了するとStatusがCompletedに変わります。学習中のLogは適宜確認可能です。
ジョブが完了すると右上に「Deploy」ボタンがEnableになるので、このボタンを押すとそのままエンドポイントを払い出すことができます。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3115707/9a836597-7a53-7e85-65b6-f2d539c3c6b9.png)

以上、データさえ用意してしまえばFTの実行がGUIをポチポチで完結することが分かっていただけるかと思います。


# 結果比較
FTの手順自体は以上ですが、せっかくなのでFT前とFT後のGemma 2B ITの出力比較をのっけておきます。
汎化性能の確認ではなく単に学習出来てるかを見るため、学習データをそのままプロンプトとして渡していることに注意してください。

### サンプル1
- Prompt：
`Mistral AIのモデルはあるのでしょうか?`
- 学習データ：
`はい、Mistral AIから利用できるモデルとして、Mistral 7B InstructとMixtral 8X7B Instructの2つがあります。これらはテキストを入力すると、テキスト形式で出力を行うモデルとなっています。`

|FT前|FT後|
|-|-|
|Mistral AIは、自然言語処理と機械学習を組み合わせたAIモデルです。自然言語処理はテキストの理解と生成を処理する技術であり、機械学習はデータに基づいてモデルを学習する技術です。<br><br>**主要な機能:**<br>* **言語理解:** 言語のテキストを理解し、その意味を理解します。<br>* **言語生成:** 理解したテキストを新しいテキストとして生成します。<br>* **言語翻訳:** 複数の言語間のテキスト翻訳を実行します。<br>* **質問回答:** 質問に回答する能力を持つモデルです。<br>* **テキスト分類:** テキストのカテゴリーを分類します。<br>* ** sentiment分析:** テキストの sentimentを分析します。<br><br>**用途:**<br><br>* **顧客サービス:** chatbotsや自動テキストサポート<br>* **コンテンツ作成:** 記事、記事、ブログ記事の作成<br>* **翻訳:** 翻訳、言語間翻訳<br>* **質問回答:** クライアントに回答を提供するシステム<br>* **データセット作成:** テキストデータの作成<br>|はい、Mistral AIから利用できるモデルは、Mistral 7B InstructとMixtral 8X7B Instructの2つがあります。|


あきらかにFT後では学習データに近い出力が出来ています。このサンプルだけを見れば適切に学習できているように見えますね。
一般的な回答としてはFT前が適切ですが、BedrockのQAとしてはFT後が正しいです。


### サンプル2
- Prompt：
`Amazon Bedrockでサポートされている基盤モデルのプロバイダーを教えてください。`
- 学習データ：
`Amazon BedrockでサポートされているプロバイダーのいくつかはAI21 Labs、Anthropic、Cohere、Meta、Stability AI、Amazonなどです。`

|FT前|FT後|
|-|-|
|Amazon Bedrockは、Amazon Marketplaceでサポートされています。<br><br>* **Amazon Aurora**<br>* **Amazon DynamoDB**<br>* **Amazon Elasticsearch**<br>* **Amazon Lex**<br>* **Amazon Kendra**<br>* **Amazon Neptune**<br>* **Amazon Redshift**<br>* **Amazon DynamoDB Streams**<br>* **Amazon Cognito**<br>* **Amazon SNS**<br>* **Amazon SQS**<br>* **Amazon SNS Lambda**<br>* **Amazon ECS**<br>* **Amazon RDS**<br>* **Amazon DynamoDB Streams**<br>* **Amazon Comprehend**<br>* **Amazon Comprehend Conversational AI**|Amazon Bedrockでは、AI21 Labs、Anthropic、Cohere、Meta、Stability AI、Amazonなどのプロバイダーからの基盤モデルを利用できます。|


FT前ではBedrockの知識がないからなのか質問の意図も捉えていませんが、FT後は適切に回答が出来ています。

### サンプル3
- Prompt：
`Amazon Bedrockの料金体系を教えてくれますか?`
- 学習データ：
`Amazon Bedrockの料金は、使用する基盤モデル、入力と出力のトークン数、プロビジョンドスループットの購入の有無で変わってきます。詳細については「Amazon Bedrockの料金体系」のページをご確認ください。`

|FT前|FT後|
|-|-|
|Amazon Bedrockの料金体系は、サービスによって異なる価格で提供されています。<br><br>**基本的な料金:**<br><br>* **Starter:** $120/monthで開始し、1000チャットメッセージや100GBのファイル共有に制限されます。<br>* **Standard:** $180/monthで開始し、10,000チャットメッセージや1TBのファイル共有に制限されます。<br>* **Professional:** $280/monthで開始し、 ilimit|Amazon Bedrockの料金体系は、使用する基盤モデル、入力と出力のトークン数、プロビジョンドスループットの購入有無によって変わります。<br><br>* **基盤モデル:** Base, Standard、Premiumの3種類があります。Baseは一般的な用途に適しています。Standardはプロビジョンドスループットを購入して|


FT前とFT後で、似た内容の出力になっており、Bedrockには3段階の料金設定があるとハルシネーションを起こしています。しかし、FT後の出力の前半部分は、学習データの内容を取り込めており、多少の調整が行われていることがうかがえます。まだ学習が足りずFT前と近い内容も出ていますが、epoch数を増やしていけば改善しそうなパターンだと思います。



### サンプル4
- Prompt：
`ガードレールとは何のことですか?`
- 学習データ：
`ガードレールとは、生成AIアプリケーションから不適切または望ましくない内容が出力されるのを防ぐための機能です。Amazon Bedrockではこのガードレール機能を利用して、保護対策を適用することができます。`

|FT前|FT後|
|-|-|
|ガードレールとは、建築物の外縁や壁に設置された、障害物を防止するための枠や枠のことです。|ガードレールとは、人の安全を確保するために設置された、人や物への衝突を防止するための構造のことです。道路や建築物などの構造に組み込まれ、人の体や物に衝突した場合に保護する役割を果たします。|

こちらはうまく学習が進んでいないパターンです。ガードレールと聞かれたらAmazon BedrockのGuardrailsを説明してほしいのですが、FT前、FT後とも既存の知識に引っ張られています。
今回のFTでは、殆どの場合でこのように出力は変わったものの適切な出力が出来ない結果となりました。
サンプル1,2のように学習データ通りに回答出来るパターンの方が少数派です。



# 学習時間と費用

今回は、ml.g5.4xlargeのインスタンスでデータ数117件に対してFTを行いました。
1epochと10epochで学習を行った際の学習時間を以下に示します。

|epochs|学習時間(分)|概算費用|
|-|-|-|
|1|7.7|0.26 USD|
|10|10.7|0.36 USD|

Gemma 2B ITであれば、100件のデータで1epoch 50円くらいということですね。
当然ながらモデルやデータやインスタンスによって変わってきますし、あくまでサンプルとしてではありますが、最低料金として100円以下でもFT出来る可能性があるのはとてもうれしいことですね。（ちなみにネットワーク料金、ストレージ料金もかかりますが10分程度の実行なので大した額にならないだろうと無視しています）




# おわりに

今回はJumpStartでFine tuningする手順を確認しました。
Fine-tuningの本髄はここからの精度評価、パラメータチューニングやデータ品質改善にあると思いますが、ひとまずGUIポチポチするだけでFine-tuningが走らせられました。こういったサービスは筆者のような素人にとってはとても有難い機能であるとともに、経験者にしても本質的な部分に労力をかけられるので便利な機能なのではないかなと思います。
LLMの学習といえば大金がかかる印象ですが、LoRAという手法が開発され、性能の良い軽量モデルが公開され、一般人でもミニマムなFTなら手が出せる領域になってきました。これからもっと多くのモデルがJumpStartで使えるようになることを期待しています。