「Llama2とAmazon SageMakerを使用したLoRAのファインチューニングモデルのモデル管理」

「Llama2とAmazon SageMakerを活用したLoRAのファインチューニングモデルのモデル管理」

ビッグデータとAIの時代において、企業は競争力を得るためにこれらの技術を使用する方法を常に探しています。現在、AIの最も注目されている分野の一つが生成AIです。その理由は十分に理解できます。生成AIは、創造性と革新の可能性を押し広げる強力なソリューションを提供します。これらの最先端のソリューションの核には、ファウンデーションモデル(FM)があります。FMは、膨大な量のデータに事前学習された高度な機械学習モデルです。これらのファウンデーションモデルの多くは、人間のようなテキストを理解し生成する驚異的な能力を示しており、コンテンツ作成から顧客サポートの自動化まで、さまざまなアプリケーションにおいて貴重なツールとなっています。

しかし、これらのモデルには課題も存在します。それらは非常に大きく、トレーニングには大量のデータと計算リソースが必要です。さらに、トレーニングプロセスの最適化やパラメータの調整は複雑で反復的なプロセスとなり、専門知識と慎重な実験が必要です。これらは、独自のファウンデーションモデルを構築しようとする多くの組織にとってハードルとなる場合があります。この課題に対処するために、多くの顧客は既存のファウンデーションモデルを微調整することを検討しています。これは、モデルの一部のパラメータを特定のアプリケーションに合わせて調整する人気のある手法ですが、モデルにすでにエンコードされている知識を保持しながら適用することができます。これにより、組織はこれらのモデルのパワーを活用する一方で、特定のドメインやタスクにカスタマイズするために必要なリソースを削減することができます。

ファウンデーションモデルを微調整するための2つの主要なアプローチがあります。それは、従来の微調整とパラメータ効率の高い微調整です。従来の微調整では、特定の下流タスクに対して事前学習モデルのすべてのパラメータを更新します。一方、パラメータ効率の高い微調整では、元のモデルのすべてのパラメータを更新せずにモデルのカスタマイズを可能にするさまざまな技術があります。そのような技術の一つがLow-rank Adaptation（LoRA）です。これは、小さなタスク固有のモジュールを事前学習済みモデルに追加し、残りのパラメータを固定したままそれらをトレーニングする手法です。以下の画像に示すようになります。

• 「大規模な言語モデルを使ったフェイクニュースの検出」を活用する
• 「SAS認定マシンラーニングエンジニアになるために必要なすべて」
• 「Amazon SageMaker JumpStartでMistral 7Bを調整して展開する」

出典: Generative AI on AWS (O’Reilly, 2023)

最近、LoRAはいくつかの理由から人気があります。より高速なトレーニング、低メモリ要件、および事前トレーニング済みモデルを複数の下流タスクに再利用できる能力を提供します。さらに重要なのは、ベースモデルとアダプタを別々に保存し、いつでも組み合わせることができるため、微調整されたバージョンの保存、配布、共有が容易になることです。ただし、これには新たな課題も生じます。ベースモデルとアダプタを組み合わせるか別々に保持するか、どのようにこれらの新しい微調整モデルを適切に管理するかです。この記事では、この新興の問題に対処するために、Amazon SageMaker上でLoRA微調整モデルを適切に管理するためのベストプラクティスを解説します。

SageMakerモデルレジストリでのFMの取り扱い

この記事では、QLoRAメソッドを使用して、Llama2大規模言語モデル(LLM)の微調整のエンドツーエンドの例を説明します。 QLoRAは、パラメータ効率の高い微調整と4ビット/8ビットの量子化を組み合わせ、FMを特定のタスクやユースケースに微調整するために必要なリソースをさらに削減するメソッドです。これには、事前トレーニングされた70億パラメータのLlama2モデルを使用し、databricks-dolly-15kデータセットで微調整します。Llama2のようなLLMは数十億のパラメータを持ち、巨大なテキストデータセットで事前トレーニングされています。微調整は、より小さなデータセットを使用してLLMを下流タスクに適応させるものです。ただし、大規模なモデルの微調整には計算リソースが高いため、QLoRAメソッドを使用してウェイトを量子化し、トレーニング時の計算コストを削減します。

例では、2つのノートブック(llm-finetune-combined-with-registry.ipynbおよびllm-finetune-separate-with-registry.ipynb)が用意されています。それぞれは、以下のダイアグラムに示されているような異なる方法でLoRA微調整モデルを処理する方法を説明します。

1. まず、SageMaker Studio Notebooksを使用して、7,000,000,000のパラメータを持つ事前トレーニング済みのLlama2モデルをダウンロードします。Llama2などのLLMは、ドメイン固有のデータでファインチューニングされた場合に、自然言語処理（NLP）のタスクで最先端のパフォーマンスを示しています。
2. 次に、QLoRAメソッドを使用して、databricks-dolly-15kデータセットでLlama2をファインチューニングします。 QLoRAは、モデルの重みを量子化することにより、ファインチューニングの計算コストを削減します。
3. ファインチューニング中、SageMaker Experiments PlusをTransformers APIと統合して、勾配、損失などのメトリクスを自動ログします。
4. 次に、2つのアプローチを使用して、ファインチューニングされたLlama2モデルをSageMaker Model Registryにバージョン管理します：
   1. フルモデルを保存する
   2. アダプターとベースモデルを別々に保存する。
5. 最後に、Deep Java Library（DJL）Servingを使用して、SageMakerリアルタイムエンドポイント上でファインチューニングされたLlama2モデルをホストします。

次のセクションでは、これらの手順のそれぞれについて詳しく説明し、SageMakerの柔軟性と異なるLLMワークフローにおけるこれらの機能がモデルの操作を改善するのにどのように役立つかを示します。

前提条件

以下の前提条件を完了して、コードの実験を開始してください。

• SageMaker Studioドメインを作成する：特にSageMaker Studioノートブックを使用して、Llama2のファインチューニングタスクを開始し、SageMaker Model Registry内でモデルを登録および表示するために使用されます。 SageMaker Model Registry。 SageMaker Experimentsを使用して、Llama2のファインチューニングジョブログ（トレーニング損失/テスト損失など）を表示および比較するためにも使用されます。
• Amazon Simple Storage Service（S3）バケットを作成する：S3バケットへのアクセスが必要であり、トレーニングアーティファクトとモデルの重みを保存するためのS3バケットが必要です。手順については、バケットの作成を参照してください。この投稿ではサンプルコードはSageMakerのデフォルトのS3バケットを使用しますが、関連するS3バケットを使用するようにカスタマイズすることもできます。
• モデルコレクションを設定する（IAMの権限）：SageMaker Execution Roleを更新して、モデルコレクションを使用してModel Registryグループ化を実装するためのModel Registry Collectionsデベロッパーガイドでリストされているリソースグループへのアクセス権を付与します。
• Llama2の利用規約を承認する：Llama2基盤モデルを使用するために、エンドユーザーライセンス契約と使用許可ポリシーに同意する必要があります。

これらの例はGitHubリポジトリで利用できます。. ノートブックファイルは、PyTorch 2.0.0 Python 3.10 GPU最適化カーネルとml.g4dn.xlargeインスタンスタイプで実行されるStudioノートブックを使用してテストされています。

Experiments plusコールバックの統合

Amazon SageMaker Experimentsでは、SageMaker Python SDKまたはboto3を使用して、統合開発環境（IDE）から機械学習（ML）の実験とモデルバージョンの組織化、追跡、比較、評価を行うことができます。SageMaker Experimentsは、モデルのメトリクス、パラメータ、ファイル、アーティファクトをログに記録し、異なるメトリクスのプロットチャート、さまざまなメタデータのキャプチャ、それらを検索し、モデルの再現性をサポートします。データサイエンティストは、ビジュアルチャートやテーブルを通じてモデルのパフォーマンスとハイパーパラメータを迅速に比較できます。また、SageMaker Experimentsを使用して作成したチャートをダウンロードし、モデル評価を関係者と共有することもできます。

LLMのトレーニングは、遅くて高価で反復的なプロセスです。ユーザーは大規模なLLM実験を追跡して、一貫性のないモデル調整の経験を防ぐために、追跡することが非常に重要です。 HuggingFace transformers APIを使用すれば、Callbacksを介してトレーニングタスク中のメトリクスを追跡することができます。コールバックは、トレーニングループの動作をカスタマイズできる「読み取り専用」のコードであり、トレーニングループの状態を検査して進行状況を報告し、TensorBoardまたはSageMaker Experiments Plusを介してカスタムロジックを使用してログを行うことができます（このコードベースの一部として含まれています）。

この投稿のコードリポジトリに含まれるSageMaker Experiments コールバックコードを以下のコードブロックに示します:

# Experiments コールバックのカスタム実装をインポートfrom smexperiments_callback import SageMakerExperimentsCallback......# SageMaker Experiments コールバックを使用して Trainer インスタンスを作成trainer = Trainer(    model=model,    args=training_args,    train_dataset=train_dataset,    eval_dataset=validation_dataset,    data_collator=default_data_collator,    callbacks=[SageMakerExperimentsCallback] # Experiments プラス コールバック関数を追加)

このコールバックは、トレーニングランの一部として以下の情報をSageMaker Experiments に自動的に記録します:

• トレーニングパラメータとハイパーパラメータ
• ステップ、エポック、最終のモデルトレーニングとバリデーションのロス
• モデルの入力および出力アーティファクト（トレーニングデータセット、バリデーションデータセット、モデルの出力場所、トレーニングデバッガーなど）

以下のグラフは、この情報を使用して表示できるチャートの例を示しています。

これにより、SageMaker Experiments の解析機能を使用して複数のランを簡単に比較できます。比較したい実験ランを選択すると、自動的に比較グラフに表示されます。

ファインチューニングモデルをモデルレジストリコレクションに登録する

モデルレジストリコレクションは、SageMaker モデルレジストリの機能であり、関連する登録済みモデルをグループ化し、モデルのスケールでの探索性を向上させるために階層的に整理することができます。モデルレジストリコレクションを使用して、ベースモデルとファインチューニングされたバリアントを追跡します。

フルモデルコピーのメソッド

最初のメソッドはベースモデルと LoRA アダプターを組み合わせてフルファインチューニングモデルを保存します。以下のコードはモデルのマージプロセスを示し、model.save_pretrained()を使用して結合モデルを保存します。

if args.merge_weights:            trainer.model.save_pretrained(temp_dir, safe_serialization=False)    # メモリをクリア    del model    del trainer    torch.cuda.empty_cache()        from peft import AutoPeftModelForCausalLM    # モデルを fp16 で読み込む    model = AutoPeftModelForCausalLM.from_pretrained(        temp_dir,        low_cpu_mem_usage=True,        torch_dtype=torch.float16,    )      # LoRA とベースモデルをマージして保存    model = model.merge_and_unload()            model.save_pretrained(        args.sm_model_dir, safe_serialization=True, max_shard_size="2GB"    )

LoRA アダプターとベースモデルをファインチューニング後に単一のモデルアーティファクトに結合することには利点と欠点があります。結合モデルは自己完結型であり、元のベースモデルを必要とせずに独立して管理および展開することができます。モデルは、ベースモデルとファインチューニングデータを反映したバージョン名を持つ単独のエンティティとして追跡できます。モデルグループの整理には、base_model_name + ファインチューニングされたdataset_nameの命名規則を採用することができます。オプションで、モデルコレクションでは元のモデルとファインチューニングモデルを関連付けることもできますが、これは結合モデルが独立しているため必要ではない場合もあります。以下のコードスニペットでは、ファインチューニングされたモデルを登録する方法を示しています。

# Model Package Group Varsft_package_group_name = f"{model_id.replace('/', '--')}-{dataset_name}"ft_package_group_desc = "QLoRA for model Mikael110/llama-2-7b-{dataset_name}-fp16".........model_package_group_input_dict = {    "ModelPackageGroupName" : ft_package_group_name,    "ModelPackageGroupDescription" : ft_package_group_desc,    "Tags": ft_tags}create_model_pacakge_group_response = sm_client.create_model_package_group(**model_package_group_input_dict)

モデルを登録するために、トレーニングエスティメータを使用できます。

inference_image_uri = sagemaker.image_uris.retrieve(    "djl-deepspeed", region=region, version="0.23.0")print(f"使用するイメージは ---- > {inference_image_uri}")model_package = huggingface_estimator.register(    content_types=["application/json"],    response_types=["application/json"],    inference_instances=[        "ml.p2.16xlarge", .........    ],    image_uri = inference_image_uri,    customer_metadata_properties = {"training-image-uri": huggingface_estimator.training_image_uri()},  #トレーニングイメージ URL を保存    model_package_group_name=ft_model_pkg_group_name,    approval_status="Approved")model_package_arn = model_package.model_package_arnprint("モデルパッケージ ARN : ", model_package_arn)

モデルレジストリから、モデルパッケージを取得してそのモデルを直接デプロイすることができます。

endpoint_name = f"{name_from_base(model_group_for_base)}-endpoint"model_package.deploy(    initial_instance_count=1,    instance_type="ml.g5.12xlarge",    endpoint_name=endpoint_name)

ただし、このアプローチには欠点もあります。モデルを組み合わせると、ベースモデルが各ファインチューニングバージョンで重複するため、ストレージの非効率性と冗長性が発生します。モデルのサイズとファインチューニングモデルの数が増えると、ストレージの需要が指数関数的に増加します。例えば、llama2 7bモデルの場合、ベースモデルは約13GBで、ファインチューニングモデルは13.6GBです。ファインチューニング後にモデルの96%が重複する必要があります。さらに、非常に大きなモデルファイルの配布と共有も困難になり、モデルのサイズとファインチューニングジョブの増加に伴い、ファイル転送と管理コストが増加するため、運用上の課題も発生します。

アダプターとベースメソッドを分離する

2番目の方法は、ランタイム時にベースウェイトとアダプターウェイトを別々のモデルコンポーネントとして保存し、シーケンシャルに読み込むことに焦点を当てています。

    ..    ..    ..    else:           # ファインチューニングされたLoRAモデルと、推論用のトークナイザーを保存        trainer.model.save_pretrained(            args.sm_model_dir,             safe_serialization=True        )    tokenizer.save_pretrained(        args.sm_model_dir    )

ベースウェイトとアダプターウェイトを保存することには、フルモデルコピーの方法と同様の利点と欠点があります。1つの利点は、ストレージスペースを節約できることです。ファインチューニングモデルの最も大きなコンポーネントであるベースウェイトは、1度だけ保存され、異なるタスクに対してチューニングされた他のアダプターウェイトと再利用することができます。例えば、Llama2-7Bのベースウェイトは約13GBですが、各ファインチューニングタスクでは約0.6GBのアダプターウェイトのみを保存する必要があり、95%のスペースを節約できます。もう1つの利点は、ベースウェイトをアダプターウェイトから別々に管理できるため、ベースウェイト専用のモデルレジストリを使用できることです。これは、ベースウェイトをインターネットを介さずにアクセスできるVPCのみモードで稼働しているSageMakerドメインにおいて便利です。

ベースウェイトのためのモデルパッケージグループを作成する

### モデルパッケージグループを作成するbase_package_group_name = model_id.replace('/', '--')base_package_group_desc = "Source: https://huggingface.co/Mikael110/llama-2-7b-guanaco-fp16".........model_package_group_input_dict = {    "ModelPackageGroupName" : base_package_group_name,    "ModelPackageGroupDescription" : base_package_group_desc,    "Tags": base_tags}create_model_pacakge_group_response = sm_client.create_model_package_group(**model_package_group_input_dict)>>>Created ModelPackageGroup Arn : arn:aws:sagemaker:us-west-2:376678947624:model-package-group/Mikael110--llama-2-7b-guanaco-fp16.........### ベースウェイトの登録from sagemaker.huggingface import HuggingFaceModel# Hugging Faceモデルクラスを作成するhuggingface_model = HuggingFaceModel(    transformers_version='4.28',    pytorch_version='2.0',    py_version='py310',    model_data=model_data_uri, # ベースウェイトのS3パス（*.tar.gz形式）    role=role,)_response = huggingface_model.register(    content_types=["application/json"],    response_types=["application/json"],    inference_instances=[    "ml.p2.16xlarge",    ...    ],    transform_instances=[    "ml.p2.16xlarge",    ...    ],    model_package_group_name=base_model_pkg_group_name,    approval_status="Approved" )

QLoRAの重み用にモデルパッケージグループを作成する

以下のコードは、QLoRAの重みにデータセット/タスクの種類をタグ付けし、微調整済みのデルタ重みを別個のモデルレジストリに登録し、デルタ重みを個別に追跡する方法を示しています。

### デルタ重みのためのモデルパッケージグループを作成するft_package_group_name = f"{model_id.replace('/', '--')}-finetuned-sql"ft_package_group_desc = "モデル Mikael110/llama-2-7b-guanaco-fp16の QLoRA"ft_tags = [    {    "Key": "modelType",    "Value": "QLoRAModel"    },    {    "Key": "fineTuned",    "Value": "True"    },    {    "Key": "sourceDataset",    "Value": f"{dataset_name}"    }]model_package_group_input_dict = {    "ModelPackageGroupName" : ft_package_group_name,    "ModelPackageGroupDescription" : ft_package_group_desc,    "Tags": ft_tags}create_model_pacakge_group_response = sm_client.create_model_package_group(**model_package_group_input_dict)print(f'作成されたModelPackageGroup Arn : {create_model_pacakge_group_response["ModelPackageGroupArn"]}')ft_model_pkg_group_name = create_model_pacakge_group_response["ModelPackageGroupArn"]>>> 作成されたModelPackageGroup Arn : arn:aws:sagemaker:us-east-1:811828458885:model-package-group/mikael110--llama-2-7b-guanaco-fp16-finetuned-sql.........### デルタ重みを登録する QLoRA モデルの重みhuggingface_model = HuggingFaceModel(    transformers_version='4.28',    pytorch_version='2.0',      py_version='py310',    model_data="s3://sagemaker-us-east-1-811828458885/huggingface-qlora-2308180454/output/model.tar.gz", OR #huggingface_estimator.model_data    role=role,)_response = huggingface_model.register(    content_types=["application/json"],    response_types=["application/json"],    inference_instances=[    "ml.p2.16xlarge",    ...    ],    transform_instances=[    "ml.p2.16xlarge",    ...    ],    model_package_group_name=ft_model_pkg_group_name,    approval_status="Approved")>>> モデルコレクションの作成状態: {'added_groups': ['arn:aws:sagemaker:us-east-1:811828458885:model-package-group/mikael110--llama-2-7b-guanaco-fp16-finetuned-sql'], 'failure': []}

以下のスニペットは、モデルレジストリからモデルをベースと微調整の重みに分割したビューを示しています。

ハイパーカスタマイズされたLLM用のモデル、データセット、タスクの管理はすぐに多くの負荷となる場合があります。 SageMakerモデルレジストリコレクションを使用すると、関連するモデルをグループ化して階層的に整理し、モデルの探索性を向上させることができます。これにより、ベースの重み、アダプターの重み、および微調整タスクのデータセット間の関係を追跡しやすくなります。また、モデル間の複雑な関係とリンクを作成することもできます。

新しいコレクションを作成し、ベースのモデルの重みをこのコレクションに追加する

# モデルコレクションを作成するbase_collection = model_collector.create(    collection_name=model_group_for_base # 例: "Website_Customer_QnA_Bot_Model")# ベースの重みをモデルコレクションの最初のレベルに追加する_weights_response = model_collector.add_model_groups(    collection_name=base_collection["Arn"],    model_groups=[base_model_pkg_group_name])print(f"モデルコレクションの作成状態: {_response}")>>> モデルコレクションの作成状態: {'added_groups': ['arn:aws:sagemaker:us-west-2:376678947624:model-package-group/Mikael110--llama-2-7b-guanaco-fp16'], 'failure': []}

タスクと/またはデータセットによって、Fine-Tuned LoRA Adapter Delta Weightsをこのコレクションにリンクする

# create model collection for finetuned and link it back to the basefinetuned_collection = model_collector.create(    collection_name=model_group_for_finetune,    parent_collection_name=model_group_for_base)# add finetuned model package group to the new finetuned collection_response = model_collector.add_model_groups(    collection_name=model_group_for_finetune,    model_groups=[ft_model_pkg_group_name])print(f"Model collection creation status: {_response}")>>>モデルコレクション作成のステータス: {'added_groups': ['arn:aws:sagemaker:us-east-1:811828458885:model-package-group/mikael110--llama-2-7b-guanaco-fp16-finetuned-sql'], 'failure': []}

これにより、モデル/タスクタイプとベースモデルをファインチューニングするために使用されたデータセットによってリンクされるコレクションの階層が作成されます。

ベースモデルとアダプタモデルを分離するこの方法にはいくつかの欠点があります。1つの欠点は、モデルの展開の複雑さです。別々のモデルアーティファクトがあるため、モデルレジストリから直接デプロイする代わりに、モデルを再パッケージするための追加のステップが必要です。次のコード例では、まずベースモデルの最新バージョンをダウンロードして再パッケージします。

!aws s3 cp {base_model_package.model_data} .!tar -xvf {model_tar_filename} -C ./deepspeed/!mv ./deepspeed/{model_id} ./deepspeed/base!rm -rf ./deepspeed/{model_id}

次に、最新のファインチューンされたLoRAアダプタの重みをダウンロードして再パッケージします。

!aws s3 cp {LoRA_package.model_data} .!mkdir -p ./deepspeed/lora/!tar -xzf model.tar.gz -C ./deepspeed/lora/

モデルをホストするために、DJLサービングとdeepspeedを使用する予定の場合、推論ディレクトリは次のようになる必要があります。

deepspeed    |-serving.properties    |-requirements.txt    |-model.py    |-base/        |-...    |-lora/        |-...

最後に、カスタム推論コード、ベースモデル、およびLoRAアダプタを単一の.tar.gzファイルにパッケージ化してデプロイします。

!rm -f model.tar.gz!tar czvf model.tar.gz -C deepspeed .s3_code_artifact_deepspeed = sagemaker_session.upload_data("model.tar.gz", default_bucket, f"{s3_key_prefix}/inference")print(f"S3 Code or Model tar for deepspeed uploaded to --- > {s3_code_artifact_deepspeed}")

クリーンアップ

ノートブックのクリーンアップセクションの手順に従ってリソースをクリーンアップしてください。推論インスタンスのコストについては、Amazon SageMaker Pricingを参照してください。

結論

この記事では、Amazon SageMaker上でLoRAファインチューンモデルを管理するためのベストプラクティスについて説明しました。ベースとアダプタの重みを1つの自己完結型モデルに組み合わせる方法と、ベースとアダプタの重みを分離する方法の2つの主要な方法をカバーしました。どちらのアプローチにもトレードオフがありますが、重みを分離することでストレージの最適化が可能になり、SageMakerモデルレジストリコレクションのような高度なモデル管理技術を使用して組織化と検索性を向上させることができます。これにより、実験を追跡し、タスクに適したモデルを見つけ、スケール時に効率的に特殊なLLMを管理することができます。試してみるために、GitHubリポジトリのサンプルコードを使用して、これらの方法を実験してみることをお勧めします。ジェネレーティブAIが急速に進化する中で、モデル管理のベストプラクティスに従うことで、実験を追跡し、タスクに適したモデルを見つけ、スケール時に特化したLLMを効率的に管理することができます。

参考文献

• 次世代Amazon SageMaker Experiments
• HuggingFace Callbacks

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