Lightning-AI/lightning-thunder

何を解決するか

Thunder は、PyTorch の最適化をトレーニングおよび推論の両方で簡素化することを目的としたソースツーソースコンパイラです。従来のコンパイラの不透明性を避けつつ、"非最適化された" PyTorch エイジ mode と高度に最適化された実行の間のギャップを埋め、パフォーマンス向上を透明に適用できるようにします。

動作方法

Thunder は3つの主要な段階で動作します：

1. 取得：Python bytecode を解釈して、モデルを直線的な Python プログラム（中間表現）に変換します。
2. 変換：この IR に変換システムを適用し、分散処理、数値精度の変更、または計算グラフの修正を行います。
3. 実行：結果のトレースを、融合エンジン（NVFuser や torch.compile）、専用ライブラリ（cuDNN SDPA、TransformerEngine）、カスタム Triton/CUDA カーネル、または標準の PyTorch エイジ操作にルーティングします。

対象ユーザー

PyTorch モデルのパフォーマンスを最大化したい開発者や研究者向けです。特に NVIDIA ハードウェア（Blackwell を含む）上で LLM やその他の大規模モデルを扱っている人、および量子化、ミックスド・プレシジョン、分散戦略を検査可能で拡張可能な方法で適用したい人におすすめです。

主な特徴

• パフォーマンス向上：PyTorch モデルを最大 40% 速く実行可能。
• 柔軟な精度：FP4、FP6、FP8、およびミックスド・プレシジョン戦略をサポート。
• 分散スケーリング：Tensor Parallelism (TP)、Pipeline Parallelism (PP)、Data Parallelism (DP) を内蔵サポート。
• 検査可能な IR：Python風の IR を使用し、ユーザーが演算をプロファイリング、カーネルにマッピング、インタラクティブにプログラムを検査できる。
• 合成可能なレシピ：最適化をレシピとしてバンドルし、異なるモデルファミリー間で簡単に移行可能。
• CUDAGraphs 対応：CUDAGraphs を通じて CPU オーバーヘッドを削減するプラグインを提供。