AWS Trainium を活用した日本語大規模言語モデルの分散学習と AWS Inferentia2 上での推論環境構築

はじめに

Anthropic が 2024年 3月に Claude3 を発表し、4月には Meta が Llama3 を発表するなど、生成 AI の進化は止まるところを知りません。一方、生成 AI の進化を支える大規模言語モデルの開発及び運用に掛かるコスト、計算機リソースの確保は多くの企業が抱える大きな課題です。AWS では機械学習 (ML) アクセラレーターチップ AWS Trainium、AWS Inferentia2 を自社開発し、これらの課題解決に取り組んでいます。（Anthropic では AWS Trainium、Inferentia の活用を表明しています）

本ブログでは、前半で、AWS Trainium 搭載 Amazon EC2 Trn1 インスタンスを活用した日本語大規模言語モデルの開発事例、大規模分散学習の課題及び実現方法について解説します。

ブログ後半では、公開された日本語大規模モデルを Inferentia2 搭載 Amazon EC2 Inf2 インスタンス上で推論実行する方法について、手順を追って解説します。１時間以内で完走できる内容ですので、公開された各種モデルの検証、活用に興味をお持ちの方は後半から読み進めて頂いても結構です。

AWS 自社設計 ML アクセラレーター搭載インスタンス

本章では、ブログの本題に入る前に、AWS が提供している生成 AI / ML モデルのトレーニング及び推論向けインスタンスについて紹介します。

AWS では、生成 AI / ML モデルのトレーニング向けのインスタンスとして、業界で広く利用されている NVIDIA 社 A100 Tensor Core GPU を搭載した Amazon EC2 P4d インスタンス、最新 H100 GPU を搭載した P5 インスタンスを提供しています。2024年 4月には L4 GPU を搭載した 推論向け G6 インスタンスの一般提供を開始しました。

一方、AWS は、クラウド向けに最適化されたカスタムシリコンの設計に何年も投資しており、お客様の幅広いワークロードにおいて、自社設計チップファミリー AWS Graviton プロセッサ、AWS Trainium、Inferentia 機械学習アクセラレーターが活用されています。2023年 11月には年次イベント AWS re:Invent において、それぞれの次世代版となる AWS Graviton4 プロセッサと AWS Trainium2 アクセラレーターを発表しました。

AWS Trainium 搭載 Amazon EC2 Inf2 インスタンス

Amazon EC2 Trn1 インスタンスは、生成 AI / ML モデルのトレーニング向けに開発した AWS Trainium アクセラレーターを搭載したインスタンスです。 他の同等の ML トレーニング向けインスタンスと比較し、コストを最大 50% 削減します。Trn1 インスタンスでは、小規模モデルのトレーニングを想定して Trainium を１つ搭載したtrn1.2xlargeと、大規模モデルの分散学習を想定して Trainium を 16個搭載したtrn1.32xlargeの２つのサイズを用意しています。また Trn1 インスタンスと比較してネットワーク帯域を 2 倍にし、1600 Gbps の Elastic Fabric Adapter (EFAv2) に対応したtrn1n.32xlargeインスタンスも提供しています。

Trn1 インスタンス及び Trainium アクセラレーターの概要及び、小規模モデルのトレーニング（ファインチューニング）については、「Trn1 独自設計チップ AWS Trainium 搭載 Amazon EC2 Trn1 インスタンスで ML トレーニングを高速実行（基礎編、実践編）」をご参照下さい。

AWS Inferentia2 搭載 Amazon EC2 Inf2 インスタンス

Amazon EC2 Inf2 インスタンスは、初代 Inf1 インスタンスと比較し、最大 3 倍のコンピューティング性能、最大 4 倍のアクセラレーターメモリ、最大 4 倍のスループット、10 分の 1 以下の低レイテンシーを実現します。 AWS Inferentia2 には Trainium と同じ世代となる Neuron コア v2 を搭載しており、大規模モデルの推論だけでなく、小規模モデルのファインチューニングも実行可能です。Inf2 インスタンスは既に東京リージョンでもローンチ済みなので、学習モデルを国内に保持しておきたい場合や、ネットワークレイテンシーを最小化したい場合には、東京リージョンにてデプロイ可能です。

AWS Trainium を活用した大規模言語モデルの分散学習

アマゾン ウェブ サービス ジャパンは 2023年 7月、国内に法人または拠点を持つ企業・団体の生成 AI 基盤モデル・大規模言語モデルの開発を支援する「AWS LLM 開発支援プログラム」を立ち上げました。本プログラムに採択された17 社の採択企業の多く（11社）では、AWS Trainium を活用し、短期間で日本語大規模言語モデルの開発を完了しました。ストックマーク株式会社、カラクリ株式会社、株式会社わたしは、 rinna株式会社ではプログラムを通して開発したモデルを Hugging Face Hub 上で公開しており、広くご活用頂けるようになっています。中でもstockmark-13b、karakuri-lm-70b-chat-v0.1は 2000 回 以上ダウンロードされた人気の日本語モデルとなっています。

AWS LLM 開発支援プログラムの総括については「基盤モデル開発に挑む各社が成果を共有。AWS LLM 開発支援プログラム 成果発表会」をご参照下さい。

大規模言語モデルの分散学習は、アクセラレーターチップ、インスタンスといった計算機リソースだけでは実現できません。効率的に実行するためには、分散学習ライブラリ、クラスタ管理フレームワーク、 Amazon FSx for Lustre 共有ファイルシステムなど、複数のサービスを組み合わせることが不可欠です。

本章では、はじめに Trainium 上で分散学習を実行する上で鍵となる分散学習ライブラリについて、サンプルスクリプトとともに解説します。続いて、AWS が提供するクラスタ管理フレームワークの選択肢と、大規模分散学習では避けられないノード障害に対する対策について解説していきます。

分散学習ライブラリ

AWS Neuron は Trainium、Inferentia 上での機械学習ワークロードを最適化するための SDK です。PyTorch を用いた大規模言語モデルの分散学習では、AWS Neuron 上で実行する分散トレーニングライブラリ、AWS Neuron Reference for NeMo Megatron 及び NeuronX Distributed を提供しています。その際、AWS Neuron では、PyTorch XLA をバックエンドのコンパイラとして利用しています。

AWS Neuron Reference for NeMo Megatron

AWS Neuron Reference for NeMo Megatron（以下、neuronx-nemo-megatron） は、既存のオープンソースパッケージ Nemo と Apex を AWS EC2 Trn1 インスタンスおよび AWS Neuron に対応させたライブラリです。いち早く Llama2 モデルの事前学習に対応したこともあり、2023年 8月と 9月に AWS LLM 開発支援プログラム採択企業向けに実施した Prototyping Camp では、neuronx-nemo-megatron と AWS ParallelCluster を使用し、Llama2 モデルのマルチノード分散学習環境の構築から推論実行までの一連の流れを完走しました。実際、採択企業の多くで neuronx-nemo-megatron を利用した大規模言語モデルの分散学習を実施しています。Llama2-7B、13B、70B それぞれのモデルの継続事前学習を実行するチュートリアルを用意しているため、大規模言語モデルの分散学習で必要となる環境構築の負荷を大幅に削減することが可能です。

NeuronX Distributed

NeuronX Distributed は、Neuron デバイス上で分散学習、分散推論を実行するための XLA ベースのライブラリです。モデルの大規模化に伴い、モデルを単一デバイス上でトレーニングすることは不可能になるので、シャーディング技術を利用してモデルを複数のデバイスに分割する必要があります。NeuronX Distributed では、テンソル並列、パイプライン並列、データ並列からなる 3D Parallelism と ZeRO-1 を使った Optimizer の重みのシャーディングに対応しています。また、アクティベーションメモリを削減するためのシーケンス並列と、フォワードパスでアクティベーションの一部を保存せずバックワードパス中に再計算する事でメモリの使用量を削減するアクティベーション再計算（selective activation checkpointing）にも対応しています。

以下は Llama2-70B モデルの事前学習を実行するサンプルスクリプトの一部です。こちらのサンプルスクリプトでは、trn1.32xlargeインスタンス上の合計 32 個の Neuron コア全てをテンソル並列で利用し、４つのインスタンスでパイプライン並列を、残り（16ノードでの分散学習の場合は 32 x 16 / 32 / 4 = 4）をデータ並列として利用しています。さらに ZeRO-1、シーケンス並列（以下のスクリプト内では設定していないがデフォルトで利用）、アクティベーション再計算を利用しています。Llama2-13B/70B モデルでの事前学習を実行するチュートリアルを合わせてご参照下さい。

# Global batch size
: ${GBS:=1024}
# Input sequence length
SEQ_LEN=4096
# Pipeline parallel degree
PP_DEGREE=4
# Tensor parallel degree
TP_DEGREE=32
# Data paralell size
DP=$(($PROCESSES_PER_NODE * $WORLD_SIZE / $TP_DEGREE / $PP_DEGREE))
# Batch size per model replica
BS=$(($GBS / $DP))
# Number microbatches for pipeline execution
# Setting same as BS so each microbatch contains a single datasample
NUM_MICROBATCHES=$BS
.
.
torchrun $DISTRIBUTED_ARGS run_llama_nxd.py \
--train_batch_size $BS \
--use_meta_device_init 1 \
--training_dir $DATA_PATH \
--training_config $SCRIPT_DIR/70B_config \
--max_steps $max_steps \
--seq_len $SEQ_LEN \
--pipeline_parallel_size $PP_DEGREE \
--tensor_parallel_size $TP_DEGREE \
--num_microbatches $NUM_MICROBATCHES \
--lr 0.00015 \
--min_lr 1e-05 \
--beta1 0.9 \
--beta2 0.95 \
--weight_decay 0.1 \
--warmup_steps 2000 \
--constant_steps 0 \
--use_zero1_optimizer 1 \
--use_selective_checkpoint 1 \
--qkv_linear 1 \
--kv_replicator 4 \
--tb_dir $tb_dir |& tee $LOG_PATH/log

クラスタ管理フレームワーク

採択企業の多くでは Prototyping Camp で体験頂いた AWS ParallelCluster を使って、マルチノードでの分散学習環境を構築頂きました。AWS では他にも様々な選択肢を用意しています。2023年 11月、AWS re:Invent では大規模な分散学習に特化したマネージドサービス Amazon SageMaker HyperPod を発表、ローンチしています。HyperPod については「大規模な分散トレーニングに特化したインフラストラクチャ、Amazon SageMaker HyperPod のご紹介」をご参照下さい。
また、Amazon Elastic Kubernetes Service (Amazon EKS) を用いた分散学習については「Train Llama2 with AWS Trainium on Amazon EKS」を、Amazon SageMaker を用いた分散学習については「Simple guide to training Llama 2 with AWS Trainium on Amazon SageMaker」をそれぞれ合わせてご参照下さい。

耐障害性 (resiliency) に対するアプローチ

分散学習を実施する際のノード数は、学習するモデルのパラメータサイズ、学習データのサイズ（トークン数）、学習を実施する期間や計算機リソースの確保の容易性など、様々な要素から決定されます。ノード数の規模が大きくなると、ノード障害の発生を想定した耐障害性、リカバリー機能の確保は Trn1 インスタンスに限らず、一般的に発生する切実な問題です。今回、株式会社リコー様では、64 インスタンスのtrn1.32xlarge(1,024 Trainium アクセラレーター、2,048 Neuron コア) を用いた大規模分散学習を実施頂きました。ノード障害を自動的に検知し、必要に応じて正常なノードとノード交換を実施、その後、学習ジョブを直近で保存された最新のチェックポイントから自動再開する仕組みを分散学習ライブラリ上に実装することによって、大規模クラスタでの学習を成功裏に導きました。これらの実装は最新版 neuronx-nemo-megatron、NeuronX Distributed ライブラリ内に既に組み込まれています。

AWS Inferentia2 上で日本語大規模言語モデルの推論環境を構築

AWS LLM 開発支援プログラムの支援を受けて、AWS Trainium 上で開発されたモデルのいくつかは、Hugging Face Hub 上で公開されています。Trainium 上で学習したモデルのデプロイ先に制限はありません。AWS Inferentia2 を搭載した Amazon EC2 Inf2 インスタンス、GPU インスタンスのどちらでも推論実行が可能です。また、GPU 上で学習されたモデルを Inf2 インスタンス上にデプロイすることも可能です。

本章では、Inferentia2 上で大規模言語モデルの分散推論を実現するための分散推論ライブラリについて解説した後、公開された日本語大規模言語モデルを Inferentia2 上にデプロイする方法を解説します。また、本章の最後では本番環境への導入に向けた次のステップを紹介します。

分散推論ライブラリ

Transformers NeuronX

大規模言語モデルの分散推論（自己回帰サンプリングによるテキスト生成）では、AWS Inferentia2 及び Trainium に最適化した Transformers NeuronX ライブラリを提供しています。小規模なモデルであればモデルを単一デバイス上にデプロイすることが可能ですが、大規模言語モデルの場合には、分散学習時と同様にテンソル並列を用いて、モデルを複数の Neuron コア間に分割して実装します。

最大サイズのinf2.48xlargeには、12 Inferentia2 アクセラレーター、24 Neuron コアが搭載されています。より高い性能（低レイテンシー）が求められる場合は、並列度を最大限上げる事が求められるので、テンソル並列度を 24 で実装することが推奨されます（さらに性能を求める場合にはtrn1.32xlargeをテンソル並列度 32 で利用することも可能）。一方、運用コストを重視する場合には、アクセラレーターメモリの容量を超えない限りの少ないテンソル並列度にする必要があります。70億、130億パラメータのモデルであれば、単一の Inferentia2 上にテンソル並列度 2 で、700億パラメータのモデルであれば、inf2.24xlargeもしくはinf2.48xlarge上に並列度 8 で実装可能です。

事前準備

70億、130億パラメータのモデルであれば、32GBのアクセラレーターメモリを持つ単一の Inferentia2 アクセラレーター上にデプロイ可能です。ここではinf2.8xlargeインスタンスを起動し、130億パラメータのモデルを実際に活用してみましょう。EC2 の立ち上げ方についてはこちらのドキュメントをご参照下さい。本ブログ内の検証は、東京リージョン (ap-northeast-1) で、以下の Deep Learning AMI Neuron (Ubuntu 22.04) AMI を利用し実施しました。また大規模モデルを扱う場合はストレージ容量に注意が必要です。ここでは EBS 512GBを用意しました。

インスタンスを起動後、SSH で接続します。モデルの推論を実行する際に Jupyter Notebook を使用しますので、SSH ポートフォワーディングを設定しておきましょう。

ssh -i "<pem file>" ubuntu@<instance DNS name> -L 8888:127.0.0.1:8888

インスタンスに接続すると、以下のような画面が表示されます。2024年 5月 10日現在、最新となる Neuron 2.18.2 がプリインストールされていることが分かります。

Welcome to Ubuntu 22.04.4 LTS (GNU/Linux 5.15.0-1031-aws x86_64)

 * Documentation:  https://help.ubuntu.com
 * Management:     https://landscape.canonical.com
 * Support:        https://ubuntu.com/pro

=============================================================================
       __|  __|_  )
       _|  (     /   Deep Learning AMI Neuron 2.18.2 (Ubuntu 22.04)
      ___|\___|___|
=============================================================================

* Supported EC2 instances: Inf1, Inf2, Trn1, Trn1n.
* To activate pre-built pytorch-1.13 environment for Inf2, Trn*, run: 'source /opt/aws_neuronx_venv_pytorch_1_13/bin/activate'
* To activate pre-built pytorch-1.13 environment for Inf1, run: 'source /opt/aws_neuron_venv_pytorch_1_13_inf1/bin/activate'
* To activate pre-built pytorch-2.1 environment for Inf2, Trn*, run: 'source /opt/aws_neuronx_venv_pytorch_2_1/bin/activate'
* To activate pre-built transformers environment for Inf2, Trn*, run: 'source /opt/aws_neuronx_venv_transformers_neuronx/bin/activate'
* To activate pre-built tensorflow-2.10 environment for Inf2, Trn*, run: 'source /opt/aws_neuronx_venv_tensorflow_2_10/bin/activate'
* To activate pre-built tensorflow-2.10 environment for Inf1, run: 'source /opt/aws_neuron_venv_tensorflow_2_10_inf1/bin/activate'
* Neuron driver version:: 2.16.7.0

Neuron は定期的にアップデートされ、新しい機能の追加、対応するモデルの拡充、性能向上、バグフィックスなどが提供されます。それでは、次のコマンドを実行して仮想環境をアクティブ化しましょう。

source /opt/aws_neuronx_venv_transformers_neuronx/bin/activate

ここで、最新の PyTorch Neuron 環境が正しくインストールできているかどうか確認しておきましょう。Neuron ドキュメント上で各パッケージの最新バージョンが確認可能です。

(aws_neuronx_venv_transformers_neuronx) ubuntu@ip-xx-xx-xx-xxx:~$ pip list | grep neuron
aws-neuronx-runtime-discovery 2.9
libneuronxla                  2.0.965
neuronx-cc                    2.13.72.0+78a426937
torch-neuronx                 2.1.2.2.1.0
transformers-neuronx          0.10.0.360

(aws_neuronx_venv_transformers_neuronx) ubuntu@ip-xx-xx-xx-xxx:~$ dpkg -l | grep neuron
hi  aws-neuronx-collectives                2.20.22.0-c101c322e                     amd64        neuron_ccom built using CMake
hi  aws-neuronx-dkms                       2.16.7.0                                amd64        aws-neuronx driver in DKMS format.
hi  aws-neuronx-oci-hook                   2.3.0.0                                 amd64        neuron_oci_hook built using CMake
hi  aws-neuronx-runtime-lib                2.20.22.0-1b3ca6425                     amd64        neuron_runtime built using CMake
hi  aws-neuronx-tools                      2.17.1.0                                amd64        Neuron profile and debug tools

130 億パラメータモデルを用いた推論実行

それでは、株式会社わたしはが公開している大喜利言語モデル（Watashiha-Llama-2-13B-Ogiri-sft）を実行してみましょう。

Jupyter Notebook を起動し、以降のセルを実行していきます。

jupyter notebook

最初にモデルをダウンロードし、Neuron コア上で動作できるようにコンパイルを実行します。コンパイルする際には、バッチサイズとテンソル並列度を指定する必要があります。inf2.8xlargeでは Inferentia2 を 1 個（Neuronコアを 2 個）搭載しているので、ここではtp_degree=2と設定しています。

from transformers_neuronx import NeuronAutoModelForCausalLM
from transformers import AutoTokenizer

model_path = "watashiha/Watashiha-Llama-2-13B-Ogiri-sft"

neuron_model = NeuronAutoModelForCausalLM.from_pretrained(
    model_path, 
    batch_size=1, 
    tp_degree=2, 
    amp='bf16'
)
neuron_model.to_neuron()

Watashiha-Llama-2-13B-Ogiri-sftは、大喜利データでファインチューニングされたモデルです。指定のフォーマットでプロンプトを用意しましょう。

prompt = """
以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。

### 指示:
入力の文は大喜利のお題です。お題に沿った面白いボケを生成してください。

### 入力:
マジシャンのショーでアシスタントが消えたまま戻ってこない時の一言。

### 応答:
"""

tokenizer = AutoTokenizer.from_pretrained("watashiha/Watashiha-Llama-2-13B-Ogiri-sft")
input_ids = tokenizer.encode(prompt, return_tensors="pt")

最後に、推論を実行します。

generated_sequences = neuron_model.sample(input_ids, sequence_length=128, top_k=50)

output = tokenizer.decode(generated_sequences[0], skip_special_tokens=True)
print(output)

推論結果の例

以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。

### 指示:
入力の文は大喜利のお題です。お題に沿った面白いボケを生成してください。

### 入力:
マジシャンのショーでアシスタントが消えたまま戻ってこない時の一言。

### 応答:
帰りの支度をしております。

面白いボケは生成されましたでしょうか？プロンプトを変更したり、繰り返し推論を実行しどのようなボケを生成できるか試してみましょう。

他にも、CyberAgentLM2-7B (CALM2-7B)、Stockmark-13b、ELYZA-japanese-Llama-2-13b-fastなど、Llama2 をベースとしたモデルであれば、Trainium で学習したモデルかGPUで学習したモデルかに関わらず、同様の手順で推論実行可能です。

ここではコンパイルを実行するためinf2.8xlargeを使用しましたが、あらかじめコンパイル済みのモデルを利用することで、推論処理自体はより安価なinf2.xlargeでも実行可能です。

700 億パラメータモデルを用いた推論実行

700 億パラメータのモデルを実行するためには、アクセラレーターメモリの制約上、複数の Inferentia2 によりテンソル並列度を上げて実行する必要があります。ここではカラクリ株式会社が公開しているKARAKURI LM 70B Chat v0.1をinf2.48xlage上で実行してみましょう。必要な手順は先程の 130 億パラメータモデルの場合と同じです。inf2.48xlageでは Inferentia2 を12 個（Neuronコアを 24 個）搭載しているので、ここでは全ての Neuron コアを活用すべく、tp_degree=24と設定しています。

from transformers_neuronx import NeuronAutoModelForCausalLM
from transformers import AutoTokenizer

model_path = "karakuri-ai/karakuri-lm-70b-chat-v0.1" 

neuron_model = NeuronAutoModelForCausalLM.from_pretrained(
    model_path, 
    batch_size=1, 
    tp_degree=24, 
    amp='bf16'
)
neuron_model.to_neuron()

先程と同様に、モデルが期待するフォーマットに合わせてプロンプトを用意します。

prompt = """以下は人間とAIアシスタントとの会話です。

Human: こんにちは。
AI: こんにちは、私はAIアシスタントです。何かお手伝いできることはありますか？
Human: 週末に日帰りで東京に遊びに行こうと思っています。日帰りなので、短時間で回れるおすすめの観光プランを教えてください。
AI: """

tokenizer = AutoTokenizer.from_pretrained(model_path)
input_ids = tokenizer.encode(prompt, return_tensors="pt")

推論実行についても、先程の 130 億パラメータモデルの例と同じです。700 億パラメータということで、より複雑なテキスト生成が実現できていることが分かります。

generated_sequences = neuron_model.sample(input_ids, sequence_length=512, top_k=50)

output = tokenizer.decode(generated_sequences[0], skip_special_tokens=True)
print(output)

推論結果の例

以下は人間とAIアシスタントとの会話です。

Human: こんにちは。
AI: こんにちは、私はAIアシスタントです。何かお手伝いできることはありますか？
Human: 週末に日帰りで東京に遊びに行こうと思っています。日帰りなので、短時間で回れるおすすめの観光プランを教えてください。
AI: 素晴らしいアイデアですね！東京には短時間で回れる多くの観光スポットがあります。たとえば、次のスポットをお勧めします：

皇居 - 東京の美しい庭園と巨大な城を眺めながら散策できます。
浅草寺 - 東京の最も有名な寺院の1つで、伝統的な日本の文化を体験できます。
東京タワー - 東京の街並みを見渡すことができるタワーで、東京タワーの周辺はショッピングや食事に最適です。
原宿&表参道 - ショッピング、ストリートアート、食べ物のための流行の発祥の地です。
明治神宮 - 東京の壮大な神社の1つで、静かで平和な雰囲気があります。
これらのスポットはいずれも都心周辺に位置し、短時間でアクセスできます。また、それぞれのスポットには、観光バス、電車、タクシー、レンタル自転車などを使用してアクセスできます。 

同じ Llama2 ベースで 700億パラメータを持つSwallow-70b-hfも同様に実行可能です。また、株式会社ELYZAでは日本語大規模言語モデルELYZA-japanese-Llama-2-70bを体験可能なデモサイトを、Amazon EC2 Inf2 インスタンスを用いて運用しています。

本番環境への導入に向けて

ここまで Amazon EC2 Inf2 インスタンス上で大規模言語モデルの推論（テキスト生成）を実行する手順について説明してきました。GPU インスタンス上での実行とさほど変わらない手順で実行可能なことがお分かり頂けたと思います。

Transformers NeuronX では、重みの 8 ビット量子化に対応しており、量子化する事で 700 億パラメータモデルをtp_degree=8で実行する事が可能です。この場合、inf2.48xlargeではモデルのサービングを 3 ワーカーで、trn1.32xlargeでは 4 ワーカーで運用可能となり、よりコスト性能を重視した運用が可能になります。また、性能向上を求める際の機能として speculative sampling にも対応しています。

ここでは、実際に本番環境でモデルのサービングを行う際に活用頂ける２つのサービスを紹介します。

Amazon SageMaker LMI コンテナ

Amazon SageMaker では、モデルの並列処理と大規模モデル推論用の Large Model Inference (LMI) Containers を提供しています。AWS Inferentia2、Trainium用の LMI コンテナでは、内部に Neuron SDK、Transformers NeuronX ライブラリ、DJLServing 高性能モデルサーバーを取り込み、SageMaker 上での本番環境での推論運用を実現しています。より詳細については
「大規模モデル推論コンテナを使って AWS Inferentia2 に大規模言語モデルをデプロイ」、「Amazon SageMaker 上で AWS Inferentia2 と AWS Trainium を使って、低コストで高性能な生成系 AI 推論を実現」をご参照下さい。

Hugging Face Text Generation Inference (TGI)

Hugging Face社では、大規模言語モデルの本番推論ワークロードに利用可能な Text Generation Inference (TGI) を提供しています。TGI には最新の Transformers NeuronX ライブラリを搭載し、Trainium、Inferentia2 の高い性能、低いコストを維持したまま Transformers ライブラリを簡単に利用するためのインタフェースである Optimum Neuron と共に、Neuron デバイスに最適化したデプロイ環境を提供しています。「Deploy Llama 2 70B on AWS Inferentia2 with Hugging Face Optimum」も併せてご参照下さい。

まとめ

本ブログでは Llama2 ベースの日本語モデルについて取り上げましたが、他の新しいモデルにも順次対応しています。
Transformers NeuronX では、Mistral-7B、MIxtral 8x7B にも対応しており、Swallow-MS-7b-v0.1、Swallow-MX-8x7b-NVE-v0.1、karakuri-lm-8x7b-chat-v0.1といった日本語モデルにも対応可能です。また、2024年 4月に発表された Llama3 8B、70B モデルにもいち早く対応し、Llama3 ベースの日本語モデルllama-3-youko-8bも Inferentia2 上で動作可能です。

モデルの対応状況については Neuron ドキュメントおよび AWS Neuron レポジトリをご参照ください。また AWS Trainium、Inferentia2 を活用しモデルの開発、運用を行っている企業の活用事例はこちらのリポジトリ上にまとめていますので合わせてご参照下さい。

本ブログが、AWS Trainium、Inferentia2 を活用して日本語大規模言語モデルを扱う際の道標となれば幸いです。

著者について

常世 大史 は、AWS Annapurna Labs のソリューションアーキテクトです。日本を拠点とし、AWS による買収以前から Annapurna Labs が提供する技術でお客様を支援してきました。ここ最近は、AWS Trainium、Inferentia の技術支援に注力しています。