AI21 Labs が Amazon EC2 P4d インスタンス、PyTorch を使用して、1,780 億のパラメータを備えた言語モデルをトレーニング

AI21 Labs は、意味の理解に重点を置いた言語モデルを開発するために機械学習を利用しており、2021 年には、最近リリースされた Jurassic-1 Jumbo をトレーニングするという目標を設定しました。これは、1,780 億のパラメータを持つ自己回帰言語モデルです。ベータテストに登録したデベロッパーは Jurassic-1 Jumbo にアクセスできるようになり、ユースケースに合わせてモデルのカスタマイズをすぐに開始できます。このソフトウェアスタートアップは、モデルを効率的にトレーニングしたいと考えていたため、Amazon Web Services (AWS) に目を向け、Amazon Elastic Compute Cloud (Amazon EC2) を利用したソリューションを構築しました。Amazon EC2 は、クラウドにおいて安全でサイズ変更可能なコンピューティングキャパシティを提供するウェブサービスです。Amazon EC2 を選択したことで、同社はノードの割り当てを含むトレーニングプロセスを制御できるようになりました。

強力なコンピューティング機能とネットワーキング機能を実現するために、同社は Amazon EC2 P4d インスタンスを選択しました。これは、クラウドでの機械学習トレーニングとハイパフォーマンスコンピューティングアプリケーション向けに、高スループットかつ低レイテンシーのネットワーキングを提供します。AI21 Labs は Amazon EC2 P4d インスタンスを利用して、Jurassic-1 Jumbo モデルを通じてサービスとしての自然言語処理を提供するために、数百の GPU にモデルトレーニングを分散することで、必要なパフォーマンスとメモリを得ることができました。同社は現在、自社の大規模モデルをトレーニングおよび管理しているため、同規模で新しいモデルの開発に取り組み、より簡単にイノベーションを実現できます。

2017 年に設立された AI21 Labs は、自然言語処理の研究を実施し、読み書き用の人工知能を活用した製品を開発するというハイブリッドなミッションの実現に取り組んでいます。同社の主力製品である Wordtune は、2020 年 10 月にリリースされたインテリジェントな執筆および編集アシスタントであり、100 万近くのユーザーをサポートするまでに成長しました。もう 1 つの主な製品である AI21 Studio は、同社の Jurassic-1 言語モデルへの API アクセスとカスタムモデル開発を提供します。「当社は、言語モデルをサービスとして提供し、独立したデベロッパーから多国籍企業に至るまで、誰もが高度な自然言語処理テクノロジーを基盤としたアプリやサービスを構築できるようにしている、数少ない企業のうちの 1 社です」と AI21 Labs の共同創業者であり、共同 CEO でもある Yoav Shoham 氏は述べています。「さらに、当社は科学的イノベーションを追及しており、この規模と複雑さのモデルで生じるソフトウェアエンジニアリングの課題にも取り組んでいます」。

最初の深層学習メガモデルを効率的にトレーニングし、モデルの高いスケーリングとパフォーマンスのニーズをサポートするために、AI21 Labs は、強力なコンピューティング、効率的なネットワーキング速度、技術サポートとガイダンスを活用できるようにしたいと考えていました。これらの理由により、同社は 2021 年初めに AWS 上でソリューションの実装を開始し、Amazon EC2 P4d インスタンスを利用してモデルをトレーニングすることにしました。これらのインスタンスは、Amazon EC2 UltraClusters と呼ばれるハイパースケールクラスターにデプロイされ、4,000 を超える NVIDIA A100 GPU、ペタビットスケールのノンブロッキングネットワーキングインフラストラクチャ、および高スループットかつ低レイテンシーのストレージを提供します。 

同社のアプローチは、低レイテンシーかつ高帯域幅の GPUDirectRDMA と、Amazon EC2 インスタンスのネットワークインターフェイスである Elastic Fabric Adapter (EFA) によってさらに最適化されました。EFA を使用することで、お客様は高レベルのノード間通信を必要とするアプリケーションを AWS 上で大規模に実行できます。モデルのサイズに鑑みて、チームはトレーニング時間を効率化するために並列処理を利用する必要があり、分散トレーニングとモデルの並列処理をサポートするために AWS 上のネットワーキング機能に注目しました。「Amazon EC2 P4d インスタンスは、EFA 上で 400 Gbps のハイパフォーマンスネットワーキングを提供します」と AI21 Labs の Jurassic テクニカルリードである Opher Lieber 氏は述べています。「GPU 間のネットワーク速度は、数百の GPU へのスケーリングにおいて、コスト効率を維持しながら効率的にスケールする能力に直接影響します」。

AI21 Labs は、EFA 用にアクティブ化された Amazon EC2 P4d インスタンスでコードベースを起動することから始めました。その後、マルチノードトレーニングアプローチのパフォーマンスと効率的なスケーリングをテストおよび検証しました。次にチームは、機能とパフォーマンスを検証するために、数百の GPU を利用するフルサイズモデルの簡単なトレーニングを開始しました。そこから、AWS 上で Jurassic-1 Jumbo モデルのトレーニングを開始することができました。オーケストレーションのために、同社は AWS ソフトウェア開発キット (AWS SDK for Python (Boto3)) を使用してインスタンスを割り当てる社内ソリューションを選択しました。AWS SDK for Python (Boto3) を利用することで、お客様の Python アプリケーション、ライブラリ、またはスクリプトをさまざまな AWS のサービスと簡単に統合できます。

ストレージのために、AI21 Labs は、Amazon Simple Storage Service (Amazon S3) を選択しました。Amazon S3 は、業界をリードするスケーラビリティ、データ可用性、セキュリティ、およびパフォーマンスを提供します。「AWS チームのサポートにより、Amazon S3 で非常に優れたパフォーマンスを実現できました。パフォーマンスと料金の両方を考慮すると、Amazon S3 を選択するのは当然でした」と Lieber 氏は述べています。チームは Amazon S3 バケットを利用して、効率的かつ分散された方法でチェックポイントを保存およびロードします。トレーニングの進行状況とイベントをログ記録するために、チームはモニタリングおよびオブザーバビリティサービスである Amazon CloudWatch を利用しています。 

AI21 Labs は、ソリューションを実装する際に AWS のサポートを活用しました。同社のチームが相談した AWS のスペシャリストは、サービスレベル、アーキテクチャ、およびハードウェアに関連する質問や懸念事項についてのガイダンスを提供しました。さらに、同社は AWS での PyTorch を利用して Jurassic-1 Jumbo のパフォーマンスを改善しました。AWS での PyTorch は、機械学習モデルの開発と本番環境へのデプロイを容易にするオープンソースの深層学習フレームワークです。 

AI21 Labs は数か月にわたるトレーニングを 2021 年 6 月に完了しました。新しいメガモデルである自己回帰言語モデルには 1,780 億のパラメータが備わっており、これは同社の競合他社の製品に匹敵する数です。また、拡張されたテキスト表現機能と名前付きエンティティのサポートを提供する、差別化された 256,000 項目の語彙も提供します。同社は現在、同社の AI21 Studio 製品を通じて、Jurassic-1 Jumbo (および 70 億のパラメータを備えた同等のモデルである Jurassic-1 Large) をオープンベータ版で提供しています。このサービスを利用すると、幅広いデベロッパーが Jurassic-1 Jumbo モデルで製品を構築でき、AI21 Labs は既に、マーケティング、コンテンツ作成、ゲーム、医療研究、自動車、電気通信、金融などの多くの業界にこのサービスを提供しています。

AI21 Labs はそのモデルを所有し、直接アクセスできるため、サードパーティーに依拠することなく変更を加えたり、イノベーションを起こしたりできます。これにより、同社のミッションの重要な部分である継続的なイノベーションの目標を追及できます。AI21 Labs は現在、追加モデルのプロトタイプを作成しており、これを大規模にトレーニングすることも計画しています。「独自のメガモデルをトレーニングし、所有していることは、今後も Wordtune と AI21 Studio の両方の製品において重要な差別化要因であり続けることでしょう」と Shoham 氏は述べています。