re:Invent 2024: Amazon SageMaker HyperPodで大規模AIモデル開発
はじめに
海外の様々な講演を日本語記事に書き起こすことで、隠れた良質な情報をもっと身近なものに。そんなコンセプトで進める本企画で今回取り上げるプレゼンテーションはこちら！
📖 AWS re:Invent 2024 - Train large models on Amazon SageMaker for scale and performance (AIM308)
この動画では、Amazon SageMaker HyperPodによる大規模モデルのトレーニングについて解説しています。HyperPodは、Generative AI開発に特化したインフラストラクチャで、モデルのトレーニング時間を最大40%削減し、数千個のアクセラレーターへのスケーリングを可能にします。EKSサポートやContainer Insightsとの統合など、最新のアップデートについても紹介されています。また、Hippocratic AIのCTOであるSaad Godilが登壇し、医療分野でのGenerative AIの実践例として、Auto-pilot agentsによる患者へのアウトバウンドコールの自動化や、Constellation Architectureによる安全性の確保について具体的に説明しています。
※ 動画から自動生成した記事になります。誤字脱字や誤った内容が記載される可能性がありますので、正確な情報は動画本編をご覧ください。
※ 画像をクリックすると、動画中の該当シーンに遷移します。
re:Invent 2024関連の書き起こし記事については、こちらのSpreadsheet に情報をまとめています。合わせてご確認ください！
本編
Amazon SageMaker HyperPodの概要と背景
みなさん、ようこそお越しくださいました。本日は Amazon SageMaker での大規模モデルのトレーニングについてお話しさせていただきます。私は AWS のシニアプロダクトマネージャーの Rekha Seshadrinathan です。本日は同じく SageMaker のシニアマネージャーである Rama Thamman と一緒に登壇させていただきます。また、私たちは Hippocratic AI の CTO である Saad Godil さんをお迎えしています。Saad と彼のチームは SageMaker を使ってモデルトレーニングを加速させており、本日はその方法についてお話しいただく予定です。
数週間前、Puget Sound 地域をボムサイクロンが襲いました。この暴風雨はハリケン並みの強風をもたらし、送電線に前例のない被害を与え、市内の至る所で木々が倒れ、道路も損傷しました。その結果、50万世帯以上が数日間にわたって停電を強いられました。この地域の多くの人々が、電力のない日常生活をどう乗り切るか、携帯電話なしで友人や家族とどう連絡を取るか、インターネットなしでどう仕事をするか、お湯などの基本的なニーズをどう確保するか、車の充電ができない状況でどう移動するかといった課題に直面する中で、私は私たちの生活や仕事の仕方を根本的に変えてきた技術について考えを巡らせました。
メインフレームコンピュータの能力を各家庭にもたらし、スプレッドシートや文書処理といった技術を導入したパソコンにせよ、 元々は軍事通信用に開発されながら、eコマースやビデオストリーミングを生み出したインターネットにせよ、 コンピュータやGPS、カメラを一人一人の手の中に収めた携帯電話にせよ。MapQuest で経路を印刷して持ち歩いていた時代を覚えていますか？そんな時代もありましたよね。
私たちは、テクノロジー業界にまた一つの大きな変革が起きている歴史的な時代に生きる幸運に恵まれています。Generative AI は世界中を席巻し、消費者は驚くべき速さでこれを採用しています。 AWS では、消費者によるこの技術の採用には大きな期待が寄せられていますが、それ以上に、この技術が組織による顧客や従業員のためのイノベーションの方法を根本的に変えると確信しています。実際、Goldman Sachs は、この技術により今後10年間でGDPが7兆ドル増加すると予測しています。
Generative AI の核心は、機械学習とAIにおける最新のイノベーションを活用することです。Transformer アーキテクチャの登場により、Foundation Model（基盤モデル）や大規模言語モデルが前面に出てきました。 Foundation Model とは何でしょうか？それは、インターネットスケールのデータで学習された、しばしば数十億のパラメータを持つ非常に大規模なモデルです。 このような大規模なデータでの事前学習により、これらのモデルは 幅広い文脈でタスクを実行することができます。事前学習済みモデルには多くの魅力的なユースケースがありますが、私たちが話をする組織は、自社のデータでこのようなモデルを学習させる能力にも大きな関心を示しています。これは彼らの業界において差別化要因となり得るからです。
SageMaker Training JobsからHyperPodへの進化
AWSでは、常にお客様の視点から考えて開発を進めてきました。約5年前、より多くのお客様がDeep Learningを始めるようになった時、私たちはSageMaker Training Jobsをリリースしました。 SageMaker Training Jobsは、インフラストラクチャを管理することなく機械学習モデルをトレーニングしたいお客様向けの、フルマネージド型APIを提供します。この仕組みは非常にシンプルで、必要なインスタンスの種類、数、そしてトレーニングスクリプトを指定するだけです。
SageMakerが自動的にクラスターのスピンアップ、必要なコンテナのクラスターへのダウンロード、トレーニングスクリプトの実行を処理します。モデルのトレーニングが完了すると、モデルの成果物を指定した出力先にコピーします。トレーニング完了後は、 インフラストラクチャを自動的にスピンダウンし、使用した分のみお支払いいただく仕組みになっています。
SageMaker Training Jobsは、実験管理ツールと統合されており、 低レイテンシーのためのWarm Poolsや、Spotインスタンスとの連携機能も備えています。Training Jobsの使いやすさから、インフラストラクチャを管理したくないお客様のモデルトレーニングにおいて、今でも非常に人気の高いソリューションとなっています。しかし、長期間にわたって大規模なモデルをトレーニングするお客様との対話を重ねる中で、ここ数年、新たな課題が浮上してきました。
お客様からよく耳にする最初の課題はハードウェアに関するものです。ここ数年、 ハードウェア市場では数多くのイノベーションが起きており、より高速なモデルトレーニングを可能にする新しいチップが数ヶ月ごとにリリースされています。しかし、こうした最新ハードウェアへのアクセスは依然として課題となっています。ハードウェアへのアクセスを獲得できたとしても、クラスターやソフトウェアをそれに対応するよう設定する必要があります。モデルのサイズと トレーニングに使用するデータセットが大きくなるにつれて、モデルトレーニングに必要な計算リソースも継続的に増加しています。実際、過去5年間で毎年4倍以上の成長を遂げており、お客様はこれらのクラスターをさらにスケールアップする必要に迫られています。
クラスターの規模が大きくなるにつれて、 インフラストラクチャの障害が発生する確率も高くなります。Metaが最近発表したLLaMaトレーニングに関する論文では、3時間ごとに1つのGPUで障害が発生したと報告されており、これはデータサイエンティストがインフラストラクチャの障害のデバッグに費やす時間が相当量に上ることを意味します。計算リソースの需要が増大する中、 組織はコストを抑制し、このインフラストラクチャをできるだけ効果的に活用することを望んでいます。
HyperPodの主要機能と技術的詳細
これらの課題に対応するため、昨年のre:Inventにて、Amazon SageMakerは HyperPodを発表しました。SageMaker HyperPodは、Generative AI開発のために特別に設計されたインフラストラクチャで、モデルのトレーニング時間を最大40%削減し、数千個のアクセラレーターにわたってスケーリングすることができます。まず、SageMaker HyperPodは耐障害性の高い環境を提供します。HyperPodクラスターには、インフラの障害を監視し、障害が検出された場合にノードを自己修復するクラスター監視ソフトウェアが搭載されており、トレーニング時間を短縮することができます。
また、HyperPodにはSageMakerの分散トレーニングライブラリが付属しており、モデルとデータをクラスター全体に簡単にシャーディングでき、AWSインフラでより高速にトレーニングを行うことができます。これらの大規模モデルをトレーニングする多くのお客様から、インフラをカスタマイズし、チームが必要とするツールを自由にインストールできるよう、インフラに対するより強力な制御を求める声が寄せられていました。HyperPodでは、インスタンスにSSH接続でき、計算環境を完全にコントロールすることができます。
インフラ面では、SageMaker HyperPodはEC2クラスターを活用しています。分散トレーニングのパフォーマンスにとって重要な、同一のスパインまたはAvailability Zone上に2つのインスタンスが配置され、Elastic Fabric Adapter（EFA）で接続されています。このネットワークインターフェースは、ノード間の通信が頻繁に行われるHPCや分散トレーニングなどのアプリケーションに重要な、高スループットと低レイテンシーを実現するために特別に設計されています。また、高性能な分散ファイルシステムであるFSx for Lustreにもアクセスできます。
HyperPodのコンピュートノードは、NVIDIA CUDAと最新のディープラーニングフレームワークがプリインストールされているDeep Learning AMI上で動作します。また、HyperPodソフトウェアには、AWSのネットワークトポロジー向けに最適化された分散トレーニングライブラリとクラスターの健全性監視ソフトウェアが含まれています。お客様が簡単にこれらのクラスターを作成・管理できるよう、SageMaker HyperPodはノードの作成、更新、削除を行うためのAPIを提供しています。また、AWSコンソールとも完全に統合されており、クラスターの作成とHyperPodクラスターへのアクセス・監視の両方をUIから行うことができます。CLIやCloudFormationテンプレートを使用することもできます。
耐障害性に関して、HyperPodにはハードウェアの障害を常時監視するクラスター健全性監視ソフトウェアが搭載されています。ハードウェアの障害を検出した場合、ノードを再起動するか、再起動が不可能な場合は、AWSが管理する予備のインスタンスプールから無料でノードを交換します。トレーニングジョブをハードウェア障害後に自動再開するように設定して実行すると、最後に保存されたチェックポイントから自動的にトレーニングジョブを再開するため、データサイエンティストや機械学習エンジニアは手動で対応する必要がありません。
コストとパフォーマンスの最適化をサポートするため、Amazon Managed GrafanaとPrometheusとの統合を実現しました。ハードウェアのメトリクスをGrafanaやPrometheusワークスペースにエクスポートし、クラスターの使用状況を監視することで、より効率的な運用が可能になります。 生成AIのパートナーを選ぶ際、多くの組織にとって重要な考慮点の一つが、新しいイノベーションが数ヶ月ごとに登場する中で、いかに素早く対応できるかということです。昨年のHyperPodのローンチ以降、私たちは継続的にアップデートを重ねてきました。
お伝えしたい重要なアップデートの一つは、 数ヶ月前にSageMaker HyperPodでAmazon EKSのサポートを開始したことです。昨年のローンチ時には、Slurmを使用してワークロードをオーケストレーションできましたが、現在ではEKSを使用したオーケストレーションも可能になりました。新しいクラスターを作成するか、 既存のEKSコントロールプレーンにHyperPodのマネージド・コンピュートノードを接続することができます。また、HyperPodのレジリエンシーも強化され、 クラスター作成時にオプションで詳細なヘルスチェックを実行できるようになりました。これにはNCCLテストやDCGM診断が含まれており、トレーニング開始後の障害発生確率を低減することができます。 EKSユーザーに向けては、研究者の利便性も向上させました。HyperPod CLIを使用してジョブを投入できるようになり、使用状況の把握にはContainer Insightsとの統合も実現しています。
Amazon EKSのローンチにより、HyperPodクラスターをトレーニングとInferenceの両方に活用するお客様が増えています。すでにEKSでワークロードをオーケストレーションしている場合、同じノードをモデルの開発とデプロイメントの両方に使用できるため、クラスターの使用効率をさらに最適化することができます。
HyperPodのデモンストレーション
HyperPodの概要説明はここまでとし、より技術的な詳細とデモをお見せするため、Ramaに引き継ぎたいと思います。ありがとうございます、Rekha。これから20分ほどで、HyperPodの技術的な詳細についてご説明し、デモもお見せしたいと思います。挙手でお聞きしたいのですが、大規模言語モデルのトレーニングを実施している方、あるいは将来的に検討している方は何人いらっしゃいますか？かなりの数の方がいらっしゃいますね。
トレーニングのパフォーマンスについてお話ししましょう。 HyperPodは大規模言語モデルやFoundation Modelのトレーニングに特化しています。このようなモデルのトレーニングには、特殊な並列モデルトレーニング技術が必要です。例えば、80億パラメータのモデルをトレーニングする場合、約160GBのGPUメモリが必要になります。NVIDIAのハードウェアを見ると、H100は1GPU当たり80GB、A100は40GBです。つまり、80億パラメータのモデルをトレーニングするには、複数のノードが必要になります。そこで必要となるのがこれらの技術であり、HyperPodではこれらがすぐに使える状態で提供されています。さらに重要なのは、ハードウェアとリソースを効率的に使用するために、スタック全体が最適化されている必要があるということです。
ネットワーク帯域幅は非常に重要です。数十から数百のノードで学習を行う場合、これらのノードは頻繁に通信を行います。バッチ処理の後には、ノード間でグラディエントや重みの交換が行われます。ネットワークがボトルネックになると、処理速度が上がりません。Elastic Fabric Adapter (EFA)を使用すると、400 Gbpsのネットワーク帯域幅を得ることができます。EFAはOSをバイパスし、TCP/IPを使用せず、Libfabricを使用します。また、EFAはNVIDIAのGPUDirectも使用します。リモートダイレクトメモリアクセスが可能で、ネットワークカードは他のGPUのメモリにアクセスできます。これにより、高いスループットを実現しています。Libfabricを使用することで、複数のメッセージ処理インターフェースやNVIDIAのCollective Communications Libraryもサポートされます。
ストレージも非常に重要です。大規模言語モデルの学習を行う場合、大量のデータを使用することが多く、学習アルゴリズムはこのデータを高速に読み込める必要があります。そのためには、データを素早く読み込めるサブミリ秒の遅延で動作する高性能なファイルシステムが必要です。データをインスタンスにダウンロードしてから学習を始めるのでは、GPUがアイドル状態になってしまい、それは望ましくありません。データを素早く読み込める必要があります。FSxを使用すると、数千のインスタンスと数十万のコアに同時にアクセスできます。これら3つの要素により、お客様はより高速な学習が可能になります。ボトルネックが少なく高速な学習が可能になれば、全体的な学習コストを削減でき、そのためお客様はHyperPodを使用することで最大40%高速な学習を実現できているのです。
次に、レジリエンスについて話しましょう。これは非常に重要です。クラスターを起動する前に、ハードウェアが正常であることを確認する必要があります。ハードウェアの障害は処理速度を低下させるため、避けたいものです。そのため、HyperPodではクラスターにインスタンスを提供する前に、詳細なヘルスチェックを実行します。NVIDIAのData Center GPU Management toolsを使用して診断を行い、GPUの健全性を確認するためのストレステストを実施します。また、CPUの健全性もチェックします。これは特にデータのロード時に重要で、CPU使用率を100%使用できること、そしてI/O操作が正常であることを確認する必要があります。これらはインスタンスが提供される前に実施されます。
インスタンスがクラスターに組み込まれると、GPUを継続的にモニタリングします。データセンター管理ツールを使用して、特定のしきい値をモニタリングし、特定のエラーや電力、クロック管理、そして温度に関するしきい値の超過がないかを確認します。これらはポリシーと呼ばれ、バックグラウンドで継続的に実行されます。また、NVIDIA SMIを使用してGPUの健全性を常時チェックし、通知を行います。 ネットワークのヘルスチェックも重要です。ヘルスチェッカーは、インスタンス間のカード間の接続を確認し、これらのインスタンスが効率的に動作できることを確認します。
では、レジリエンス機能について見ていきましょう。エラーは必ず発生するものですから、HyperPodで自己修復できる必要があります。学習を実行する際は、常にチェックポイントを作成しています。大規模言語モデルの学習は通常、何週間も何ヶ月も続くため、エラーから自動的に回復できる必要があります。 学習中にノードが1つ故障した場合、HyperPodは自動的にエラーを検出し、そのノードをクラスターから切り離します。 新しいノードを追加し、詳細なヘルスチェックを実行し、必要なコンテナと最後のチェックポイントをダウンロードして自動的に開始します。この面倒な作業は私たちが担当しますので、お客様は心配する必要はありません。これらすべての処理はバックグラウンドで自動的に行われます。
状況によっては、ノードをオンデマンドで置き換えたい場合があります。インスタンスが応答しなくなることがあるので、そういった場合にノードを置き換えることができます。Slurmを使用している場合は、そのコマンドを実行できます。EKSを使用している場合は、ノードラベルを使用してラベルを変更すると、そのノードは自動的に置き換えられます。数多くのオープンソースツールや様々な選択肢がある中で、使いやすさは非常に重要です。
私たちは、お客様が使用したいツールを自由に選択できる柔軟性を提供したいと考えています。オーケストレーションの観点では、HyperPodの一部として、マネージド型のEKSとマネージド型のSlurmを提供しています。独自のオーケストレーションレイヤー（独自のマネージドサービスなど）をデプロイすることもできます。ジョブの投入に関しては、Ray、Run.ai、kubectlを使用でき、可観測性については、Prometheus、Grafana、CloudWatchをサポートしています。場合によっては、ノートブックやラップトップからクラスターに接続したい場合もあるでしょう。その際はSageMaker StudioやJupyter notebooksを使用してクラスターに接続し、トレーニングジョブを投入することができます。
トレーニングライブラリについては、HyperPodにはPyTorchなどの人気のあるトレーニングフレームワークが、バージョニングなども含めて全て揃っているので、それらについて心配する必要はありません。トレーニング戦略については、一般的なものを全て組み込んでいます。実験については、Kubeflow、Weights and Biases、MLflowなどのツールをサポートしているので、すぐに使い始めることができます。また、マシンイメージやディープラーニングコンテナなどのファーストパーティコンテナも提供しているので、簡単にプルダウンして作業を開始できます。デバイスドライバーやツールキットも全て含まれています。
ハードウェアは非常に重要です。HyperPodでは、NVIDIAの最新かつ最高性能のGPUを利用できます。P5のH200、P5のH100、そして私たちの独自シリコンであるTrainiumインスタンスもHyperPodの一部として利用可能です。ストレージとネットワーキングについても触れましたが、ここでデモをご覧いただきましょう。デモをお見せする前に、これから説明する手順を確認させていただきます。
Hippocratic AIのCTO Saad Godilによる事例紹介
まず最初は準備段階で、クラスターを開始する前に、ネットワーク設定などの特定のリソースを用意する必要があります。次に、HyperPodクラスターの作成方法をお見せします。その一環として、ライフサイクルスクリプトを作成できます。このスクリプトは特定の設定を行うために使用されます。ノードに特定のものをインストールしたい場合は、このスクリプトで実行できます。これはオプションの手順です。その後、クラスター設定の作成方法をお見せします。ここでヘッドノードやコンピュートノードを設定します。また、推論用のグループやトレーニング用のグループなど、異なるグループの作成方法もお見せします。これらは主に管理者が一度だけ実行する設定作業です。
クラスターの設定が完了すると、データサイエンティストが作業を開始できます。トレーニングジョブの実行方法、診断コマンドを実行するためのノードへの接続方法、そして耐障害性の機能についてデモンストレーションを行います。ノードの1つにエラーを注入した際に、HyperPodが自動的にノードを置き換える様子をお見せします。また、温度やGPU使用率など、クラスターの様々なメトリクスを監視する必要があるため、可観測性は非常に重要です。この分野において、私たちは包括的な機能を提供しています。
このデモンストレーションでは、このワークショップを使用します。このリンクはプレゼンテーションの最後に共有させていただきます。このワークショップを使用すると、HyperPodクラスターを簡単にセットアップし、テストして、同じ設定を本番環境で使用することができます。それでは、デモに進みましょう。こちらが先ほど申し上げたHyperPodワークショップです。ここで様々な機能をテストすることができます。まずは前提条件から始めましょう。必要な基本リソースをいくつかセットアップしています。2つのオプションがあります：ネットワークや関連コンポーネントを含むすべてをセットアップする完全デプロイメントモード、または既存のネットワークなどのリソースを活用できる統合型または最小デプロイメントモードのいずれかを選択できます。
このデモでは、完全デプロイメントモードを使用します。これをクリックすると、CloudFormationテンプレートが必要なリソースをすべてセットアップします。アベイラビリティーゾーンやCIDRレンジなどの詳細を変更することができます。スタックを作成すると、必要なリソースがすべて自動的にバックグラウンドで作成されます。EKSクラスターが作成され、VPC、サブネット、関連コンポーネントなど、必要なネットワークリソースも作成されるのが確認できます。この処理には数分かかりますが、完了すると、クラスターの作成に必要なリソースがすべて揃います。
クラスター作成の一環として、Lifecycleの設定を作成する方法をお見せします。これは、クラスターの一部として必要なインストールや設定をセットアップするために使用されます。手順はシンプルです：スクリプトを作成し、必要なコマンドをすべて含め、このファイルを特定の場所に保存します。それをコピーするだけで、クラスター作成時にこの場所を参照します。Lifecycleの設定はこれだけシンプルです。クラスターの作成には2つのオプションがあります：CLIを使用するか、HyperPodコンソールを使用するかです。このデモでは、HyperPodコンソールの使用方法をお見せします。こちらがSageMakerコンソールです。左側にHyperPodコンソールがあります。クラスター管理に移動すると、先ほど追加したEKSクラスターが表示されます。
そのクラスターを選択して、名前を付けます。Demoとクラスターと呼ぶことにします。異なるインスタンスグループを作成できます - 推論用のグループ、トレーニング用のグループなど - そして異なるインスタンスタイプを持つ異種混合クラスターを作成できます。ここでは、G5インスタンスを使用して、1つのグループだけを作成します。8つのインスタンスを割り当て、Lifecycleスクリプトの場所を指定します。また、先ほど説明したディープヘルスチェックを含む耐障害性チェックのいくつかを有効にします。高度な設定を使用して、ストレージやその他の必要な詳細を設定できます。ここでは、リソースの一部として作成したネットワークを選択できます。
最後に、Submitボタンをクリックすると、バックグラウンドでクラスターが作成されます。クラスター内のインスタンス数に応じて、10〜15分ほどかかる場合があります。最初は、インスタンスがPending状態で表示されます。ハードウェアが正常に機能していることを確認するため、詳細なヘルスチェックを実施します。その後すぐに、インスタンスはRunning状態に移行します。この時点でクラスターの準備が整い、データサイエンティストがジョブを投入してトレーニングを開始できるようになります。次は、トレーニングジョブの投入方法をご紹介します。
このワークショップには様々なサンプルがありますが、ここではPyTorch FSDPのサンプルを使用します。もちろん、必要に応じて変更することもできます。ここでは、Kubeflowの構文を使用してジョブを投入する方法をお見せします。
Applyコマンドを使用すると、バックグラウンドでジョブが投入され、8つのワーカーが表示されます。ログを確認することで、どのノードがリーダーノードなのかを確認できます。ご覧の通り、トレーニングが開始され、損失の最小化が進行しています。このように非常にシンプルです。ジョブを投入すればトレーニングが開始されます。場合によっては、各ノードにアクセスしてGPUの使用状況を確認したいこともあるでしょう。
これも非常に簡単です。アクセスしたいノード名さえわかれば、AWS Session Managerを使用してそのインスタンスに接続できます。セッションを開始すると、通常のEC2インスタンスと同様に、別のユーザーとしてログインしたり、topコマンドでプロセスを確認したりできます。また、NVIDIA-SMIを実行してGPUの使用率やヘルス状態を確認することもできます。このように、インスタンスへのアクセスは非常に簡単です。
では、レジリエンシー機能を見ていきましょう。レジリエンシーには2つのシナリオがあります。1つ目は、エラーが発生した際にHyperPodが自動的にノードを置き換えるケース、2つ目は、必要に応じて手動でノードを置き換えるケースです。このデモでは、HyperPodが自動的にノードを置き換える方法をお見せします。ここでは、一方にワーカーのセットアップを表示し、もう一方にノードを表示します。8つの異なるノードがあります。
これから3番目のノードにログインしてみましょう。Session Managerのセッションを開いて、ログにエラーを注入してみます。エラーメッセージをログに書き込むと、HyperPodが自動的にそのノードを検知します。3番目のノードがnot readyの状態になり、Workerがpending状態になるのが確認できます。そうするとHyperPodがそのクラスターから負荷を取り除き、新しいノードを追加して、詳細なヘルスチェックなどを実行します。ご覧の通り、下から3番目に新しいノードが追加されました。ノードが追加されると、コンテナのダウンロードなど、すべてのセットアップが実行されます。それが完了すると、中断された時点からトレーニングを再開するためにチェックポイントも取得されます。その特定のノードにログインして、トレーニングの進行状況を確認することができます。このように、HyperPodが自動的に面倒な作業を処理してくれるので、ユーザーは心配する必要がありません。
次に、可観測性について見ていきましょう。多数のノードがある場合、温度やGPU使用率、CPU使用率などを監視できることは非常に重要です。特にトレーニング時には、最適な使用率が得られない可能性があるため、繰り返し確認しながら、これらすべてを監視する必要があります。HyperPodコンソールで、Container Insightsをクリックすると、観測ページに移動し、クラスターレベルでメトリクスを監視できます。
8つのノードがあり、CPU使用率とGPU使用率が高いことが確認できます。これは最終的に見たい理想的な状態です。また、ネットワークの使用状況を確認するための入出力データ量など、他にも多くのメトリクスがあります。先ほど言及した温度も監視できます。GPU温度の監視は特に重要で、その他にも活用できるメトリクスがたくさんあり、アラートを設定することもできます。
Podレベルのメトリクスも確認できます。約95個のコンテナが実行されているのが分かります。より詳細なレベルで監視が可能です。ここで、あるノードのGPU使用率を見てみましょう。8つのWorkerがあり、GPU使用率を確認したいWorkerを選んで詳しく見ることができます。このように、豊富な機能が用意されています。
ここまでで、クラスターの迅速なセットアップ方法、ジョブの投入方法、監視方法、そして回復性の機能について見てきました。ここで、HyperPodsの利用経験についてお話しいただくため、Saad Godilさんをお招きしたいと思います。
Hippocratic AIのConstellation Architectureと技術的課題
はい、ありがとうございます、Rama。皆さん、こんにちは。Hippocratic AIの共同創業者兼CTOのSaad Godilです。Hippocratic AI について、私たちが取り組んでいることをご紹介させていただきます。私たちは、世界的な看護師不足の問題に対して、臨床業務を実行できるGenerative AIベースのAuto-pilot agentsを構築することで取り組んでいます。具体的には、患者へのアウトバウンドコールを行っています。私たちの技術を使って、患者さんに自律的に電話をかけ、様々なタスクを実行することができます。
Generative AIについて話す時、この技術には様々な使い方があります。医療分野のGenerative AI企業の多くは、私たちが「Copilot」と呼ぶソリューションを開発していることにお気づきかと思います。EMR SummarizationやAmbient listeningなどがその良い例です。これらの製品の多くは、人間の作業効率を向上させることを目的としています。しかし、その欠点は大きな効果が得られないことです。生産性を最大でも10%から15%程度しか向上させることができず、それでは看護師不足の解決にはなりません。私たちにはもっと大きなレバレッジが必要で、そこでAuto-pilot agentsの出番となります。これらの自律型エージェントはタスクを完全に実行でき、はるかに大規模なスケーリングを可能にします。
具体的な例をいくつかご紹介させていただきます。術前の電話連絡の場合、例えば大腸内視鏡検査の予約がある場合、検査の3日前に電話をかけて、腸管洗浄剤の服用開始と低繊維食の開始を確認します。検査前日には固形食を中止してクリア液体食を開始するよう連絡し、検査6時間前には水分摂取を中止するよう連絡します。このように、ケアプランの指示事項すべてを管理し、患者さんとコミュニケーションを取って、それらを確実に守っていただけるようサポートできます。
もう一つの例は術後の電話連絡です。処置後に退院される際、私たちが電話をかけて、お薬を受け取られたかどうかを確認し、合併症の有無を確認し、医師から指示されたケアプランについての質問にお答えします。また、Health risk assessmentやAnnual wellness visitの問診票など、患者さんの都合の良いタイミングで効率的に情報を収集することにも大きな関心が寄せられています。これらは、Generative AIで完全に自動化し、効率的に実行できると考えています。
最も興味深いのは、この先どのような展開が期待できるかということです。現在は、私たちが補完したり自律的に実行したりできるタスクについて話していますが、これらの電話連絡のコストがゼロに近づくにつれて、今日では実施されていない多くのケアマネジメントの使用例が生まれると考えています。その良い例として、あるお客様から熱中症リスク評価についての相談を受けました。熱波の際に、最もリスクの高い患者層を特定し、即座に全員に電話をかけることができます。これこそがGenerative AIのパワーです。必要に応じて無限にスケールアップし、毎日患者さんに電話をかけ、健康リスク評価を行い、熱中症のリスクがないことを確認できるのです。
私たちのシステムにおいて、安全性は最も重要な特徴の一つです。私たちは「Constellation Architecture」というコンセプトを導入しました。というのも、1つのモデルだけでは実現できないことが分かったからです。このConstellation Architectureは20以上の異なるモデルで構成されており、専門モデルがメインモデルと協力して、必要な安全性を確保しています。市販薬の使用制限、処方薬の服薬順守、検査結果の分析などに特化したモデルがあり、適切なアドバイスを提供します。これらの専門モデルにより、安全性に関する各軸に特化した学習が可能となり、必要な安全性レベルを実現するモデルを構築できます。
これらのモデルは、オープンソースモデルをベースに、私たちの独自の専用データセットで学習させています。この専門化によって、必要な安全性を確保するための最後の性能向上を実現することができます。
最近、私たちはConstellation Architectureをアップグレードしました。最初のPolaris 1.0からPolaris 2.0へと進化させました。 最大のモデルのパラメータ数は70億から405億へと6倍に増加しました。電話で患者さんとリアルタイムに会話するアプリケーションとして、レイテンシーは非常に重要です。これらのモデルをアップグレードした際、推論能力は大幅に向上しましたが、低レイテンシーを維持することは大きな課題でした。右側のボックス＆ウィスカープロットで示されているように、革新的な技術により同じレイテンシーを維持することができました。
実装した技術には、Calibrated FP8 Quantizationが含まれています。通常、量子化すると性能がトレードオフされますが、学習セットに合わせて調整することで、臨床的な安全性を維持することができました。また、Aggressive Prefix Cachingの実装も重要なブレークスルーでした。長い会話では入力プロンプトが非常に大きくなり、特に大規模モデルでは処理に時間がかかる可能性があります。しかし、次の通話で変更されるのは、新しく追加される患者の発言だけです。前回のクエリのプリフィルを活用することで、時間を節約しレイテンシーを低減できます。また、複数のノードで多数の通話を並行して処理するため、Conversation-based Routingを実装し、同じ会話からの通話が同じモデルに送られるようにしました。
これらのモデルはすべてAmazon SageMaker HyperPodで学習されています。 HIPAAに準拠し、患者データが極めて重要であるため、私たちは複数の異なるクラスターを維持しています。機密性の高い患者情報を扱う本番環境の推論ワークロードは、完全に別のアカウントとクラスターで実行され、学習と開発作業は別のアカウントで行われています。適切なデータ分離、セキュリティ、権限を確保するため、異なる環境には複数のアカウントを使用しています。
これが私たちのTrainingクラスターの構成です。 ストレージにはFSxとS3を使用して、GPUへの高速アクセスを実現し、モデルのトレーニングを行っています。ログとメトリクスの収集にはGrafanaとPrometheusを使用しています。これらのGPUはかなり高価なので、使用率を確認するための綿密なモニタリングが必要で、目標とする使用率を達成できているかを確認するためにこれらのメトリクスを活用しています。次に、Inferenceアーキテクチャについてですが、これも似たような構成になっています。 ご覧の通り、ユーザーアクセスはあまり多くありません。というのも、Productionクラスターへのアクセスは非常に制限されているためです。上部に追加されているのがApplication Load Balancerで、APIゲートウェイマネージャーとしてKongを使用して、モデルがAPIを通じてアプリケーション層と通信できるようにしています。
Ramaのプレゼンテーションでカバーされた機能の多くを、私たちも実際に活用しています。まず、クラスターの起動と停止が簡単にできる点は、多数のクラスターを運用している私たちにとって非常に重宝しています。クラスターの起動と停止を完全に自動化できることは極めて重要で、必要に応じてTerraformスクリプトで制御しています。複数のクラスターを運用しているため、InferenceとTrainingの異なるワークロード間でリソースのバランスを取る必要がよくあります。GPUを別のクラスターに移動できることは、限られたリソースを最適化する上で非常に重要です。新しいクラスターを立ち上げる際には、HyperPodのLifecycleスクリプトが自動的にノード設定を行います。GPUの信頼性は極めて重要で、故障したノードをバックアップ用のReserved Instanceと交換できるHyperPodの機能を活用することで、すぐに運用を再開できます。
今後の展望についてですが、 私たちは本当にElastic GPU Computeについて考えています。患者さんへの電話は、1日のうち患者さんが対応可能な6時間に限られています。現在、AWSと協力して、Inferenceワークロードをスケールアップし、ビジネス目標を達成するために、ピーク時のGPUコンピューティングリソースをどのように確保できるか検討しています。
もう一つ、私が本当に期待している機能は、SageMaker HyperPod上のEKSです。ご想像の通り、私たちのInferenceは大幅にスケールアップしています。複雑なワークロードを管理するために、スケールアップとスケールダウンの両方の自動化が必要です。先ほど話したように、オーケストレーションとスケールアップのすべてを自動化する必要があります。EKSを使用したKubernetesでの構築が正しい戦略だと考えており、次はそこに投資して構築を進めていく予定です。
医療は昔から、医療従事者の数という制約を常に抱えてきました。システム全体がその制約を前提に設計されてきたのです。そして今、Generative AIによって、初めて医療における真の豊かさを実現できる可能性が出てきました。これこそがGenerative AIが私たちの実現を助けてくれることであり、私たちはこのビジョンを実現するためにHyperPodとSageMakerチームとのパートナーシップを喜ばしく思っています。
まとめと今後の展望
最後に、主なポイントをまとめたいと思います。 SageMaker HyperPodでは、当社が提供する耐障害性機能と分散学習ライブラリにより、学習時間を最大40%短縮することができます。私たちが提供するAPIとコンソールUIを使用することで、数千個のアクセラレーターへ簡単にスケールアップできます。また、GPUやCPUの使用率を監視できるモニタリングインフラを活用して、コストとスケールを最適化することができます。さらに、HyperPodは柔軟でカスタマイズ可能なので、インスタンスにSSH接続してスタックを好みに応じてカスタマイズすることができます。
ローンチ以来、 あらゆる規模、あらゆる業界の企業がGenerative AI開発を加速するためにSageMakerを活用しています。Hugging Face、Perplexity AI、Hippocratic AIといった有力なスタートアップから、SalesforceやThomson Reutersといった大企業まで含まれています。これらの企業の多くが、 HyperPodによって学習時間が何百時間も節約され、データサイエンティストの生産性が向上したと述べています。
HyperPodを始めてみたい方のために、いくつかのリソースをご紹介します。ぜひ写真を撮っておいてください。また、これらの多くのお客様が本日14時30分から開催される顧客パネルにも参加されます。HOPPR AI、Hugging Face、WRITER AIの詳細についてお聞きになりたい方は、同じく14時30分からAIM 229で開催されるセッションにもぜひご参加ください。ご清聴ありがとうございました。引き続き質問をお受けいたします。
※ こちらの記事は Amazon Bedrock を利用することで全て自動で作成しています。
※ 生成AI記事によるインターネット汚染の懸念を踏まえ、本記事ではセッション動画を情報量をほぼ変化させずに文字と画像に変換することで、できるだけオリジナルコンテンツそのものの価値を維持しつつ、多言語でのAccessibilityやGooglabilityを高められればと考えています。
Discussion