re:Invent 2023: AWSとToyotaが語るSageMakerでの大規模ML学習
はじめに
海外の様々な講演を日本語記事に書き起こすことで、隠れた良質な情報をもっと身近なものに。そんなコンセプトで進める本企画で今回取り上げるプレゼンテーションはこちら！
📖 AWS re:Invent 2023 - Train and tune state-of-the-art ML models on Amazon SageMaker (AIM335)
この動画では、AWSのSageMakerを使った最先端の機械学習モデルのトレーニングとチューニングについて学べます。Amazon SageMakerのPrincipal Product ManagerであるGal Oshriが、大規模モデルのトレーニングにおける課題とその解決策を紹介します。さらに、Toyota Research InstituteのThomas Kollarが、SageMakerを活用した自動運転やロボティクスの研究事例を共有します。Llama 2の事前学習や日本語モデルの追加学習など、具体的な実装例も見ることができます。 ※ 動画から自動生成した記事になります。誤字脱字や誤った内容が記載される可能性がありますので、正確な情報は動画本編をご覧ください。
本編
Amazon SageMakerを用いた大規模機械学習モデルのトレーニングとチューニング
みなさん、こんにちは。私はGal Oshriと申します。AWSのSageMakerを担当するPrincipal Product Managerです。本日は、Emily WebberさんとToyota Research InstituteのMachine Learning ManagerであるThomas Kollarさんと一緒に、Amazon SageMakerを使用した最先端の機械学習モデルのトレーニングとチューニングについてお話しします。始める前に、現在すでに機械学習モデルのトレーニングを行っている方は何人いらっしゃいますか？素晴らしいですね。10個以上のGPUやアクセラレーターを使ってトレーニングしている方は？そうですか、100個以上の方は？1000個以上は？いないようですね。今日は、そういったことについて少し学んでいきましょう。
大規模な機械学習モデルのトレーニングにおける課題と、SageMakerがどのようにそれらのモデルのトレーニングを支援できるかについてお話しします。その後、Emilyが大規模言語モデルのファインチューニングとプリトレーニングについて説明し、SageMaker上でLlama 7Bをトレーニングするデモをお見せします。最後に、TomからToyota Research Instituteと彼らの機械学習ユースケースについてお話を伺います。
機械学習モデルの進化と大規模化の課題
機械学習は、レコメンデーションから信用リスク予測、自動運転からドキュメント分析まで、幅広い分野ですでにその有用性が証明されています。しかし最近、コンピュータビジョンと自然言語処理のディープラーニングモデルへの関心が爆発的に高まっています。 ほんの数年前まで、機械学習モデルを使って非常にきれいなリビングルームを生成しようとすると、このような画像が得られていました。偽物だとわかりますし、一貫性がなく、リビングルームだとはっきりとはわかりません。
わずか数年後の今、このような画像を生成できるようになりました。本物の画像ではないと見分けるには、かなり注意深く見る必要があります。re:Inventで1年前に同様のものをお見せしましたが、その時はこのような画像がモデルから生成されるのを見て、ある意味衝撃的でした。しかし1年後の今、会場の多くの方々がすでにこのタイプの画像が生成されるのを見たことがあり、機械学習モデルで得られる品質を知っています。
では、これはどのように実現されたのでしょうか？まず、ここ数年でアルゴリズムに顕著な改善がありました。特に、今日耳にする多くの大規模モデルで使用されているtransformerアーキテクチャです。しかし、これらのモデルをトレーニングするためのデータセット、モデルサイズ、計算量も増加しています。多くの研究が示すように、これらの要素を拡大し続けることで、より良い結果を得ることができます。そのため、競争力を維持するには、これらの進歩を活用して機械学習で顧客に最高の体験を提供する方法を真剣に考える必要があります。
大規模モデルトレーニングにおける課題とSageMakerの役割
さて、大規模モデルのトレーニングは素晴らしいものです。すぐに最大のモデルをトレーニングして終わりにしてしまおう、と思うかもしれません。しかし、実際にはそれほど単純ではありません。いくつかの課題があります。 まず、最新のハードウェアを使用したいということです。数年ごとにハードウェアの革新が起こり、トレーニングの効率が2倍から9倍に向上します。しかし、最新のハードウェアにアクセスできるだけでは十分ではありません。それがどれだけうまく機能するかを考える必要があります。 machine learningチームへの中断を最小限に抑えてトレーニングを継続できるほど、障害に強いものでしょうか？オーケストレーションについて、そして利用可能なリソースを最も効果的に使用する方法について考える必要があります。
特に、多くのモデルを並行してトレーニングしたい大規模なdata scientistチームがある場合はそうです。より大規模なデータセットが必要だという話をしましたが、 これらの大規模データセットの保存、読み込み、処理には多くの作業が必要になる可能性があります。そこには多くの落とし穴があります。モデルのトレーニングのためにより多くのコンピューティングリソースを得るためのインフラストラクチャと、使用するアルゴリズムの両方をスケールアップすることを考える必要があります。今日我々がこれらのユースケースのためにトレーニングするモデルは、単一のアクセラレータには収まりません。そのため、スケールアップに必要なアルゴリズムについて考える必要があります。
最後に、コストについて考える必要があります。これらのモデルのトレーニングには、数十万ドルから数百万ドルのコストがかかる可能性があります。そのため、特に散発的な実験を行い、さまざまなアイデアを試している初期段階では、モデルのトレーニング時の効率性について考える必要があります。効率的に使用する方法を考える必要があります。そして、それは金銭的なコストだけでなく、チームの時間も同様です。多くの顧客が、ML engineerがインフラストラクチャの問題に時間を費やさないようにすることが最優先事項の1つだと私たちに伝えています。
SageMakerのトレーニングジョブAPIと機能
しかし、希望がないわけではありません。Amazon SageMakerは、これらの多くの課題に対処するのに役立ちます。SageMakerがどのように機能するかの概要を説明しますが、 Emilyのデモでより詳細に見ることができます。まず、create training job APIを呼び出します。このAPIは、データセット、コンピューティングリソースの構成、および使用するトレーニングアルゴリズムに関する情報を取得します。SageMakerは、モデルのトレーニングのためのクラスターを設定し、 デフォルトで適切なVPCとネットワーク構成を行い、多くの時間を節約します。ただし、必要に応じてすべてを自分で構成し、柔軟性を追加することもできます。
デフォルトでは、Amazon SageMakerは多くの時間を節約するように設計されています。クラスターのスピンアップの一環として、SageMakerはジョブが開始される前、そして課金が始まる前に、ハードウェアのヘルスチェックを実行して、すべてが効果的に機能していることを確認します。これにより、時間とコストを節約し、トレーニングが効率的に継続できることを保証します。
SageMakerは、Amazon S3、Amazon EFS、またはAmazon FSx for Lustreからデータを読み込み、データをコピーするかストリーミングするかのオプションがあります。データセットのサイズによって、どちらかがより適切になるでしょう。SageMakerは素早く始められる優れたオプションを提供していますが、同時に柔軟性も備えており、他のソースからデータを読み込むこともできます。その後、Amazon ECRからトレーニングイメージをダウンロードし、SageMakerの組み込みアルゴリズムを選択できます。SageMakerのディープラーニングコンテナを使用してPyTorch、TensorFlow、Hugging Faceをすぐに利用することもできますし、独自のアルゴリズムを含む独自のトレーニングイメージを持ち込むこともできます。
SageMakerはまた、データ並列とモデル並列の両方のトレーニングを加速させる分散トレーニングライブラリも提供しています。SageMakerはトレーニングを開始し、プロセス全体を通じてログをAmazon CloudWatchにストリーミングします。メタデータとハイパーパラメータを保存し、後で確認できるようにします。TensorBoardなどのツールを使用して実験を可視化するオプションもあります。トレーニング中、チェックポイントをストレージに同期させます。これは障害に強くなるために重要です。トレーニング中に何か問題が発生した場合、その時点までの進捗を失わないようにするためです。
トレーニングの終了時、SageMakerはモデルと他の出力データを保存し、後で再確認できるようにします。その後、SageMakerはすべてのコンピュートをシャットダウンします。つまり、ジョブが深夜3時に失敗しても、誰かが起きて電源を切る必要はなく、モデルのトレーニングに使用されていないハードウェアやインスタンスに対して料金を支払い続けることもありません。同じパラダイムで、トレーニングを簡単に多くのインスタンスにスケールアップすることができ、大規模なモデルを得ることができます。
SageMakerの回復力と使いやすさ
今年リリースされた素晴らしい機能の1つが、クラスタ修復機能です。トレーニング中にインスタンスの1つが失敗した場合、そのインスタンスに何が起こったかを確認し、再起動が必要か、別のインスタンスに置き換える必要があるかを判断し、数分以内にトレーニングを再開します。これらの回復力機能により、トレーニングが可能な限り迅速に、人手を介さずに継続されることが保証されます。
もしこれらが少し難しく感じられたとしても、実際に始めるのは非常に簡単です。モデルトレーニングをSageMakerトレーニングジョブに変換するための最も重要なコードは、estimator APIです。後ほどデモで詳しく見ていきますが、大まかに言えば、このAPIはPythonファイルまたはエントリーポイント（この場合はcifar10.py）を受け取ります。これは、ラップトップでモデルトレーニングを行う方法とよく似ています。また、使用したいインスタンスタイプ、その数、そしてハイパーパラメータ（後で簡単に変更して追加のトレーニングジョブを試すことができます）も指定します。さらに、トレーニング中にCloudWatchでそれらのメトリクスを表示できるように、メトリクス定義を追加します。最後に、データへのパスを提供し、estimator.fitを呼び出します。
さらに良いニュースがあります。最近、始めるのがさらに簡単になりました。既存のPythonコードに remote python デコレータを追加するだけで、ランタイム、パッケージ、関数、その他すべてを自動的にシリアライズし、estimator APIについて学ぶ必要もなく、SageMaker training jobとして実行できるようになりました。トレーニングジョブが開始されると、メタデータを簡単に表示し、後で再現したり、そのトレーニングジョブをクローンしたりすることができます。
実験の追跡は非常に重要です。過去の経験から学ぶことが大切です。トレーニング結果をスプレッドシートや文書に記録してチーム内で共有している人をよく見かけますが、それでは協力して過去の実験から学ぶのが難しくなります。これらすべてを自動的に一箇所に保管することで、失敗から学び、より良いモデルを構築することがはるかに容易になります。
SageMaker Profilerとスマートシフティング機能
では、トレーニングの追跡から、パフォーマンス、特にトレーニング速度の向上に話を移しましょう。これは、モデルのトレーニングにインスタンスを使用する時間、プロジェクト全体の完了時間、そしてコストに影響します。SageMaker Profilerは、ハードウェアの使用状況を理解し、パフォーマンスの問題の根本原因を特定するためのML観測ツールです。
モデルトレーニングの効率を最大化するために、ここに示すダッシュボードには、GPU使用率に関する全体的な指標が表示されます。これはできるだけ高くしたいものです。また、クラスター内の各ノードにおけるトレーニングジョブ全体を通じてのGPU使用率も表示されます。つまり、全体的な使用率が高くても、一部の間隔で使用率が低い場合があり、それをもう少し詳しく調べたいかもしれません。ダッシュボードの下部には、他の指標もあります。例えば、各GPUカーネルにかかった合計時間などです。これにより、トレーニングをさらに改善するために次に何を最適化すべきかについての追加のヒントが得られます。
Profilerには、より詳細なタイムラインビューを示す別のページがあり、ホストとデバイスのすべてのレベルからデータを取得できるので、各ポイントで何が起こっているかをより深く理解できます。ここで、Amazon SageMakerの新機能のプレビューをお知らせできることを嬉しく思います。それは、データのスマートシフティングです。スマートシフティングは、オンラインデータ精製技術で、ディープラーニングのトレーニング時間とコストを最大35%削減できます。モデルをトレーニングする際、より大きなデータセットを使用したいと話しましたが、データセットの質も重要です。データ内の一部のサンプルは、モデルトレーニングにとってあまり有益でない場合があります。あるいは、すでに見たサンプルかもしれません。重複データや類似データがある可能性もあります。
データの前処理や不要なデータの削除は、しばしば困難です。そこで、スマートシフティングが役立ちます。これは、トレーニングジョブ中にデータを分析し、モデルにとってあまり有益でない低損失サンプルをフィルタリングします。データのサブセットでトレーニングすることで、トレーニングにかかる時間とコストを最大35%削減できます。また、低損失サンプルのみをフィルタリングするため、最終的なトレーニング精度への影響は最小限か、ほとんどありません。さらに、データやトレーニングパイプラインを変更する必要がないため、簡単に始められます。
こちらは、スマートシフティングを使用した簡単な例です。Amazon SageMaker deep learning containerを使用し、シフティングデータローダーをロードします。そして、既存のデータローダーをシフティングデータローダーでラップし、少し設定を追加するだけで使用を開始できます。モデルやデータパイプラインの他の部分を変更する必要はありません。すでに、この機能で素晴らしい結果を得ているお客様もいます。例えば、LG AI Researchは、モデルの精度を変えることなく、トレーニングのパフォーマンスを大幅に向上させることができました。
Amazon SageMaker HyperPodとファインチューニングの利点
re:Inventでは、Amazon SageMaker HyperPodも発表しました。これにより、大規模モデルをトレーニングするお客様は、今日お話ししたSageMakerの管理された利点をすべて得られると同時に、インスタンスに直接アクセスしたり、Slurmを使用したりするなど、より馴染みのあるUXを利用できます。トレーニングについて説明した、故障したインスタンスの交換やトレーニングの迅速な開始など、同様の耐障害性機能を備えており、時間を最大20%節約できます。また、モデル並列トレーニングとデータ並列トレーニングの両方でパフォーマンスを向上させる、最適化された分散トレーニングライブラリの恩恵も受けられます。
クラスターに対するより細かな制御が可能で、インスタンスに直接アクセスしたり、追加のソフトウェアをインストールしたり、クラスターに必要な変更を加えたりして、トレーニングをさらに微調整できると言いました。さて、大規模モデルのトレーニングについて話してきましたが、それが常に必要というわけではありません。時には、既存のモデルをファインチューニングするだけで十分な場合もあります。これは、既存のモデル、つまりファウンデーションモデルがあり、それを特定のユースケースに合わせてファインチューニングしたい場合に有効です。自分のデータを導入することで、既製のファウンデーションモデルを使用するよりも、モデルを改善できます。しかし、モデル全体を一から学習させる必要がないため、多くの時間とコストを節約できます。
ただし、課題もあります。一部のモデルはオープンソースではないため、モデルの重みをダウンロードして、既存のSageMakerトレーニングジョブで自分でファインチューニングすることができません。しかし、SageMakerのアルゴリズムとモデルパッケージに加えられた改良により、この状況は変わりました。今では、サードパーティのモデルを、プライベートデータでファインチューニングすることで、簡単かつ安全にカスタマイズできるようになりました。これにより、エンドツーエンドのセキュリティが提供されます。モデルプロバイダーは、モデルの重みを公開することなくモデルを提供でき、お客様は、データをモデルプロバイダーに公開することなく、自分のデータを使ってそのモデルをファインチューニングできます。さらに、ファインチューニング後の最終的なモデルの重みも、お客様だけがアクセスできます。
LLMのファインチューニングとプリトレーニング：Emily Webberの解説
これは様々なモデルやアルゴリズムで実現できます。例えば、Cohereモデルなどがあります。これらはすべて、先ほど説明したPython SageMaker SDKを通じて簡単に使用でき、SageMaker experimentsやpipelinesなど、他のSageMaker機能とも統合されています。SageMaker inferenceを使用すると、最終的にモデルをデプロイして、本番シナリオで安全な方法で推論に使用することができます。
ここで、SageMakerでのLLMのファインチューニングとプリトレーニングについて、Emilyに説明を譲りたいと思います。ありがとう、Gal。皆さんも私と同じくらい、これらの新しいローンチと機能に興奮していることを願っています。自己紹介させていただきます。私はEmily Webberで、AWSのgenerative AI foundationsテクニカルフィールドコミュニティをリードしています。特に、これらのローンチの多くは、皆さんとの直接の対話から生まれました。私たちはお客様の声に耳を傾け、トレーニングスタックに求められる主要な機能を理解するために皆さんと対話を重ね、その結果、今ご紹介した多くの機能が生まれました。
大規模言語モデルをカスタマイズする方法は多岐にわたります。ここでは、2つの軸で表現しています。下の軸は大まかに複雑さとコストを表しています。当然、左側に近いほど望ましいです。LLMのカスタマイズ技術は比較的簡単であることが望ましく、それによって素早く始められ、コストも抑えられます。しかし、これらの技術には段階があります。ほとんどのお客様は、プロンプトエンジニアリングから始めます。これは大規模言語モデルを簡単に改善・カスタマイズする良い方法です。ただし、他の追加技術ほど正確ではありません。
多くのお客様は、プロンプトエンジニアリングから、retrieval augmented generationスタックと呼ばれるものに移行します。ここでは、データセットがあり、そのデータを埋め込み（dense representation）に変換し、それらのドキュメントを取得してユーザーとやり取りします。これは次にファインチューニングスタックに変わる可能性があります。ここには少し重複がありますが、Galが言及したように、カスタムデータを使用してモデルをファインチューニングし、追加の知識を加えることができます。しかし、これらの技術はすべて、究極の目標であるプリトレーニング、つまり新しい基盤モデルの作成と比べると見劣りします。これらの技術はすべてSageMakerで利用可能であり、スタックによってしっかりサポートされています。今日のセッションでは、ファインチューニングからプリトレーニングへの移行方法を学びます。
さて、小規模モデルのファインチューニングは非常に効果的です。小規模モデルをファインチューニングする理由をいくつか挙げてみましょう。まず、コストが低いです。より小さなデータセット、場合によってはより小さなモデルを使用しながら、ファインチューニングによって精度を向上させつつ、コストを抑えることができます。7億パラメータ程度の小規模モデルを扱う場合、モデル自体が物理的に小さいため、より大きなモデルと比べてトレーニング時間が短く、推論時間も速くなります。つまり、より多くのモデルをトレーニングし、より多くの推論を行うことができます。
オブジェクトが小さくなるため、管理が容易になります。ストレージ要件が小さくなるので、モデルのコピーが簡単になります。アプリケーション、パッケージ、CICDパイプライン、リポジトリにモデルを組み込むのも容易になります。多くの顧客は、特にfine-tuningやpre-trainingを通じて新しいモデルを作成することで得られる所有権を本質的に好みます。これにより、企業の知的財産を増やすことができます。また、小さなモデルをfine-tuningする場合、デプロイメントのオプションも増えます。デプロイメントのオプションには、serverlessも含まれます。私の顧客の中には、これらの小さな70億パラメータのモデルを作成してfine-tuningし、コンパイルしてLambdaでホストし、serverlessで推論を実行している人もいます。小さな領域に精通したこれらの小さなモデルを扱う場合、非常に柔軟性があります。
SageMakerを用いたLlama 2のプリトレーニングデモ
Pre-trainingは、非常に大規模なデータセットに最適です。オンラインで簡単に入手できない数百GBや複数テラバイトのカスタム言語データがある場合、言語モデルを真剣にカスタマイズすることを検討する必要があるかもしれません。場合によっては、ゼロから新しいモデルを作成することも考えられます。
これは特に、使用している主要な言語がWikipediaやRedditのようなソースでうまく表現されていない場合や、これらのプラットフォームが特定の分野の専門用語やコンテキストを理解していない場合に関連します。
では、なぜpre-trainingがそれほど強力なのでしょうか？その理由の一つは、pre-trainingの損失関数がより汎用性が高いからです。ゼロから新しいfoundation modelを作成する場合、学習は少し異なります。より一般的で、ニューラルネットワークの深い部分で行われます。さらに、pre-trainingは教師なしデータで行うことができるため、何百万もの記録にラベルを付ける必要がありません。単に自分の言語データのテラバイトを取得してトークン化し、追加の監督なしでそれをネットワークに入力することができます。
また、pre-trainingの効率性の向上も評価しています。多くの人が、データサイエンスやモデリングのためのリソースが限られた小規模なチームを持っています。数多くのアプリケーションで使用できる1つの巨大で強力なfoundation modelの作成に努力を集中することで、transformerが登場する前の一般的な慣行であった複数の小規模な機械学習オペレーションを最適化し維持するよりも効率的であることがよくあります。
AWSでのファウンデーションモデルの事前学習プロセス
新しいfoundation modelを事前学習するには何が必要でしょうか？一見難しそうに思えるかもしれませんが、洗練された成熟した学習インフラのおかげで、実際にはかなり手の届くものになっています。Amazon SageMakerで、特に私たちの学習インフラ上でゼロから事前学習され作成されたモデルの例を3つご紹介しましょう。画像モデルのStable Diffusionは、50億枚の画像と240テラバイトの画像データで学習されました。Technology Innovation Instituteが開発した最大のオープンソース言語モデルであるFalconは、3テラバイト弱の言語データから1兆のトークンを使用しました。このモデルは400億のパラメータを持ち、48台のp4dインスタンスを使用して2ヶ月かかって学習されました。最後に、SageMaker上で学習された金融分野の大規模言語モデルがあり、これは2テラバイト弱の言語データを使用しています。
これらの要件は意外にも手の届くものです。多くの企業がおそらくそれだけの量の言語データを持っており、SageMakerで提供している機能により、幅広い組織にとって学習体験が非常に身近なものになっています。
では、AWSでこれらのfoundation modelを作成し事前学習するにはどうすればよいでしょうか？最初のステップはデータの収集とアクセスです。少なくとも1テラバイトの自社の言語データを用意することをお勧めします。これには文書、デジタル化されたPDF、会話、その他の豊富で堅牢な言語ストリームが含まれます。多くの企業がこれをオープンソースのデータと組み合わせることで、モデルが自社特有の言語のニュアンスと一般的な知識の両方を理解できるようにしています。
情報を収集して保存した後は、データの前処理を行います。SageMakerはデータセットの前処理にも優れた機能を備えています。
SageMakerを活用したToyota Research Institute (TRI)の機械学習の取り組み
私たちのビルダーの一人であるJennyが、Amazon SageMakerで多くの前処理とデータ変換ジョブを実行するのを手伝ってくれました。先ほど学んだリモート関数を含むトレーニングジョブAPIを使用して、トークン化や前処理のジョブを並列で実行できます。このコアとなるトレーニングジョブの構成は、新しいモデルをゼロから作成する場合と、一般的なデータ変換や処理の両方に適用できます。
データセットを前処理した後、お好みのデータストレージを使用して最適化します。多くのお客様がAmazon FSx for Lustreを使用しているのを見かけます。これは、データを1か所に保存し、このボリュームをトレーニングジョブの実行に簡単に接続できるからです。異なるモデルサイズ、インフラストラクチャ、実験的な選択肢を繰り返し試す際に、同じ場所でデータを使用し保存できます。
その後、お客様はトレーニングスクリプトの開発と反復を行う必要があります。SageMakerのインフラストラクチャで得られる柔軟性は素晴らしいものです。T3 mediumのような小さなインスタンスを使用できます。これはノートブックインスタンスに最適で、非常にコスト効率が良く小規模です。そして、ボタンを数回クリックするだけで、G4やG5シリーズなどの小型GPUにスケールアップできます。チームはこれを使って開発し、トレーニングループの細かな部分を調整できます。最終的には、同じプラットフォーム上で数百から数千のGPUにスケールアウトできます。
これがステップ4からステップ5への移行です。徐々に大きなインスタンスで開発とテストを行い、最終的にはスケールアップして、SageMakerが提供する大規模なトレーニングインフラストラクチャを使用します。モデルのアーティファクトを段階的に評価し、SageMakerがメタデータ、スクリプト、ハイパーパラメータを保持し、すべてのアーティファクトをAmazon S3に保存する方法により、過去の作業を簡単に参照できます。6か月前や3年前の実験を再現しようとしても、AWSで行われたものであれば、そのジョブの結果を簡単に参照し、一部のアーティファクトを取得して新しい実験を実行できます。
これが、AWSで基盤モデルを事前学習する方法の概要です。これらはすべて、Amazon SageMakerで提供している分散学習ライブラリによって可能になっています。これらの機能は、効率性と拡張性を提供するデータ並列処理とモデル並列処理の分散学習ライブラリを含め、長年にわたって構築してきたものです。モデル並列処理は、ニューラルネットワークを複数のアクセラレータやGPUに分散させ、最適化されたパフォーマンスを提供する方法です。データ並列パッケージを使用すると、大規模なクラスター全体でモデルのコピーを作成できます。
AWSのネットワークトポロジに最適化されたカスタム通信コレクティブを提供しており、全体の学習時間を最大40%削減できます。これは、スタックのこのレイヤーで長年にわたるイノベーションの結果であり、すべてがSageMakerを通じて利用可能です。お客様も私たちに同意しています。昨日のSwamiの基調講演で聞いたように、Perplexity AIのCEOであるAravind Srinivasは、特にデータ並列処理とモデル並列処理の学習ライブラリを使用して、HyperPodモードで最適化されたパフォーマンスを得るために、SageMakerを喜んで使用しています。
SageMakerの別の便利な機能はWarm Poolsです。学習ジョブAPIは、モデルを学習させるときにインフラストラクチャを作成します。model.fitを呼び出すか、そのPythonの学習スクリプトを実行すると、同時にインスタンスが起動します。クラスターの作成とスクリプトの実行の呼び出しが結合されています。これはコスト効率の面で非常に有用です。つまり、正しいLustreボリュームを指定し忘れてジョブが失敗した場合でも、そのインスタンスが起動したままお金を請求されることはありません - 電源が切れます。非常にコンピューティング効率が高いのです。しかし、開発者にとっては、新しいコード行を送信するだけで8分も待つのは大変です。
この課題に対処するため、昨年Warm Pools機能をリリースしました。この機能を使用すると、同じイメージを使用して新しいジョブを数秒で実行できます。開発者にとっては非常に便利で、学習スクリプトに1行、2行、または3行の編集を加えて、数秒でジョブを実行できます。Warm Pool機能は、SageMaker学習APIを使用した開発に非常に役立ちます。
SageMakerのもう一つの主要な機能は、多くの異なるタイプのインスタンスを使用できることと、スクリプトを実行しようとする基盤となるインフラストラクチャに大きな柔軟性があることです。もちろん、その一つがAWSのカスタムアクセラレータです。TrainiumとInferentiaの機能は両方ともSageMakerで利用可能です。比較可能なAmazon EC2インスタンスと比較して、コストパフォーマンスが大幅に向上しており、Llama 2に対してTrainium 1で最大46%の向上が見られます。最近発表されたばかりのTrainium 2では、さらに優れたパフォーマンスが得られるでしょう。今日のデモでは、SageMaker上のTrainiumを見ていきます。
TRIにおけるSageMakerを用いた大規模言語モデルの開発
このデモは何でしょうか？Llama 2の事前学習を行います。もちろん、視覚的な要素も用意しました。これはBedrockから提供された、Las Vegas stripにいるサングラスをかけた漫画のラマです。SageMakerで70億パラメータのLlamaの事前学習を行います。
なぜこれを行うのでしょうか？なぜこれが有用な演習なのでしょうか？これは、少なくとも数百ギガバイトの独自データ、つまりかなり大規模な独自の言語データセットを持っていることを前提としています。これはオンラインで一般的に入手できない知識、つまり独自の専有データセットです。例えばWikipediaのアーカイブには通常見られないような知識です。
これにより、ドメイン内の精度が向上します。このスモールモデルは、そのドメイン内で驚くほど高い精度を発揮します。繰り返しになりますが、世の中のあらゆることを知っているわけではありませんが、そのデータセットとドメイン内では驚くほどの精度を発揮します。これはもちろん、所有権と柔軟性を高め、サーバーレスホスティングの利用やコスト削減の機会につながります。
いくつかのサンプルノートブックを用意しました。異なるインスタンス、つまりT3 medium、Trainiumの1つ、FSx for Lusterに保存された大規模データの最適化、そしていくつかのWarm Pools、さらに分散トレーニングインフラストラクチャについて説明していきます。では、見てみましょう。
さて、ここからスタートです。サンプルノートブックから始めます。これは公開されているGitHubリポジトリで、Amazon Annapurna MLチームのNeuron Xと、NVIDIAのNemo Megatronを使用しています。これはNVIDIAの分散トレーニングフレームワークであるNemo Megatronです。そして、これをTrainiumアクセラレータ上で実行します。また、PyTorchとcoreのtorchrunフレームワークも使用します。
まず、これをノートブックインスタンス上で実行しています。これはSageMakerのノートブックインスタンスです。私の頼もしいM4インスタンスです。便利なのは、Dockerイメージを作成できることです。ここでノートブックインスタンス上で行うのは、ディープラーニングコンテナを指定することです。これはAWSで構築、追跡、管理、更新を行う完全マネージド型のディープラーニングコンテナです。ソフトウェアフレームワークが変更されると、それに応じて更新します。これらのディープラーニングコンテナを管理し、Dockerfileでこのコンテナを継承して、その上に構築することができます。これは、スクリプトがAWSフレームワークで確実に動作するようにする簡単な方法です。
まず、Dockerイメージをビルドします。次に、最適化された実行のためにFSx for Lustreをセットアップします。その後、SageMaker training APIを使用してトークン化によりデータセットを準備し、Hugging Faceの重みをNeuronフレームワークに変換します。そして、Llama 2のトレーニングを行います。本当にシンプルなDockerイメージで、SageMakerをインポートし、deep learning containerアカウントを指定して、そのイメージを取得します。 そして、これはDockerイメージをビルドする便利なシェルスクリプトです。これで、ローカルにDockerイメージができました。このイメージは、DockerイメージをホストするAWSの別のサービスであるElastic Container Registry (ECR)にプッシュされます。
このプロセスでは、ノートブックインスタンスからローカルにあるDockerイメージをAWSのECRにプッシュします。
Dockerイメージがホストされたら、次はFSx for Lustreのセットアップです。これはCloudFormationテンプレートを使用して自分のアカウントにLustreをデプロイすることで行いますが、コンソールからも可能です。LustreはS3バケットとの双方向データリポジトリを確立します。データがS3バケットにある限り、双方向データリポジトリを有効にできます。データとメタデータがバケットからLustreにコピーされ、Lustreにファイルが書き込まれると、それらがS3に戻されます。 S3からダウンロードするか、ネットワーキングが適切に設定されていれば、Lustreボリュームをノートブックインスタンスにマウントして、ノートブックインスタンスから直接内容を確認することができます。
スタックを作成した後、LLMトレーニング用のデータセットを準備します。 SageMaker SDKを使用して、データセットやモデルのポインタ、キー、FSx for Lustreへのポインタなど、多数のハイパーパラメータを作成します。また、ファイルシステム入力を設定し、PyTorch APIを使用します。これは複雑な例ですが、SageMakerを初めて使う人向けに、Hugging FaceやPythonトレーニングスクリプトを使用して数行のコードでモデルをトレーニングできるようなシンプルな例もあります。
コンテナへのポインタをインポートしますが、これは実際にはPyTorch用のdeep learning containersを指しています。このAPIでは、スクリプト内で独自の前処理関数を定義できます。ソースディレクトリを指定しますが、これはローカルファイルでもGitリポジトリでも構いません。そして、インフラストラクチャとしてTrainium dot 32 xlインスタンスを1つ定義します。
データの前処理を行い、Lustreに保存し、それがS3に複製されると、同じLustreボリューム上でモデルのトレーニングを進めます。 Lustreは特に有用です。なぜなら、数秒でマウントでき、大規模なデータセットを扱う際に便利だからです。Lustreは、複数のジョブに対して素早くマウントできる高性能なファイルシステムを提供し、マウント数に応じてバンドワイズがスケールします。
先ほどと同様に、すべてのハイパーパラメータを指定します。4つのインスタンスを使用し、各インスタンスに32個のTrainiumアクセラレータがあることを示します。トレーニングする70億パラメータのハイパーパラメータを設定し、 FSx for Lustreのポインタを含め、トレーニングジョブを開始します。
残りのハイパーパラメータを設定し、 PyTorch APIを指定してスクリプトをロードした後、model.fitを呼び出します。約束通り、このジョブのすべてのコンテンツは SageMaker control planeに直接ロードされます。これにより、ジョブの開始時刻、ステータス、アーティファクトの保存場所、実行時間を正確に確認できます。出力を表示し、ログを確認し、モデルに必要なあらゆる情報にアクセスできます。その後、モデルをダウンロードし、その上にアプリ全体を構築できます。
では、ステージをTomに引き継ぎます。彼がToyotaについての情報を皆さんと共有します。
TRIでのSageMaker活用事例：実験からモデルサービングまで
ありがとう、Emily。今日は、TRIでのmachine learningの加速にSageMakerをどのように活用しているかについてお話しします。まず、TRIが実際にどのような組織なのか、少しご説明しましょう。現在進行中のプロジェクトの例をいくつか挙げます。最初は、Toyota Supraを使用した自律ドリフト走行のプロジェクトです。ここで、ビデオを再生させていただきます。
これが一つの例です。もちろん、AIがこれらの作業の基盤を築く上で重要な役割を果たしています。二つ目は、私たちが多くのチャレンジ問題に取り組んでいるということです。 大規模なロボティクスグループがあり、研究室で始めたチャレンジ問題に焦点を当てていますが、実際の環境にも出て行き、 そこでシステムを評価しています。これは、私たちが一から社内で構築したロボットシステムの例で、スーパーマーケットの棚から商品を取り出し、補充することができます。最近では、この技術は工場環境にも応用されています。
私たちは250人のチームで、複数の拠点に分かれています。Los AltosとCambridge, Massにチームがあり、human-centered AIやmaterial scienceのチームもあります。最近、generative AIについて私たちが発見したことの一つは、 ロボティクスの文脈において、これまで不可能だと思われていた幅広いタスクを実行できるようになったということです。 これは diffusion policy と呼ばれる技術で、人間の少数の例から非常に複雑なタスクを学習することができるようになりました。
これを基に、TRIのmachine learningチームは、言語、視覚、行動にわたる基盤の構築を目指しています。 言語については、常識的な知識と幅広いアプリケーションの両面から取り組んでいます。例えば、言語はToyotaの企業アプリケーションの文脈で広く応用されていますが、code generationの面でも活用されています。視覚は言語と連携してロボットに「目」を与え、そして行動は、様々なプラットフォームにわたって幅広いタスクを実行するために使用されます。
しかし、この講演はSageMakerについてのものです。そこで、TRIでSageMakerをどのように活用して進歩を加速させているかについてお話ししたいと思います。まず一つ目は一般的な実験です。1から8のインスタンスを使用してtraining jobをスケールアップしています。二つ目は、これらのアイデアを活用して、非常に迅速にスケールアップする方法です。 数個のGPUだけでなく、一度に数百のGPUを使用することができます。最後に、SageMakerをさらに幅広いアプリケーションに使用することもできます。 例えば、モデルのサービング等です。これらはローカルのデバイスでサービングするのが難しいものです。
では、TRIでSageMakerを使って行っている実験について少しお話しします。まず、高レベルの概要として、SageMakerで訓練している 様々なアプリケーションやモデルがあります。一つ目は large language models です。二つ目は monodepth model です。例えば、RGB画像から深度を推論するものです。 三つ目は stable diffusion のようなものです。より良い特徴表現のための言語から画像生成です。四つ目は3D表現です。言語から3D構造を生成するようなものです。これはロボティクスや他の多くのアプリケーションで有用です。
しかし、これらすべてにおいて、SageMakerが様々な理由で非常に有用であることがわかりました。
私たちが直面する課題の一つは、既存のトレーニングインフラストラクチャやクラスターを再利用する必要性です。先ほど説明があったWarm Poolsは、これを実現する一つの方法です。私たちは日々これを活用して、リソースを引き戻し、トレーニングジョブの反復を続けています。
二つ目の課題はスケーリングです。1台から多数のインスタンスへ迅速に移行できる必要があり、また、インスタンスタイプを素早く変更できることも求められます。 さらに、高性能なシステムが必要ですが、SageMakerはバックエンドで非常に最適化されています。 最後に、すべてのサイエンスグループにわたって様々なジョブを実行する柔軟性が必要です。
先ほど見ていただいたように、これらのスケールアップがいかに簡単かを繰り返しお伝えしたいと思います。1台のインスタンスから始めて、そのインスタンスでトレーニングジョブを反復できます。スケールアップが必要な場合は、非常に簡単な変更で可能です。 例えば、インスタンス数を1から8に変更するだけで、実行を素早くスケールアップできます。Galが言及したように、新しいハードウェアが登場した際には、 ハードウェアタイプも迅速に変更できます。この場合、p4インスタンスからp5に切り替えることで、トレーニングジョブのスループットが2倍になり、トレーニング時間を短縮できます。
これらのシステムがいかに高性能かを示すために、複数のインスタンスにわたるスケーリングを見てみると、スケーラビリティはほぼ線形です。 SageMakerは、トレーニングジョブのスケールアップにおいて非常に高性能であることがわかりました。Emilyが先ほど言及したように、これらのデータセットは巨大です。 私たちは数テラバイトのデータセットから始めましたが、Amazon FSx for Lustreですぐに始められるのは便利です。しかし、トレーニングジョブをスケールアップし、必要なデータ量が数テラバイトから0.5ペタバイト以上に増加するにつれて、WebDatasetのような他のリソースを取り込むSageMakerの柔軟性は本当に素晴らしく、トレーニングの実行を加速させました。
繰り返しになりますが、グループは1インスタンスから8インスタンスまでジョブを実行しており、これらはそのアプリケーションの一部です。しかし、それ以上に、私たちはトレーニングの実行をさらに大規模に拡張することができています。TRIでそれを実現している方法の1つを紹介しますと、SageMakerを使用して最先端のLarge Language Models (LLMs)を構築しています。最初は、すべてのシステムを検証するために、Llama 2モデルの一部を再現してきました。このためには、これらすべてのインスタンスにわたってスケーラビリティとパフォーマンスが必要です。
SageMakerが私たちに提供してくれたのは、まさにそのスケーラビリティです。これは最新のハードウェアであるH100を使用しており、ノード数が増加するにつれてほぼ線形のスケーリングが見られます。これは約256台のH100に相当し、このようなトレーニングジョブをLlama 2モデルの事前学習に使用すると、30インスタンスまでスケールアウトした場合、約1週間かかります。1兆以上のトークンを使用することで、最先端のモデルを再現できます。現在、SageMaker上で70億パラメータモデルから13億、34億、70億へとスケールアップしています。
SageMakerの重要な機能の1つは、時間を無駄にしないようにするための修復作業です。Galが先ほど言及したかもしれませんが、これらのジョブをスケールアップすると、ハードウェアが故障することがよくあります。ハードウェアが故障すると、ダウンタイムが発生し、ダウンタイムがあるとモデルのトレーニングができないため、お金がかかってしまいます。私たちにとって、ここでSageMakerのクラスター修復オプションが価値を発揮しました。
フィードマシンの1つが故障した際、クラスターは約10分で復旧しました。トレーニングの実行を非常に迅速に継続することができました。これが事前学習についてです。
次に、アップトレーニングについて話しましょう。これは、事前学習に必要なほど大規模ではないものの、大きなデータセットがあり、特定のドメインに焦点を当てたい場合に使用します。Toyota Research Institute (TRI)はトヨタ中心の組織であるため、日本語は私たちが非常に興味を持っていた分野の1つでした。最先端のモデルの中には、実際には日本語に特化してトレーニングされていないものがあります。確かに少量の日本語トレーニングデータは含まれていますが、それほど多くはありません。
オープンソースで入手可能なすべてのデータを収集すると、100億から1000億のトークンが得られます。これは日本語モデルの追加学習に十分な量です。私たちが発見したのは、130億パラメータを持つLlama 2を追加学習させることで、性能が向上するということです。これは、一流のクローズドソースモデルに対する勝率指標で示されています。
次のステップは、モデルの指示微調整です。これにより、大規模言語モデルが指示に従い、より会話的になります。指示データを使用して微調整を行います。その際、最初の部分が指示で、2番目の部分が期待される応答となります。日本語での追加事前学習と指示微調整を行うことで、日本語において最先端の性能を得ることができます。
このモデルは、例えばLlama 70Bと比べてはるかに小さいにもかかわらず、より高性能です。SageMakerは、TRIでこれらの実験を非常に迅速に行うことを可能にしてくれました。最後に言及したいのは、モデルのサービング等、他のワークロードも実行できるということです。私たちは、オープンソースモデルと社内で開発したモデルの両方をTRI全体で、そして将来的には外部にも提供するために、SageMakerエンドポイントを活用しています。
まとめると、SageMakerを使用する主な3つの領域があります：1〜8ノードの小規模実験、32、64、またはそれ以上のインスタンスを使用する大規模トレーニング、そしてモデルのサービングです。SageMakerは、これらの多様なモデルのトレーニングと一般的な実験において非常に重要な役割を果たしています。
最後に、これらすべてのモデルのトレーニングにSageMakerを使用できたことは素晴らしい経験でした。次回AWS re:Inventに参加する頃には、一度のトレーニングで言語やその他の入力に応じて多様なロボティクスタスクを実行できる基盤モデルが完成しているかもしれません。以上で私の発表を終わり、Galにバトンタッチします。
まとめと今後の展望
Tom、ありがとうございます。最後に、SageMakerについてもっと学び、使い方を知るためのリンクとQRコードをいくつかご紹介したいと思います。皆様、お時間をいただきありがとうございました。 もう少しこの場に残っておりますので、ご質問がある方はどうぞお声がけください。Smart Siftingチームのメンバーも何人かここにいますので、そちらについて質問がある方や詳しく知りたい方はお声がけください。
※ こちらの記事は Amazon Bedrock を様々なタスクで利用することで全て自動で作成しています。
※ どこかの機会で記事作成の試行錯誤についても記事化する予定ですが、直近技術的な部分でご興味がある場合はTwitterの方にDMください。
Discussion