IBM Researchデータ・ローダーはオープンソースコミュニティのAI モデルトレーニングを強化

大量のデータでAIモデルをトレーニングする際に、ボトルネックをどのように克服すればよいでしょうか。今年のPyTorchの会議で、IBM Researchは大規模LLMトレーニングのための画期的なデータローダーを発表しました。現在 PyTorch ユーザーが利用できるこのツールは、できるだけ幅広いオーディエンスを対象に大規模なトレーニングを簡素化することを目的としています。

高スループットのデータ・ローダーのアイデアは、研究者がモデルのトレーニング中に観察した実用的な問題から生まれました。作業には、進歩的に効率的なGPUを維持しながら、複数のデバイスで大量のデータを処理できるツールが必要だったからです。IBM Researchがこのリリースに関するブログで述べているように、「これはすべて、単に仕事を成し遂げるために必要なツールを作っていた研究者チームのおかげである」。

IBM ResearchのDavis Werheimer氏は、大規模なトレーニングの際に発生する可能性のある課題のいくつかを次のように説明しています。「大規模なトレーニングに関しては、80/20ルールのようなものがあります。公開されているすべての文献の80％は、GPUメモリと通信および計算の間のアルゴリズムのトレードオフに関するものです。しかし、実際に何かを構築しようとすると、時間の80%は、パイプラインが最も狭いボトルネックの速度で実行されるため、他のすべての現実的な問題の長大なロングテールに依存する可能性があります。」

IBMチームはトレーニング・プラットフォームを開発しながら、ボトルネックに直面し続けていました。「GPUの使用がますます向上するにつれて、データ・ローダーがボトルネックになることが増えています」とWertheimer氏は観察しています。

この認識により、二重の開発プロセスが生まれました。「一方で、トレーニング・プラットフォームを進化させ、もう一方ではデータ・ローダーを絶えず進化させて、トレーニング・プラットフォームからの速度要求に遅れを取らず、ボトルネックを避けるようにすることができました」と彼は説明します。

IBM ResearchのLinsong Chuは、データ・ローダーの主要な機能の概要を次のように説明しています。

ステートフルでチェックポイント可能: 「モデルを保存するたびに、データ ローダーの状態も保存されます。また、チェックポイントから回復するたびに、モデルの状態とデータ ローダーの状態の両方を同時に回復する必要があります」と Chu 氏は言います。

チェックポイントの自動再スケーリング： データ・ローダーは、長時間のトレーニング・セッション中に発生したワークロードの変化に自動的に適応します。「トレーニングには数週間から数か月かかる場合が簡単にあり、途中でワークロードを再拡張しなければならない理由はたくさんあります」とChu氏は述べています。

効率的なデータ・ストリーム： このシステムは、シャフリングのためのゼロ・ビルド・オーバーヘッドでデータ・ストリームをサポートします。

非同期分散オペレーション: 「データローダーはノンブロッキングにしたい」とChu氏は説明しています。「データローダーの状態を保存する際、通信を一切行わない形で保存内容を配布したいと考えています。」

動的データ混合： データ・ローダーはさまざまなデータ混合率に適応できるため、進化するトレーニングのニーズに役立ちます。

効率的なグローバル・シャフリング： このツールは、大規模なデータセットを処理する際のメモリーのボトルネックに対処し、データが増大してもシャフリングを効率的に行います。

PyTorchがネイティブ、モジュール式、広範： 適応性と拡張性を考慮して設計されたデータ・ローダーは、将来の成長に備えています。「来年、30兆、50兆、100兆のトークンを処理しなければならないとしたらどうでしょうか？」とChuは言います。「世界は急速に変化しているため、今日を乗り切るだけでなく、明日にも生き残りられるように、データローダーを構築する必要があります。」

IBM研究チームは、数か月にわたってデータローダーを厳密にテストし、大小を問わず何百ものジョブを実行しました。彼らは、安定したスムーズなコード番号を観察しました。さらに、データ・ローダー全体は非同期で動作し、ノンブロッキングです。

「これらすべてを実現するために、PyTorchの組み込み機能をたくさん活用しました」とWertheimer氏は言います。「だからこそ、私たちは貢献し、恩返しをしているのです。」