大規模なデータセット（Papers100M のような）でグラフニューラルネットワークを訓練すると、痛みをよくご存知だと思います。エッジリストと特徴量行列を読み込もうとすると、GPU が作業を始める前に瞬時に 24GB 以上の OOM 割り当てでクラッシュします。
私はこれを解決するために作成したカスタム C++ データエンジン GraphZero v0.2 を、公開オープンソース化しました。これによってシステム RAM を完全にバイパスします。
仕組み: 標準ライブラリはすべてをメモリに読み込もうとします。GraphZero は代わりに、生 CSV を 2 種類の高度に最適化されたバイナリ形式（.gl はトポロジー用、.gd は特徴量用）にコンパイルします。
次に POSIX mmap を使ってこれら巨大なファイルを SSD から直接メモリマップします。nanobind を使用して、C++ エンジンは生のメモリポインタを PyTorch にゼロコピーの NumPy 配列として直接渡します。
訓練ループ（GraphSAGE のような場合）では、PyTorch は RAM に 50GB のテンソルがあるとみなします。ターゲットノードのバッチをインデックス付けすると、OS のページフォールトが発生します。オペレーティングシステムは必要な 4KB ブロックのみを NVMe ドライブから自動的に取得します。
パイプラインを飽和させるため、C++ エンジンは batch_random_fanout の隣接サンプリングをマルチスレッド化して、OpenMP を用い、ディスク I/O、CPU サンプリング、GPU 演算を完全に並列化します。
結果: Python がデータセット自体に対して文字通り 0 バイトの RAM を割り当てることなく、50GB のデータセットを訓練できます。
私は低レベルのシステムエンジニアリングとメモリ管理を学ぶためにこれを作りました。リポジトリには、ゼロコピーのマウントをローカルでテストできる GraphSAGE の訓練スクリプトと、合成データセット生成器が含まれています。
このコミュニティには、これを徹底的に検証して、Python API の設計やパフォーマンスについて厳しいフィードバックをいただければと思います！
GitHub: リポジトリ
[リンク] [コメント]
