AI21 Labs が Amazon SageMaker を利用して生成系 AI モデルの導入を加速

2017 年に設立された AI21 は、AI21 Studio を利用して独自の言語モデルへのアクセスを企業に提供しています。30,000 を超えるデベロッパーが AI21 Studio を利用して、独自の生成系 AI アプリケーションを構築しています。同社はまた、AI を活用した読み書きアシスタントである Wordtune も提供しています。Wordtune は、世界中の何千万ものユーザーが書き言葉に関するタスクを処理するのをサポートしています。

2021 年 8 月、AI21 は 2 つのサイズの Jurassic-1 言語モデルをリリースしました。Large モデルは 75 億のパラメータを備えており、高速でコスト効率が高く、Jumbo モデルは 1,780 億のパラメータを備えており、コストは高くなりますが、質の高いテキスト出力を提供します。大きい方のモデルは最高の質を提供しますが、大規模に実行するには高コストになる可能性があり、操作する際の機敏性にも劣ります。顧客が大規模に運用する際にコストと質の間のトレードオフを最適化できるよう、AI21 は Amazon SageMaker を利用して 170 億のパラメータを備えた 3 つ目のモデルである Grande を事前トレーニングし、2022 年 12 月にリリースしました。

AI21 は、開始から 2 か月未満でプロジェクトを迅速に完了しました。この際に、モデルの事前トレーニングにかかった時間は 20 日未満でした。LLM は数十億のパラメータを備えた巨大なニューラルネットワークであるため、トレーニングは時間のかかる困難なプロジェクトであり、膨大なコンピューティングリソースが必要です。Amazon SageMaker を利用することで、AI21 はよりシンプルかつ効率的なモデルトレーニングプロセスを経験し、必要な数の GPU にわたって分散トレーニングジョブをスケールすることができました。「AWS のソリューションアーキテクトは応答が迅速で、インタラクティブであったため、適宜トラブルシューティングしながら予定どおりにプロジェクトを完了できました」と AI21 の Vice President of Platform である Dan Padnos 氏は述べています。

既に AWS の利用経験があった同社は、高いコスト効率を備えており、使いやすく、フルマネージドであることを理由として、Amazon SageMaker を選択しました。AI21 はまた、既存のトレーニングソフトウェアスタックを引き続き使用して、迅速に稼働を開始することもできました。これは、同社がビジネスを構築する際に重要なことでした。Grande モデルの事前トレーニングを 20 日未満で完了するには、AI21 は 256 個の A100 GPU を 32 個のインスタンスに分散して使用する必要がありました。大規模なトレーニングを行うには、ノードの割り当てをオーケストレートし、中心的な場所でログ記録を利用できるようにして、手動による監視を軽減できるツールが必要でした。「この規模の分散トレーニングジョブを実行している場合、些細または平凡に見えるかもしれないあらゆる種類の技術的課題が頭痛の種になり得ます」と Padnos 氏は述べています。「Amazon SageMaker は、その複雑さを管理し、チームが細かい作業に費やす必要がある労力を減らすために使用できる機能を備えています」。 例えば、Amazon SageMaker は、企業が効率を高めるために使用できるヘルスチェックや一元的なログ記録などの機能を備えています。

Amazon SageMaker を利用することで、AI21 は新しいモデルを迅速にリリースできました。同社は、以前のトレーニング方法と比較して数週間の時間を節約できると想定しています。「Amazon SageMaker はノードの障害を処理し、スムーズに再起動して、大規模な分散実行をオーケストレートするため、モデルの事前トレーニングに取り組むチームはコアタスクに注力できます」と Padnos 氏は述べています。「技術的な課題に対処する代わりに、モデルがどのように実行されているか、トレーニングがどのように進行しているかを評価できるのです」。

Grande モデルの機能は AI21 の多くの顧客のニーズをより良く満たすものであったため、スケジュールの加速は重要でした。E メールの自動ドラフティングなどの消費者向けのユースケースに対応する必要がある顧客は、その規模の大きさによりにコスト効率が求められることから、Jumbo モデルから Grande モデルへの移行に価値を見出していました。Grande モデルは、その導入からわずか数か月後に、同社のトラフィックの約 3 分の 2 に使用されるようになりました。「急速に導入が進みました。当社はその結果に非常に満足しています」と Padnos 氏は述べています。「Amazon SageMaker を利用した私たちの経験は非常に好ましいものでした。予期せぬ困難が生じることもなく、予定どおり、予算内で、求めていた成果を実現できました」。

ユーザーエクスペリエンスはスムーズである必要があるため、生成系 AI アプリケーションの主な考慮事項は推論レイテンシーが低いことです。ユーザーは、Wordtune などのツールを利用してコンテンツをドラフティングする際、思考プロセスに遅延を生じさせることなく、AI がクイックリファレンスとして機能することを望んでいます。Amazon SageMaker を利用することにより、AI21 は顧客のニーズに最適な Grande モデルで低い推論レイテンシーを実現しました。これにより、ある大規模なクライアントは、レイテンシーを 4 分の 1 に低減することができました。その結果、AI21 の顧客は、ユーザーエクスペリエンスの質を落とすことなく、日々、ほぼリアルタイムで数百万のユーザーにサービスを提供できます。「当社のある大規模なクライアントでは、ユーザー満足度のメトリクスの大幅な改善が見られました。このクライアントは、Jumbo モデルから Grande モデルへの移行に伴ってレイテンシーが大幅に低減したことで、この改善を実現できたと考えています」と Padnos 氏は述べています。

また、Grande モデルのリリースは、AI21 とその顧客の両方の成長に貢献しました。「Amazon SageMaker を利用してトレーニングされた Grande モデルをリリースした後、全体的なトラフィックの増加が見られました」と Padnos 氏は述べています。「Grande モデルに移行した個別のクライアントのトラフィックも増加しました」。

Grande モデル (現在は Mid と呼ばれています) は、Amazon SageMaker JumpStart で利用できます。Amazon SageMaker JumpStart は、組み込みアルゴリズム、基盤モデル、事前構築済みの ML ソリューションを備えた ML ハブであり、Amazon SageMaker のユーザーは数回クリックするだけでデプロイできます。データライフサイクルはプライバシーを維持するためにユーザーの環境内に含まれており、組織はコードを記述したり、コードプレイグラウンドを必要としたりせずに、言語モデルをデータに適用できます。AI21 の次世代シリーズの基盤モデルである Jurassic-2 とタスク固有のモデルも、Amazon SageMaker JumpStart で利用できます。

AI21 は、今後数か月から数年で、世界中で生成系 AI の導入が増加していくことに高揚感を覚えています。同社は AWS サービスを利用して、より高速かつ正確で、高い信頼性とコスト効率を備えた LLM に積極的に取り組んでいます。「当社は AWS のチームと非常に良い関係を築いています」と Padnos 氏は述べています。「チームメンバーは当社と連携して技術的な詳細を深く掘り下げ、困難なタスクでも協力してくれました。プロセス全体を通じて、AWS のチームはクリエイティブに取り組んでおり、当社の課題と目標を認識していました」。

詳細については、https://aws.amazon.com/sagemaker にアクセスしてください。