[論文レビュー] Flashlight: Enabling Innovation in Tools for Machine Learning

Jacob Kahn, Vineel Pratap+12 more|arXiv (Cornell University)|Jan 29, 2022

Explainable Artificial Intelligence (XAI)被引用数 14

ひとこと要約

Flashlightは、新しいテンソル計算、メモリ管理、分散学習技術の迅速なプロトタイピングを可能にする、軽量でモジュール式の機械学習ライブラリです。最小限のオーバーヘッドと内部構造の完全な制御を備えており、テンソル加算のようなコア演算をカスタム実装に簡単に切り替えることができます。PyTorch や TensorFlow と同等のパフォーマンスを達成するとともに、メジャーなフレームワークへの上流貢献を可能にします。

ABSTRACT

As the computational requirements for machine learning systems and the size and complexity of machine learning frameworks increases, essential framework innovation has become challenging. While computational needs have driven recent compiler, networking, and hardware advancements, utilization of those advancements by machine learning tools is occurring at a slower pace. This is in part due to the difficulties involved in prototyping new computational paradigms with existing frameworks. Large frameworks prioritize machine learning researchers and practitioners as end users and pay comparatively little attention to systems researchers who can push frameworks forward -- we argue that both are equally important stakeholders. We introduce Flashlight, an open-source library built to spur innovation in machine learning tools and systems by prioritizing open, modular, customizable internals and state-of-the-art, research-ready models and training setups across a variety of domains. Flashlight allows systems researchers to rapidly prototype and experiment with novel ideas in machine learning computation and has low overhead, competing with and often outperforming other popular machine learning frameworks. We see Flashlight as a tool enabling research that can benefit widely used libraries downstream and bring machine learning and systems researchers closer together. Flashlight is available at https://github.com/flashlight/flashlight .

研究の動機と目的

• フレームワークのサイズと複雑さの増大に伴い、機械学習におけるイノベーションが難しくなっている現状に対処する。
• 機械学習研究者とシステム研究者との間の溝を埋めるために、システムイノベーションに最適化されたフレームワークを提供する。
• カスタムメモリマネージャーや分散計算モデルを含む、新しい計算パラダイムの迅速なプロトタイピングを可能にする。
• PyTorch や TensorFlow のようなモノリシックなフレームワークの代替として、最小限でモジュール式であり、内部コンponentのエンドツーエンド実験を可能にする。
• 生産環境で動作する高パフォーマンスな環境で新技術をテスト・検証することで、メジャーなフレームワークへの上流貢献を促進する。

提案手法

• テンソル演算と計算グラフのための明確で合成可能なAPIを持つ、最小限でモジュール式のC++ライブラリとしてFlashlightを設計する。
• 複数の分野（例：自然言語処理、コンピュータービジョン）にわたる完全で生産環境対応のリファレンスモデルを実装し、システム研究のためのテストベッドを提供する。
• プラグイン可能なメモリマネージャーやカスタムアロケーターバックエンドを通じて、メモリ割り当てとテンソルレイアウトに対する低レベルの制御を提供する。
• 任意のデバイスマッピングと計算スケジューリングをサポートし、新しい分散学習戦略の実験を可能にする。
• 継承やコンpositionを介して、既存のオペレーターやベンチマークを変更せずにコア演算（例：テンソル加算）を簡単に置き換えられるようにする。
• 実験中に演算レベルでのメモリ動作とパフォーマンスを監視できる、カスタムのテレメトリおよびプロファイリングツールを統合する。

実験結果

リサーチクエスチョン

• RQ1 最小限でモジュール式の機械学習ライブラリは、モノリシックなフレームワークを上回るシステムレベルのイノベーションを可能にするか？
• RQ2 割り当てとフラグメンテーションに対する完全な制御のもとで、研究者が新しいメモリ管理技術をどれほど効果的にプロトタイピング・ベンチマークできるか？
• RQ3 カスタムテンソル演算（例：最適化された加算）をフレームワークに統合する際、既存のコードを破壊せずに、どの程度の柔軟性で統合できるか？
• RQ4 一般化された分散計算およびオプティマイザー状態のシャーディングアプローチを、柔軟で拡張可能なフレームワークで実装・評価できるか？
• RQ5 大規模な複雑なコードベースに依存せずに、新しい計算プリミティブに関する研究をエンドツーエンドでスケールしてテストするにはどうすればよいか？

主な発見

• Flashlightは、PyTorch や TensorFlow と比較して、コードベースとバイナリサイズが桁違いに小さく、イテレーションの高速化と保守の容易化を実現した。
• 研究者は、大きなキャッシュブロックの分割を制限するメモリマネージャーを迅速にプロトタイピング・テストでき、複数のモデルで内部フラグメンテーションを20％以上削減した。
• 既存のコードを変更せずに、テンソル加算の新しい実装をすべてのオペレータにシームレスに置き換えることができ、完全な相互運用性を示した。
• フレームワークのモジュール式設計により、バッファのシャーディングや動的デバイススケジューリングを含む、一般化された分散学習技術の開発が可能になり、ハードコードされた仮定が一切不要だった。
• パフォーマンスベンチマークの結果、Flashlightのリファレンス実装は、エンドツーエンドのトレーニングワークロードにおいて、PyTorch や TensorFlow の実装と同等またはそれ以上の性能を示した。
• Flashlightの研究から得られた知見は、キャッシュアロケータの改善を含め、PyTorch に成功裏に上流貢献され、フレームワークの実世界への影響を裏付けた。