Amazon SageMaker 上で MPT-7B を微調整する

'Fine-tuning MPT-7B on Amazon SageMaker.'

Amazon SageMaker 上で MPT-7B のチューニングに必要なデータセットを準備し、トレーニングジョブを作成する方法を学ぶ

新しい大規模言語モデル（LLM）が毎週発表されており、それぞれが前任者を打ち負かして評価リーダーボードを引き継ごうとしています。その中でも、MosaicML がリリースした MPT-7B は最新のモデルの一つです。この7,000,000,000のパラメータを持つモデルは、他の同様のモデルとは異なり、オープンソースで商用利用のためにライセンスされています（Apache 2.0 ライセンス）🚀。

MPT-7B のような基礎モデルは、Web からクロールされた数兆のトークン（100トークン〜75語）から構成されるデータセットで事前にトレーニングされており、最適に提示された場合、印象的な出力を生成できます。しかし、実際の世界のアプリケーションで大規模言語モデルの価値を本当に引き出すためには、スマートなプロンプトエンジニアリングだけでは自分のユースケースに合わないかもしれないため、特定のドメイン用のデータセットで基礎モデルをファインチューニングする必要があります。

LLM には数十億のパラメータがあり、そのため、このような大規模なモデルをファインチューニングすることは困難です。幸いなことに、ドメイン固有のデータセットが「小さく」、ファインチューニングにはトレーニングデータをわずか数回しか通過させないため、ファインチューニングは事前トレーニングされた基礎モデルよりもはるかに安価で高速です。

この記事では、以下の内容を学びます。

• 線形回帰と勾配降下法
• Boto3 vs AWS Wrangler PythonによるS3操作の簡素化
• 非教師あり学習シリーズ：階層クラスタリングの探索

• 大規模言語モデルをファインチューニングするためのデータセットを作成および構造化する方法。
• 分散トレーニングジョブを完全にシャードされたデータ並列で構成する方法とは何か。
• 😊 HuggingFace Estimator を定義する方法。
• Amazon SageMaker 上で MPT-7B をファインチューニングするトレーニングジョブを起動する方法。

1. 依存関係をインストールし S3 パスを設定する

SageMaker Python SDK とその他のパッケージをインストールして開始します。この SDK を使用すると、Python コードを数行書くだけで AWS 上で機械学習モデルをトレーニングおよびデプロイすることができます。以下のコードは、Github の sagemaker_finetuning.ipynb ノートブックで利用できます。SageMaker Studio、SageMaker ノートブックインスタンス、または AWS アカウントに認証した後、ノートブックを実行します。

!pip install "sagemaker==2.162.0" s3path boto3 --quietfrom sagemaker.huggingface import HuggingFacefrom sagemaker.inputs import TrainingInputfrom sagemaker import s3_utilsimport sagemakerimport boto3import json

次に、データが S3 に保存されるパスを定義し、SageMaker セッションを作成します。

# S3 パスを定義するbucket             = "<YOUR-S3-BUCKET>"training_data_path = f"s3://{bucket}/toy_data/train/data.jsonl"test_data_path     = f"s3://{bucket}/toy_data/test/data.jsonl"output_path        = f"s3://{bucket}/outputs"code_location      = f"s3://{bucket}/code"# SageMaker セッションを作成するsagemaker_session  = sagemaker.Session()region             = sagemaker_session.boto_region_namerole               = sagemaker.get_execution_role()

2. ファインチューニング用のデータセットを作成する

ファインチューニングのデモンストレーションのために、ダミーデータセットを作成します。このサイズのモデルを完全なデータセットでトレーニングすることは時間がかかり、コストがかかるため、まず小さなデータセットでトレーニングジョブをテストしてデバッグすること、そしてその後トレーニングを完全なデータセットにスケールアップすることが良いアイデアです。

• データセットを辞書のリスト形式でフォーマットする — データセットは、各例がキーと値の構造を持つ辞書のリスト形式でフォーマットする必要があります。例えば、以下のようになります。

{    "prompt": "Pastel de Nata とは何ですか？",    "response": "Pastel de Nata は、ポルトガル風のエッグカスタードタルトパイで、シナモンがかけられることもある。"}

promptはモデルに与えられる入力(例えば、質問)です。responseはモデルが予測する出力(例えば、promptの質問に対する回答)です。生のpromptは、モデルがより良い出力を生成するのを助けるプロンプトテンプレートに合わせて前処理されることが多いです。モデルは因果関係言語モデリングのためにトレーニングされているため、「ドキュメントを完了する」と考えることができます。モデルがドキュメントを完了していると思わせるように、プロンプトテンプレートを設計することが良いアイデアです。Andrej Karpathyは、彼の講演「State of GPT」で、このメカニズムをよく説明しています。

prompt_template = """以下の質問に適切に答えるレスポンスを書いてください。### 質問:{question}### レスポンス:"""dataset = [    {"prompt": "Pastel de Nataとは何ですか？",     "response": "Pastel de Nataは、シナモンで振りかけられる場合があるポルトガルの卵のカスタードタルトのペイストリーです。"},    {"prompt": "アムステルダムで有名な美術館は何ですか？",     "response": "アムステルダムには、世界的に有名な様々な美術館があり、Rijksmuseum、Van Gogh Museum、Stedelijk Museumに立ち寄ることなくして市内観光は完了しません。"},    {"prompt": "欧州議会はどこにありますか？",     "response": "欧州議会の正式な所在地はストラスブールです。"},    {"prompt": "オランダの天気はどうですか？",     "response": "オランダは、穏やかな夏と寒い冬を特徴とする典型的な海洋性気候の国です。"},    {"prompt": "Poffertjesとは何ですか？",     "response": "Poffertjesは、イーストとそば粉を使用した伝統的なオランダのバッター菓子です。小さくてふわふわしたパンケーキに似ています。"},]# テンプレートに基づいてプロンプトをフォーマットfor example in dataset:    example["prompt"] = prompt_template.format(question=example["prompt"])training_data, test_data = dataset[0:4], dataset[4:]print(f"トレーニングデータのサイズ: {len(training_data)}\nテストデータのサイズ: {len(test_data)}")

• トレーニングデータとテストデータをS3にアップロード — トレーニングデータとテストデータがリストとして準備され、JSONラインとしてS3にアップロードされます。以下のユーティリティ関数を使用します。

def write_jsonlines_to_s3(data, s3_path):    """辞書のリストをJSONラインファイルとしてS3に書き込む"""    json_string = ""    for d in data:        json_string += json.dumps(d) + "\n"    s3_client   = boto3.client("s3")        bucket, key = s3_utils.parse_s3_url(s3_path)    s3_client.put_object(         Body   = json_string,         Bucket = bucket,         Key    = key,    )write_jsonlines_to_s3(training_data, training_data_path)write_jsonlines_to_s3(test_data, test_data_path)

3. SageMakerトレーニングジョブ

S3にデータセットがあるため、Amazon SageMakerでトレーニングジョブを作成します。そのために、エントリポイントスクリプトを作成し、トレーニング設定を指定する構成ファイルを変更し、HuggingFaceエスティメータを定義する必要があります。トレーニング環境を設定するLLM FoundryおよびComposerライブラリのCLIランチャーからトレーニングスクリプトを(再)使用します。これらのパッケージは、MPT-7Bの背後にあるMosaicMLがメンテナンスしています。作業フォルダは以下のように構成する必要があります。

└── fine-tune-mpt-7b-sagemaker/    ├── training_script_launcher.sh    ├── fine_tuning_config.yaml    ├── sagemaker_finetuning.ipynb

これらのファイルのそれぞれについて深く掘り下げていきます。

• 構成ファイルを作成する finetuning_config.yaml — LLM Foundryリポジトリで提供されるテンプレート、特にmpt-7b-dolly-sft.yamlファイルを使用すると、良い出発点になります。ただし、データセットのサイズやトレーニングインスタンスによっては、バッチサイズなどのこれらの設定を調整する必要があるかもしれません。私はファイントゥーニングのためにファイルを変更しました(check finetuning_config.yaml)。注目すべきパラメータは以下のとおりです。

max_seq_len: 512global_seed: 17...# Dataloaderstrain_loader:  name: finetuning  dataset:    hf_name: json    hf_kwargs:        data_dir: /opt/ml/input/data/train/...eval_loader:  name: finetuning  dataset:    hf_name: json    hf_kwargs:        data_dir: /opt/ml/input/data/test/ ...max_duration: 3epeval_interval: 1ep...global_train_batch_size: 128...# FSDPfsdp_config:  sharding_strategy: FULL_SHARD  mixed_precision: PURE  activation_checkpointing: true  activation_checkpointing_reentrant: false  activation_cpu_offload: false  limit_all_gathers: true  verbose: false# Checkpoint to local filesystem or remote object storesave_folder: /tmp/checkpointsdist_timeout: 2000

max_seq_length は入力トークンの最大数を示します (100 トークン ≒ 75 単語を覚えておいてください)。トレーニングデータとテストデータは、トレーニングジョブに関連付けられたコンテナ内の /opt/ml/input/data/{train, test} ディレクトリから Datasets ライブラリを使用してロードされます。コンテナディレクトリの構造については、SageMaker Training Storage Folders のドキュメントを確認してください。 max_duration はファインチューニングのエポック数を指定します。 2 〜 3 エポックが通常の選択肢です。 eval_interval は、モデルがテストセットで評価される頻度を示します。

分散トレーニング戦略は、Fully Sharded Data Parallel (FSDP) であり、MPT-7B のような大きなモデルの効率的なトレーニングを可能にします。各 GPU にモデルのコピーを保持する従来のデータ並列戦略とは異なり、FSDP はデータ並列ワーカー間でモデルパラメータ、オプティマイザステータス、および勾配をシャードします。FSDP について詳しくは、この洞察的な PyTorch イントロポストを確認してください。FSDP は、LLM Foundry で使用される分散トレーニングライブラリである Composer に統合されています。

save_folder はモデルのチェックポイント ( .pt ファイル) が保存される場所を決定します。ここでは、一時フォルダ /tmp/checkpoints に設定します。

• エントリーポイントスクリプトの作成 launcher.sh — エントリーポイントとして bash スクリプトが使用されます。bash スクリプトは、LLM Foundry リポジトリをクローンし、必要なパッケージをインストールし、Composer ライブラリの分散ランチャーを使用してトレーニングスクリプトを実行します。通常、SageMaker のトレーニングジョブでは、python train.py のようなコマンドを使用してトレーニングスクリプトを実行します。ただし、エントリーポイントに bash スクリプトを渡すこともでき、このシナリオではより柔軟性があります。最後に、/tmp/checkpoints に保存されたモデルチェックポイントを HuggingFace モデル形式に変換し、最終アーティファクトを /opt/ml/model/ に保存します。SageMaker は、このディレクトリ内のすべてのファイルを圧縮し、tarball model.tar.gz を作成して S3 にアップロードします。tarball は推論に役立ちます。

# MosaicML から llm-foundry パッケージをクローンします# ここにトレーニングスクリプトがホストされていますgit clone https://github.com/mosaicml/llm-foundry.gitcd llm-foundry# 必要なパッケージをインストールしますpip install -e ".[gpu]"pip install git+https://github.com/mosaicml/composer.git@dev# ファインチューニング構成を使用してトレーニングスクリプトを実行しますcomposer scripts/train/train.py /opt/ml/code/finetuning_config.yaml# Composer チェックポイントを HuggingFace モデル形式に変換しますpython scripts/inference/convert_composer_to_hf.py \    --composer_path /tmp/checkpoints/latest-rank0.pt \    --hf_output_path /opt/ml/model/hf_fine_tuned_model \    --output_precision bf16# モデルアーティファクトディレクトリの内容を表示しますls /opt/ml/model/

• 😊 HuggingFace Estimator を定義する — Estimator は、トレーニングジョブを実行するために使用される Docker コンテナを設定します。PyTorch 2.0.0 および Python 3.10 を備えたイメージを使用します。bash スクリプトと構成ファイルは自動的に S3 にアップロードされ、コンテナ内で利用可能になります (SageMaker Python SDK によって処理されます)。トレーニングインスタンスを g5.48xlarge に設定し、8x NVIDIA A10G GPU を搭載しています。 p4d.24xlarge は、より高価ではありますが、8x NVIDIA A100 GPU を搭載しているため、良い選択肢です。また、トレーニングセットとテストセットでトラックするメトリックス (交差エントロピーとパープレキシティ) を示します。これらのメトリックスの値は、Regex 式でキャプチャされ、Amazon CloudWatch に送信されます。

# トレーニングジョブに使用するコンテナイメージを定義するtraining_image_uri = f"763104351884.dkr.ecr.{region}.amazonaws.com/huggingface-pytorch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04-v1.1"# CloudWatch に送信するメトリックスを定義するmetrics = [    # トレーニングセット    {"Name": "train:LanguageCrossEntropy",     "Regex": "Train metrics\/train\/LanguageCrossEntropy: ([+-]?((\d+\.?\d*)|(\.\d+)))"},    {"Name": "train:LanguagePerplexity",     "Regex": "Train metrics\/train\/LanguagePerplexity: ([+-]?((\d+\.?\d*)|(\.\d+)))"},    # テストセット    {"Name": "test:LanguageCrossEntropy",     "Regex": "Eval metrics\/eval\/LanguageCrossEntropy: ([+-]?((\d+\.?\d*)|(\.\d+)))"},    {"Name": "test:LanguagePerplexity",     "Regex": "Eval metrics\/eval\/LanguagePerplexity: ([+-]?((\d+\.?\d*)|(\.\d+)))"},]estimator_args = {    "image_uri": training_image_uri,     # トレーニングコンテナイメージ    "entry_point": "launcher.sh",        # ランチャー bash スクリプト    "source_dir": ".",                   # ランチャースクリプトと構成ファイルを含むディレクトリ    "instance_type": "ml.g5.48xlarge",   # インスタンスタイプ    "instance_count": 1,                 # トレーニングインスタンスの数    "base_job_name": "fine-tune-mpt-7b", # トレーニングジョブ名のプレフィックス    "role": role,                        # IAM ロール    "volume_size": 300,                  # インスタンスにアタッチされた EBS ボリュームのサイズ (GB)    "py_version": "py310",               # Python バージョン    "metric_definitions": metrics,       # トラックするメトリックス    "output_path": output_path,          # モデルアーティファクトがアップロードされる S3 の場所    "code_location": code_location,      # ソースコードが保存される S3 の場所    "disable_profiler": True,            # プロファイラインスタンスを作成しない    "keep_alive_period_in_seconds": 240, # 実験中に Warm Pools を有効にする}huggingface_estimator = HuggingFace(**estimator_args)

⚠️ SageMakerトレーニングの各クォータと、このクールな機能を使用する場合はWarm Poolsのクォータを要求することを確認してください。SageMakerで多くのジョブを実行する予定の場合は、SageMaker Saving Plansを見てください。

• トレーニングジョブを開始 🚀 — Amazon SageMakerでトレーニングジョブを開始するために必要なものがすべて揃いました:

huggingface_estimator.fit({    "train": TrainingInput(        s3_data=training_data_path,        content_type="application/jsonlines"),    "test": TrainingInput(        s3_data=test_data_path,        content_type="application/jsonlines"),}, wait=True)

トレーニング時間は、データセットのサイズに依存します。ダミーデータセットでは、トレーニングにはおよそ20分かかります。モデルがトレーニングされ、😊 HuggingFace形式に変換されたら、SageMakerはモデルtarボール（model.tar.gz）をS3 output_pathにアップロードします。実際には、アップロードステップにはかなりの時間（>1時間）がかかることがわかりました。これは、圧縮するモデルアーティファクトのサイズに起因する可能性があります（約25GB）。

4. 概要

この記事では、MPT-7Bをカスタマイズして、SageMakerでデータセットを準備してトレーニングジョブを作成する方法を説明しました。実装は、LLM Foundryのトレーニングスクリプトを活用し、Composerライブラリの分散トレーニングランチャーを使用しています。モデルをカスタマイズして展開したい場合は、Philipp Schmidのブログ投稿をチェックすることをお勧めします。SageMakerでLLMを展開する方法に関する多数の例があります。カスタマイズされたMPT-7Bモデルをお楽しみください！🎉

この記事で使用されたすべてのコードは、Githubで入手できます:

GitHub – jpcpereira/sagemaker-fine-tune-mpt-7b

github.com

— João Pereira

お読みいただきありがとうございます。この記事が、Amazon SageMakerでMPT-7Bのような大規模言語モデルを微調整するための手助けになることを願っています。私の将来の記事を読みたい場合は、フォローしてください。フィードバックは大変ありがたいです！何か質問などがあれば、コメントを残してください。もしくはメールまたはLinkedInで直接連絡してください。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!