Amazon SageMaker で大規模なモデル推論 DLC を使用して Falcon-40B をデプロイする

Deploy Falcon-40B using Amazon SageMaker's large-scale model inference DLC.

先週、Technology Innovation Institute（TII）はTII Falcon LLMを発表しました。これはオープンソースの基礎的な大規模言語モデル（LLM）です。Amazon SageMakerで1兆トークンでトレーニングされたFalconは、他のLLM（llama-65Bなど）よりも比較的軽量でホスティングにかかる費用が少なく、トップクラスのパフォーマンスを誇ります（執筆時点でHugging Faceリーダーボードで1位）。この記事では、SageMaker上の大規模モデル推論ディープラーニングコンテナを使用して、Falconを言語理解や自動書き換え支援などのアプリケーションに展開する方法を説明します。

FalconがSageMakerに着陸しました

TIIはアブダビのAdvanced Technology Research Council内の応用研究機関であり、科学者、研究者、エンジニアのチームは、社会を未来に対応できるようにするための革新的な技術の発見と科学的な飛躍の開発に専念しています。今年早々、TIIは最先端のオープンソースLLMをトレーニングすることを目指し、SageMakerのインフラストラクチャ、ツール、専門知識を使用しました（このモデルがSageMakerでどのようにトレーニングされたかについては、「Technology Innovation Institute trains the state-of-the-art Falcon LLM 40B foundation model on Amazon SageMaker」を参照してください）。この努力の結果、TII Falcon LLMが誕生しました。

Falconは1兆トークンでトレーニングされ、Eleuther AI Language Model Evaluation Harnessに対して最高のパフォーマンスを発揮し、現在Hugging Faceリーダーボードで1位です。このモデルはFalcon-40BとFalcon-7Bの2つの異なるサイズで利用可能であり、言語理解、会話体験、自動書き換え支援などのアプリケーションで最先端のパフォーマンスを発揮できます。この記事では、これらのドメインにおける高精度推論のために、SageMaker上でFalconを展開する方法を説明します。

SageMaker大規模モデル推論DLCはLLMホスティングを簡素化します

Falcon-40BやFalcon-7BなどのLLMをホストすることは困難です。大きなモデルは数十億のパラメータを含むため、通常よりも正確ですが、サイズが大きいため推論レイテンシーが低下したり、スループットが悪化することがあります。LLMをホストするには、より多くのGPUメモリと最適化されたカーネルが必要です。また、Falcon-7Bなどの小さなモデルは、AWS G5インスタンスタイプを駆動するNVIDIA A10Gインスタンスのような単一のGPUに一般的に適合しますが、Falcon-40Bのような大きなモデルは適合しません。その場合、テンソル並列処理などの戦略を使用して、その大きなモデルを複数の部分に分割し、複数のGPUのメモリを利用する必要があります。小さなモデルに対して使用される従来のホスティングソリューションには、このような機能がないため、難易度が高くなります。

• Amazon SageMakerを使用した生成型AIモデルにおいて、Forethoughtがコストを66％以上削減する方法
• BrainPadがAmazon Kendraを使用して内部の知識共有を促進する方法
• Amazon TranslateのActive Custom Translationを使用して、マルチリンガル自動翻訳パイプラインを構築します

SageMaker大規模モデル推論（LMI）ディープラーニングコンテナ（DLC）を使用すると、これらの問題を解決できます。LMI DLCは、Falcon-40BなどのLLMをホストするための完全なエンドツーエンドソリューションです。フロントエンドには、トークンストリーミングやインスタンス内での自動モデルレプリケーションなどの大規模モデル推論に最適化された高性能モデルサーバー（DJL Serving）が含まれています。バックエンドには、DeepSpeedやFasterTransformerなどの高性能モデルパラレルエンジンが複数含まれており、複数のGPUにまたがるモデルパラメータを分割して管理できます。これらのエンジンには、人気のあるトランスフォーマーモデルの最適化されたカーネルも含まれており、推論を最大3倍高速化することができます。LMI DLCを使用すると、SageMaker上でLLMをホストするための構成ファイルを作成するだけで済みます。SageMaker LMI DLCについては、「Model parallelism and large model inference」と「available images」のリストを参照してください。また、LMI DLCを使用してSageMakerでBloom-175Bをホストする方法については、以前の記事をご覧ください。

ソリューション概要

この記事では、DeepSpeedを使用してSageMaker上でFalcon-40Bをホストする方法を説明します。Falcon-40Bをホストするには、複数のA10 GPUを使用する必要がありますが、Falcon-7Bは単一のGPUで済みます。また、DeepSpeedとAccelerateの両方を使用してFalcon-40BとFalcon-7Bをホストするための例も用意しています。コード例はGitHubでご確認いただけます。

この例は、SageMakerノートブックインスタンスまたはAmazon SageMaker Studioノートブックで実行できます。LMIとDeepSpeedを使用してFalcon-40Bをホストするには、ml.g5.24xlargeインスタンスを使用する必要があります。これらのインスタンスには、それぞれ96 GiBのGPUメモリをサポートする4つのNVIDIA A10G GPUが搭載されています。また、ホストには96 vCPUと384 GiBのホストメモリがあります。LMIコンテナは、モデルをダウンロードし、その構成要素であるパラメータを複数のGPUに分散するためのモデルアーティファクトをパーティションに分割するなど、LLMをホストするためのオペレーションの大部分を処理することができます。

SageMakerの機械学習（ML）インスタンスのクオータはアカウントによって異なる場合があります。このポストに従っているときにg5.24xlargeインスタンスのクオータを超過したエラーが表示された場合は、Service Quotasコンソールを介して制限を増やすことができます。

ノートブックの解説

まず、例に必要な依存関係をインストールしてインポートします。Boto3 SDKとSageMaker SDKを使用します。SageMakerとLMIで使用するためのモデルアーティファクトを保存するためにAmazon Simple Storage Service（Amazon S3）を使用するため、S3プレフィックス変数を設定します。以下のコードを参照してください。

import sagemaker
import jinja2
from sagemaker import image_uris
import boto3
import os
import time
import json
from pathlib import Path
from sagemaker.utils import name_from_base

role = sagemaker.get_execution_role()  # エンドポイントの実行ロール
sess = sagemaker.session.Session()  # 異なるAWS APIとやり取りするためのSageMakerセッション
bucket = sess.default_bucket()  # アーティファクトを収容するバケット
model_bucket = sess.default_bucket()  # アーティファクトを収容するバケット
s3_code_prefix_deepspeed = "hf-large-model-djl-/code_falcon40b/deepspeed"  # コードアーティファクトが格納されるバケット内のフォルダー
region = sess._region_name
account_id = sess.account_id()
s3_client = boto3.client("s3")
sm_client = boto3.client("sagemaker")
smr_client = boto3.client("sagemaker-runtime")
jinja_env = jinja2.Environment()

次に、ワークスペースのためのローカルフォルダーを作成して、モデルアーティファクトを保存します。

!mkdir -p code_falcon40b_deepspeed

最初に、ローカルに作成したディレクトリにserving.properties設定ファイルを作成します。このserving.propertiesファイルは、LMIコンテナーとフロントエンドのDJL Servingライブラリに、どのモデル並列化と推論最適化エンジンを使用するかを示しています。DeepSpeedとHugging Face Accelerateの両方の構成オプションはConfigurations and settingsにあります。ここで、option.model_idパラメータを設定して、使用するHugging Faceモデルを定義します。SageMakerはHugging Faceモデルを扱うことが簡単で、この1行だけで済みます。さらに、4つのGPUを持つml.g5.24xlargeインスタンスを使用しているため、option.tensor_parallel_degreeを値4に設定します。このパラメータは、生成するパーティションの数を定義し、分散化します。ml.g5.48xlargeなど8つのGPUを持つより大きなインスタンスを使用し、値を4に設定した場合、LMIは自動的にモデルの2つのレプリカを作成します（それぞれが4つのGPUに分散されます）。以下のコードを参照してください。

%%writefile ./code_falcon40b_deepspeed/serving.properties
engine=Python
# falcon-40b-instructをデプロイする場合は、model_id値を'tiiuae/falcon-40b-instruct'に設定します
option.model_id=tiiuae/falcon-40b
option.tensor_parallel_degree=4
#option.s3url = {{s3url}}

必要に応じて、tiiuae/falcon-40bをtiiuae/falcon-40b-instructに置き換えることもできます。

また、必要なパッケージをインストールするために指定できるrequirements.txtファイルも含まれています。

%%writefile ./code_falcon40b_deepspeed/requirements.txt
einops
torch==2.0.1

最後に、モデルに使用されるmodel.pyファイルが必要です。

%%writefile ./code_falcon40b_deepspeed/model.py
from djl_python import Input, Output
import os
import torch
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
from typing import Any, Dict, Tuple
import warnings

predictor = None


def get_model(properties):
    model_name = properties["model_id"]
    local_rank = int(os.getenv("LOCAL_RANK", "0"))
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        low_cpu_mem_usage=True,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        device_map="auto",
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    generator = pipeline(
        task="text-generation", model=model, tokenizer=tokenizer, device_map="auto"
    )
    return generator


def handle(inputs: Input) -> None:
    global predictor
    if not predictor:
        predictor = get_model(inputs.get_properties())
    if inputs.is_empty():
        # Model server makes an empty call to warmup the model on startup
        return None
    data = inputs.get_as_json()
    text = data["text"]
    text_length = data["text_length"]
    outputs = predictor(text, do_sample=True, min_length=text_length, max_length=text_length)
    result = {"outputs": outputs}
    return Output().add_as_json(result)

これで完了です！これで、DeepSpeedを使用してFalcon-40Bをデプロイするために必要なすべてのアーティファクトを作成しました！ディレクトリを*.tar.gzファイルにパッケージ化し、Amazon S3にアップロードします。実際のモデルはダウンロードされず、このファイルにパッケージ化されていません。LMIコンテナは、Hugging Faceから直接モデルをダウンロードします。モデルのコピーをより高速にダウンロードできる場所に置く場合は、S3バケットをターゲットにするオプションもあります。LMIには、Amazon S3からの高性能なダウンロードの最適化も含まれています。次のコードを参照してください：

s3_code_artifact_deepspeed= sess.upload_data("model.tar.gz", bucket, s3_code_prefix_deepspeed)
print(f"S3 Code or Model tar for deepspeed uploaded to --- > {s3_code_artifact_deepspeed}")

これで、使用するコンテナを定義し、モデルオブジェクトを作成するだけです：

inference_image_uri = (
    f"763104351884.dkr.ecr.{region}.amazonaws.com/djl-inference:0.22.1-deepspeed0.8.3-cu118"
)
model_name_acc = name_from_base(f"falcon40b-model-ds")
create_model_response = sm_client.create_model(
    ModelName=model_name_acc,
    ExecutionRoleArn=role,
    PrimaryContainer={"Image": inference_image_uri, "ModelDataUrl": s3_code_artifact_deepspeed},
)
model_arn = create_model_response["ModelArn"]

次に、エンドポイント構成を作成し、エンドポイントを作成します：

endpoint_config_name = f"{model_name}-config"
endpoint_name = f"{model_name}-endpoint"
endpoint_config_response = sm_client.create_endpoint_config(
    EndpointConfigName=endpoint_config_name,
    ProductionVariants=[
        {
            "VariantName": "variant1",
            "ModelName": model_name,
            "InstanceType": "ml.g5.24xlarge",
            "InitialInstanceCount": 1,
            "ModelDataDownloadTimeoutInSeconds": 3600,
            "ContainerStartupHealthCheckTimeoutInSeconds": 3600,
            # "VolumeSizeInGB": 512
        },
    ],
)
endpoint_config_response

create_endpoint_response = sm_client.create_endpoint(
    EndpointName=f"{endpoint_name}", EndpointConfigName=endpoint_config_name
)
print(f"Created Endpoint: {create_endpoint_response['EndpointArn']}")

ホスティングの成功のために考慮すべき構成項目

大規模なモデルホスティングに重要な考慮事項は、Hugging Faceからのモデルのダウンロードに十分な時間を確保することです。私たちのテストでは、Falcon-40Bのインスタンスへのダウンロードに約90分かかりました。これを許容するための重要な構成セットは、ContainerStartupHealthCheckTimeoutInSecondsとModelDataDownloadTimeoutInSecondsです。SageMakerエンドポイント構成がそれぞれ3600の値を持つことを確認してください。また、LMIコンテナを使用して、S5cmdユーティリティを使用したAmazon S3からのダウンロードが元のモデルのゾオよりもはるかに簡単であり、モデルのダウンロード時間を約10分に短縮できます。

エンドポイントのステータスを確認するには、DescribeEndpointを呼び出します。すべてが完了したときに通知されます。エンドポイントは、LMIがモデルのパーティショニングとオーケストレーションを処理するため、各リクエストがml.g5.12xlargeインスタンスの4つのGPUを使用して処理されます。これにより、LLMをホストしてGPUアクセラレータを水平にスケーリングすると、パフォーマンスを向上させることができます。次のコードを参照してください：

response_model = smr_client.invoke_endpoint(
    EndpointName=endpoint_name,
    Body=json.dumps({"text": "What is the purpose of life?", "text_length": 150}),
    ContentType="application/json",
)

response_model["Body"].read().decode("utf8")

完了した場合は、エンドポイント構成、エンドポイント、およびモデルオブジェクトを削除する場合は、次のコマンドを実行できます：

sm_client.delete_endpoint(EndpointName=endpoint_name)
sm_client.delete_endpoint_config(EndpointConfigName=endpoint_config_name)
sm_client.delete_model(ModelName=model_name)

この投稿で参照したコードは、GitHubの完全なノートブックで見つけることができます。

結論

SageMaker HostingとLMI DLCは、Falcon-40BのようなLLMをホストするのを簡単にします。複数のGPUでモデルをホストするために必要なものをオーケストレーションするために、巨大で差別化された重い作業を引き受け、ニーズに合わせて構成可能なオプションを提供します。さらに、Hugging Faceモデルの使用は、これらのモデルをサポートする組み込みサポートがあるため、非常に簡単になります。

このポストでは、DeepSpeedを使用してFalcon-40Bモデルをホストする方法と、Accelerateを使用してFalcon-40B、およびより小さなFalcon-7BモデルをホストするGitHubの例を提供しました。 SageMakerでLMIを使用して、これを試して、これまでに公開された最高性能のLLMに手を出すことをお勧めします！

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