🙌

ELYZA LLM for JP (デモ版)についての解説: (1) 70Bモデルの推論基盤

Takuro Murayama (Yamazaki)

2024/05/30に公開

はじめに

こんにちは。ELYZA のML Engineeringチームの堀江 (@eemon18)、村山 (@zakktakk)です。

本記事では、弊社が2024/03/11にリリースした下記のデモについて、どのように70Bという巨大なモデルをホスティングしているのかを解説します。

まだデモを直接触れたことがないという方はぜひ一度以下のURLからアクセスし、140GBを超えるバイナリファイルからなるモデルがどのくらい高速に動作するのかを確かめてみてください。

https://elyza.ai/lp/elyza-llm-for-jp

本記事ではまず弊社推論アーキテクチャについて説明し、その後70Bを運用する際の技術選定や高速化の工夫について解説します。

推論アーキテクチャ

弊社のLLMアプリケーションのアーキテクチャは、平易に表すと以下の図のように構成されています。LLMの推論処理には時間がかかるため、アプリケーションと推論インスタンスはメッセージキューを介して非同期に繋がっています。ブラウザ上またはAPI経由でリクエストを投げると、リクエストに紐づいたモデルの識別子に応じてAPIサーバーがメッセージの送信先を振り分け、適切な推論インスタンスにルーティングされる仕組みとなっています。

nsbridge

推論インスタンスは、メッセージキューからメッセージを取得し、推論を実行した後、結果をワーカー専用のキューに送信します。これにより、処理の負荷が均等に配分され、システムの効率性を高めます。

弊社では自社モデルの開発のほかに、各企業との共同研究を行なっています。個社専用の推論インスタンスには、共同研究を通じて得られたモデルやスクリプトを社内のAIエンジニアがプロダクション向けに実装し、デプロイしています。この時、メッセージキューとの通信、ファイル / オブジェクトストレージに配置されたモデルの読み込み、機密性の高いデータの暗号化・復号化など推論処理以外に実装が必要な部分が多くあります。それら全てをAIエンジニアが実装するのはコストが高く現実的ではありません。そこでアプリケーションとの通信に必要な処理を切り出し、nsbridge（nlp service bridgeの略、アプリケーションとの「橋渡し」をするという意図で命名）という社内ライブラリを提供することとしました。

以下の図に推論インスタンス内での処理フローとnsbridgeがカバーする範囲を示しています。AIエンジニアはnsbridgeの記法に従って推論処理を実装するだけで済み、共同研究の成果を迅速に本番環境に適用できます。

nsbridge_detail

今回作成したデモでは、生成したテキストを逐次的に出力するためにStreaming技術を利用しています。当初nsbridgeはStreamingに非対応でしたが、デモのために実装を追加し、アプリケーションと推論インスタンス間でのメッセージキューを介したStreaming機構を実現しました。このStreaming機構の詳細については、後日公開予定のテックブログで詳しく紹介することになります。

次章以降でこの推論アーキテクチャ上に、過去に運用経験のない70Bの巨大モデルを載せる際の工夫や苦労について話をしていきます。

70Bという巨大なモデルの運用の難しさ

一般的にMLモデルを運用する際に考慮しなくてはならない点として主に以下の3つがあると考えられます。

コスト：主に時間あたりの料金
レイテンシ： 1レスポンスを処理するのに要する時間
スループット： 1秒あたりに処理可能なレスポンス数

コストとレイテンシ・スループットはトレードオフの関係にあり、簡単に言ってしまえばH100を8台積んだようなインスタンスを用意してしまえば所望のレイテンシ・スループットを達成するのは可能かもしれませんが、$98.320 / hour (p5.48xlarge, us-east)という途方も無い運用費が要求されます。

今回運用を想定する70Bのような巨大モデルの場合、自社内に大量のGPUを確保しているという場合以外は、これらに加えてインスタンスを どれだけ容易に確保出来るか、という要素が入ってきます。
あらかじめ利用量を見込める場合はインスタンスを予約しておくといった方法もありますが、今回のようなデモとして公開するケースでは流量を事前に見積もり切るのは困難ですし、いざというときにスケールアップさせられないという事態を引き起こしかねない設計は問題があります。

また、LLM、というよりは巨大なモデル特有の話ではありますが、運用上トレードオフを考慮しなければいけない要素としてモデルの精度が含まれます。昨今、32 bitや16 bitサイズのモデルをそのまま運用時に用いることは少なく、量子化と呼ばれる手法によってバイナリサイズを小さくし、要求されるGPU RAMの量を少なくする試みが盛んに行われています。これによって、コストを下げつつもレイテンシやスループットを維持することが可能になるケースがあるのですが、多くの場合モデルの精度劣化とトレードオフの関係になります。

そのため、十分なレイテンシ・スループットを確保しながら、コストを一定に抑えつつ、必要なときにはN台以上までスケールアップでき、なおかつ学習時の精度を出来るだけ損なわない、という非常に困難な (それでいて技術的工夫のしがいがある) 課題が出来上がります。また、勿論ソフトウェアとして運用していく以上、ここには既存資産の使いまわしやすさやメンテナンスしやすさ、という要素も含まれます。

運用環境の選定基準

今回デモ用に70Bモデルを運用すると決めた際に、選定基準として、精度、コスト、確保容易性、レイテンシ、スループットの順に優先順位を付けることとしました。精度をもっとも優先順位を高くおいたのは、構築したモデルの日本語能力の高さをそのままデモを触ってもらうユーザーに届けたいと考えたためです。精度を多少犠牲に、より高速に、安価に構築する方法もあったのですが、精度劣化のほとんどない量子化方法について検証が十分に出来なかったため、今回は8 bit / 4 bit等の量子化については採用しないこととしました。
スループットを確保容易性より劣後させたのは、確保が十分に容易である場合スループットはスケールアウトで代替可能である可能性があるためです。

さて、量子化せずに70Bモデルを運用すると決めると、かなり選択肢は狭まります。140GB程のモデルを並列化せずに運用可能なGPUは市場に存在しないので、自然とT4 / L4 / A10 / A100 / H100を複数台搭載したインスタンスか、もしくはML専用のチップを搭載したインスタンスとなります。

単純に140GB以上のVRAMを確保可能なインスタンスをコストで並べた場合、以下のようになります (us-east region, US dollar / hour)

	AWS	GCP
A10G x 8 (g5.48xlarge)	16.288	-
L4 x 8 (g2-standard-96)	-	8.003
A100 x 8 (p4d.24xlarge, a2-highgpu-8g)	32.773	29.388
H100 x 8 (p5.48xlarge)	98.320	-
AWS inferentia2 accelerator x12 (inf2.48xlarge)	12.981	-

これらの中で、ある程度長期的に現実的に運用可能なラインとしてA10G, L4, そしてinf2インスタンスに絞り込みました。また、確保容易性の観点で見た場合に、A100, A10Gインスタンスは特に日本リージョンで確保が難しいことが多く、候補から外しました。

レイテンシの計測では、L4のマルチGPUによる推論に比べ、inf2の場合レイテンシが簡単に2倍程度短く出来そうなことがわかり、また、当初想定していたよりもLlama2がinf2インスタンス上で素直に動作したことが最終的な決め手となり今回はinf2.48xlargeでの運用となりました。

Speculative Decoding

今回のデモでは、モデルの精度を落とさずに高速化するための手法としてSpeculative Decoding を採用しています。技術的な詳細については割愛しますが、対象target modelに対し、それよりは比較的小さいdraft modelによってdraftを作成させ、target modelではそのtokenを並列でforwardします。こうすることで、target modelではautoregressiveな処理を回避し高速に推論することが可能になります。

4/10にリリースされた AWS Neuron SDK 2.18 でSpeculative Decodingのサポートが発表されていますが、実は実装自体は transformers-neuronx 内に存在し、コンパイル自体も可能になっていたため、AWSの方たちのサポートを受けながら本機能をデモに組み込みました。

この手法の肝はdraft modelに何を採用するか、という部分ですが、今回ベースとしては (VRAMの都合上) elyza/ELYZA-japanese-Llama-2-7bのモデルを採用し、候補としていくつかfine tuningを行い実際の推論速度、ELYZA-tasks-100による評価および、実際の出力の定性評価を通しチェックしました。

以下の推論速度の計測実験では、elyza/ELYZA-tasks-100を用いて100回推論したものを平均しています。

	推論速度	ELYZA-tasks-100 (GPT4-評価)
70B Model + Greedy Sampling (inf2)	22.75 +/- 0.66 tokens/s	3.84
1. Speculative Decoding with elyza/ELYZA-japanese-Llama-2-7b	35.69 +/- 6.76 tokens/s	3.81
2. Speculative Decoding with elyza/ELYZA-japanese-Llama-2-7b + target modelと同様の学習を行った draft model	36.09 +/- 6.04 tokens/ s	3.84
3. Speculative Decoding with elyza/ELYZA-japanese-Llama-2-7b + target modelの出力を学習した draft model	42.560 +/- 5.52 tokens/s	3.7

実験結果を見ると、target modelの出力を学習した、いわば蒸留モデルの推論速度は他と比べても非常に高速であり元の推論速度の2倍程度となっています。しかし、ELYZA-tasks-100による評価結果は芳しく無く、また、出力の人の目でのチェックにおいても他のモデルが2/100程度誤字脱字を含むのに対し、蒸留モデルは5/100となっていました。

Speculative Decodingの仕組み上、ここまで出力の質が低下することはあまり考えにくいのですが (問題があるtokenであれば棄却されるため) 、この問題はおそらくtarget modelによるtokenの採択確率推定のロバストさに起因していると考えられます。今後、temperatureを一定以上大きくしながらも、出力の質を低下させないロバストな推論方法を検討し、本問題は明らかにしていく予定です。

スポットインスタンスによるスケールアウト

推論処理の高速化だけでなく、システム全体のスループットを向上させるために、私たちはインスタンスのスケールアウトを実施しました。これにより、より多くのリクエストを同時に処理できるようになります。

アプリケーションが配置されている東京リージョンではinf2.48xlargeのオンデマンドインスタンスは1時間あたり19.47ドル（2024年5月17日時点）かかります。1台あたりのスループットが40[tokens/s]、1対話あたりの平均出力トークン数が300[tokens]とすると1つの対話を処理するのに7.5秒かかります。仮に秒間2対話の処理を目標とすると15台のインスタンスが必要となり、1日あたり100万円以上かかる計算となります。今回70Bを提供しているのはデモ用のサイトであり、多くのユーザーに弊社のモデルを体験してもらうことを目的としていますが、厳しい予算制約の中で運用する必要があり、1日100万円もの運用費は予算を大きく超過してしまいます。

そこで、コストと推論速度を両立するためにスポットインスタンスを利用することとしました。スポットインスタンスは、AWS EC2の余剰リソースの一部を安価に利用できる購入オプションの一つです。コスト的なメリットがある一方で、余剰リソースが存在しない場合は起動できない、または起動中でも余剰リソースが不足してくると中断されるため可用性の観点でデメリットがあります。そのため、デモサイトでは可用性との兼ね合いを考え、常時稼働のオンデマンドインスタンスを1〜2台立ち上げておき、それ以降はアクセス状況に応じて必要な台数を全てスポットインスタンスで賄うという形を取りました。これにより、コストを抑えつつ、アクセス量の増加にも柔軟に対応できる体制を整えています。

開発におけるハマりどころ

発展途上のライブラリ・SDK

NeuronCore関連のライブラリやSDKは発展途上で、まだバグや不足している機能が散見されます。実際に開発中に直面したバグとして、Speculative DecodingがECS Optimized AMI（amzn2-ami-ecs-inf-hvm-2.0.20240227-x86_64-ebs）でうまく動かなかった事例がありました。

弊社の推論コンテナ管理はECS on EC2で行っています。当初inf2向けのECS Optimized AMIを推論コンテナのベースに利用していたのですが、Speculative Decodingで推論を実行したときに生成速度が本来であれば40[tokens/s] 程度になることが見込まれているところ、4 [tokens/s] ほどしか出ない事象が発生しました。コンパイルや推論時のオプション・環境差分等を精査した結果、AMIを実験環境で用いていたInferentia向けのDeep Learning AMI（Deep Learning AMI Neuron PyTorch 1.13 (Amazon Linux 2) 20240214）に差し替えたところ実験時と同等の40 [tokens/s] のパフォーマンスが出ることが明らかとなり、AMIに起因する事象であることがわかりました。（本事象はAMIのアップデートにより修正されています。）

また、transformers_neuronxのSpeculative Decodingには繰り返しを抑制するパラメータであるrepetition_penaltyを設定することができません。このため、デモでは時折冗長な回答が生成されるという問題があります。こちらについては自前でくり返しを抑制する機構を実装することを検討しています。

環境差分による出力の変化

ナイーブに推論を実行すると70Bモデルの出力が社内の実験環境（GPU）と推論環境（inf2）で大きく異なるという問題も開発中に発生しました。以下に実際の入出力の差分を記載します。

入力

クマが海辺に行ってアザラシと友達になり、最終的には家に帰るというプロットの短編小説を書いてください。

実験環境（GPU）の出力

 クマのブルーはいつも森で一人でいました。彼は他のクマと違って、海が大好きでした。彼は海辺に行き、海を眺めて過ごすのが好きでした。

ある日、ブルーが海辺に行くと、アザラシのアーサーが一人でいました。アーサーもブルーと同じで、海辺によく来ていたのです。ブルーはアーサーに話しかけました。

「僕の名前はブルー。君は？」

「僕はアーサーだよ。ブルーはここによく来るの？」

「うん。僕は海が好きなんだ。アーサーは？」

「僕も海が好きなんだ。でも、僕はアザラシだから、当たり前かもしれないね。」

ブルーとアーサーは仲良くなり、一緒に遊ぶようになりました。彼らは海辺で一緒に過ごすことが多くなりました。

ある日、ブルーが森に帰ると、他のクマがブルーを呼び止めました。

「ブルー、どこに行ってたんだ？」

「海辺に行ってたよ。」

「海辺？一人で？」

「ううん。アザラシのアーサーと一緒に遊んでたんだ。」

「アザラシと？一緒に遊んでたの？」

「うん。アーサーは僕と一緒で、海が大好きなんだ。」

他のクマは驚きました。クマがアザラシと仲良くなるなんて、聞いたことがありませんでした。

ブルーは他のクマにアーサーとの出会いを話しました。他のクマはブルーの話を聞き、ブルーとアーサーの仲の良さがわかりました。

ブルーは他のクマにアーサーを紹介することにしました。ブルーとアーサーは他のクマと一緒に遊ぶようになり、みんな仲良くなりました。

ブルーはアーサーと出会ったことで、他のクマとも仲良くなることができました。ブルーはアーサーに感謝し、これからもアーサーと仲良くしていこうと思いました。

推論環境（inf2）の出力

クマのブルーはいつも森で一人でいました。彼は他のクマと違って、海が大好きでした。彼は海辺に行き、海を眺め、独りで過ごすのが好きでした。

ある日、ブルーが海辺に行くと、アザラシのアーサーが一人でいました。アーサーもブルーと同じで、海辺に来るのが好きだったのです。ブルーはアーサーに話しかけました。

「僕はブルー。君は？」

「僕はアーサーだよ。クマとアザラシが友達なんて、初めてだね。」

「僕もだよ。でも、僕は他のクマともあまり遊ばないんだ。」

「僕もだよ。僕は海辺に来るのが好きなんだ。」

ブルーとアーサーは仲良くなり、一緒に遊ぶようになりました。彼らは海辺で一緒に遊んだり、お互いの家に遊びに行ったりしました。

ある日、ブルーは森の家に帰ることにしました。彼はアーサーに別れを告げ、森に帰っていきました。

ブルーが森に帰ると、他のクマが遊びに誘ってくれました。ブルーは嬉しくなり、一緒に遊むことにしました。ブルーは他のクマとも仲良くなり、楽しく過ごしました。

ブルーはアーサーに会いたくなり、また海辺に行くことにしました。彼はアーサーに会い、一緒に遊きました。ブルーとアーサーは仲良しの友達として、これからも一緒に過ごすのでした。

こちらもライブラリ、インフラ、コンパイル時のオプション等差分が起こりうる要素が複数あり、原因調査に多くの時間を割きましたが、結果的に原因はLlama-2-70Bとtransformers_neuronxを併用した場合特有の問題（release noteにおいてもKnown Issuesとして報告されています）であることがわかり、コンパイル時のオプションを正しく指定する（--enable-mixed-precision-accumulation を追加する）ことで解決しました。