NVIDIA/Megatron-Energon

Megatron Energon – 大規模マルチモーダルモデル学習用データローダー

概要 – このPythonパッケージ（および小さなCLI）は、NVIDIAのMegatron‑LM学習フレームワーク向けに、大規模なマルチモーダルデータセット（画像、動画、テキストなど）を準備、インデックス化、ストリーミングします。単独でも使用できます。

重要性 – 今日の大規模ビジョン言語モデルやビデオ言語モデルの学習には、多くのストレージシステムに分散したペタバイト規模のデータが必要です。Megatron Energonはその重労働を処理します：

主な機能（READMEに記載）

機能	説明
CLI (`energon`)	WebDatasetのインデックス準備、メディアメタデータの収集、データセットの読み取り専用仮想ファイルシステムとしてのマウントなど。
Python API	`get_train_dataset`や`get_loader`などの関数が、学習ループで反復可能なPyTorch互換のデータセット/ローダーを返します。
分散サポート	`WorkerConfig`を使用すると、ランク、ワールドサイズ、ワーカー数を指定でき、同じコードが単一GPUまたはフルマルチノードクラスターで実行されます。
マルチモーダル対応	画像、動画、テキストなどに対応し、`TaskEncoder`に登録された「クッカー」を使用して、生のサンプルをモデルが期待するテンソルに変換します。
再現性と再開性	インデックスは`.nv-meta`フォルダーに保存されます。ローダーはデータセット全体を再処理せずに再起動できます。
インストール	`pip install megatron-energon`（安定版）またはリポジトリから直接インストールできます。

典型的なワークフロー

データの準備 – energon prepare /path/to/dataset（またはJSONLファイルに対して）を実行して、.nv-metaインデックスとオプションの分割を生成します。
（オプション）メディア統計の収集 – energon prepare-media …で画像サイズ、動画の長さなどを収集し、フィルタリングやバケティングに使用できます。
コードでの読み込み – Python APIを使用してtrain_dsとtrain_loaderを取得し、学習スクリプト内でバッチを反復処理します。
Megatron‑LMとの統合 – バッチをMegatronのモデル並列学習ループに供給して、大規模LLMまたはマルチモーダルモデルの学習を行います。

対象読者 – 非常に大規模なビジョン言語、ビデオ言語、または他のマルチモーダルモデルを構築する研究者やエンジニアで、スケーラブルで再現可能なデータパイプラインを必要とする方。研究用コードのため、APIは変更される可能性があります。

詳細はこちら – 完全なドキュメントはhttps://nvidia.github.io/Megatron-Energon/、GitHubリポジトリで問題、テスト、例を参照してください。