IBMとPyTorchによる、AIモデル学習を大幅に改善する1行のコード変更

記事をシェアする:

IBM ResearchとPyTorchは、AI学習におけるチェックポイント作成のコスト効率改善に取り組みました。

私たちは今、基盤モデルがもたらしたAIブームの真っただ中にいます。これらのモデルの多くはサイズが巨大ですが、学習データセットはさらに巨大です。何十億ものパラメーターを持つモデルの学習に使われるデータは、テラバイトに及ぶこともあります。これらのモデルは、音楽の生成から新しい化学物質の発見、大規模な企業プロセスの自動化まで、あらゆる種類の目的に使用することができます。

IBMは、パブリック・クラウドやプライベート・クラウド上でシームレスに動作する、企業向けの信頼できる基盤AIモデルの構築に取り組んでいます。私たちはAIの未来のために、クラウドネイティブでオープンソースのスタックを構築していますが、これはPyTorchを利用してAIシステムの構築をよりシンプルにしています。一般に、モデルの学習中、データサイエンティストは障害への備えとして定期的にチェックポイントをシステムの永続的なストレージに書き込み、障害から復旧できるようにしておきます。GPUの無駄な時間を減らすために、チェックポイントの書き込みは迅速に行われる必要があって、通常、これはNFSのような高速な共有ファイルシステムに書き込むことで行われます。当然のことに、モデルのサイズが大きくなるにつれて、このチェックポイントのデータサイズも大きくなります。例えば、110億個のパラメーターを持つモデルの場合、モデルの重みと最適化アルゴリズム（optimizer）の状態を含む1回のチェックポイントには130GBのストレージが必要になります（32ビット精度の場合）。

私たちは、GPU時間を犠牲にしない、より安価なストレージを使用できるかどうかを検証しました。共有ファイルシステムからオブジェクト・ストレージに単純に切り替えただけでは、説明のつかないエラー（NCCLのタイムアウトやサイレント障害など）が発生しました。IBMとPyTorchのチームは協力して、S3 API†を持つ、IBM Cloud Object Storageのようなオブジェクト・ストレージをサポートするために、PyTorch内の分散チェックポイント・アルゴリズムを修正しました。これには、従来のファイルシステムAPIとS3FS（オブジェクトストアを使用した）ファイルシステムAPIの違いを調べる必要がありました。

現在の分散チェックポイント・アルゴリズムの実装は、強い書き込み後読み取り整合性を保証する、POSIX（Portable Operating System Interface）に準拠したファイルシステムを想定しています。PyTorch FSDP (Fully Sharded Data Parallel)では、各GPUがモデルと最適化アルゴリズムの状態の１シャードを担当し、シャードのテンソルをそのGPU用に指定された特定のディレクトリに書き込みます。共有ファイルシステムでは、オーケストレーター（ランク0のGPU）がディレクトリを作成すると、他のすべてのGPUからすぐに見えるようになります。しかし、オブジェクト・ストア上のS3FSのように、結果整合性に基づくファイルシステム・レイヤーでは、この仮定は成り立ちません。

このような整合性の問題は再現不可能なエラーとして現れるため、チームは多くのログファイルを調べ、根本原因を特定するために多くの仮説を立てなければなりませんでした。最終的にこの修正は、各GPUが、ローカルビューにディレクトリーが存在しない場合に作成できるようにするということでした。これらの書き込みは競合しないので、S3FSの結果整合性で十分に機能します。この修正により、PyTorchの分散チェックポイントはS3FS APIを持つオブジェクト・ストレージへの書き込みをサポートできるようになりました。

チェックポイントを記述するための別にありえるアプローチは、すべての重みと最適化アルゴリズムの状態を1つのノードのRAMに集めることでしたが、これは全データを集める段階でクラッシュする可能性があり、そうなるとGPU時間を大幅に浪費し、大きな障害となってしまいます。私たちの方法では、私たちの測定によると、110億パラメータモデルの場合、チェックポイントにかかる時間は1時間以上から数分に短縮されました。また、オブジェクト・ストレージは最も手頃なファイル・ストレージ・システムの1つであるため、研究者たちはこのシステムに切り替えることで、学習費用を節約することができるようになりました。

IBM Researchの主席研究員でプロジェクト・リーダーの一人であるRaghu Gantiは、チームがこれまで使用してきたストレージの種類について、次のように語りました。「いわば通勤用の普通の車で済むところを、これまでスーパーカーを使ってきたようなものです」「オブジェクト・ストレージの手頃な価格で、洗練されたモデルの学習を行うことができるようになりました」

IBMとPyTorchは、この技術的な成果と、安価なネットワーク機器を使用してメモリスケジューリングを改善するためという以前の成果を組み合わせることで、AIモデルの学習と実行を、より迅速で、よりコスト効率的で、さらに大規模なモデルにスケールできるように、そしてこれらすべてをIBMのハイブリッドクラウド・プラットフォームでできるように進めています。

† AWS S3はAmazonのSimple Storage Serviceです。S3のAPIは、IBM COSのような他のオブジェクト・ストレージ・プロバイダーにも実装できる標準です。

この記事は英語版IBM Researchブログ「IBM and PyTorch change one line of code to massively improve AI model training」（2023年5月4日公開）を翻訳し一部更新したものです。