lightning thunder

PyTorch モデルを超高速で作成します。

---

Lightning.ai • パフォーマンス • 始めましょう • インストール • 例 • Thunder の内部 • 参加してください! • ドキュメント

⚡ ライトニングサンダーへようこそ

Thunder は PyTorch モデルを超高速にします。

Thunder は、PyTorch のソース間コンパイラーです。さまざまなハードウェア エグゼキューター (nvFuser、torch.compile、cuDNN、TransformerEngine FP8 など) を一度に組み合わせて使用​​することで、PyTorch プログラムを高速化します。

シングル GPU 構成とマルチ GPU 構成の両方をサポートします。 Thunder は、使いやすく、理解しやすく、拡張可能であることを目指しています。

注記

ライトニングサンダーはアルファ版です。気軽に参加していただけますが、途中でいくつかの困難が生じることが予想されます。

シングル GPU のパフォーマンス

Thunder は、最適化とクラス最高のエグゼキュータの使用の複合効果により、コンパイルされていない標準の PyTorch コード (「PyTorch 熱心」) よりも大幅な高速化を実現できます。以下の図は、LitGPT で実装された Llama 2 7B の事前トレーニング スループットを示しています。

上のプロットに示されているように、Thunder は、nvFuser、torch.compile、cuDNN、TransformerEngine FP8 などのエグゼキューターの組み合わせを使用して、H100 の Eager コードと比較してトレーニング スループットで 40% の高速化を達成しています。

マルチGPUパフォーマンス

Thunder は、複数の GPU でモデルをトレーニングするための DDP や FSDP などの分散戦略もサポートしています。次のプロットは、FP8 混合精度を使用しない Llama 2 7B で測定された正規化されたスループットを示しています。 FSDP のサポートは進行中です。

始めましょう

追加のインストールやセットアップを必要とせず、Thunder を使い始める最も簡単な方法は、Zero to Thunder チュートリアル スタジオを使用することです。

サンダーをインストールする

Thunder はアルファ版であり、最新の開発はmainブランチで行われています。次のようにして、 mainブランチから Thunder の最新バージョンをインストールできます。

pip install git+https://github.com/Lightning-AI/lightning-thunder.git@main

最高のパフォーマンスを実現するには、次の追加の依存関係を指定して Thunder をインストールします。

• 次のように、PyTorch 2.5.1 用にビルドされたプレリリースの nvFuser をインストールします。

 # install nvFuser built for the matching stable PyTorch
pip install --pre nvfuser-cu121-torch25

• 次のように cudnn をインストールします。

 # install cudnn
pip install nvidia-cudnn-frontend

高度なインストール オプション

インストールしていじって貢献する

Thunder をいじったり、Thunder に貢献したりすることに興味がある場合は、Thunder リポジトリのクローンを作成し、pip の編集可能モードでインストールすることをお勧めします。

git clone https://github.com/Lightning-AI/lightning-thunder.git
cd lightning-thunder
pip install -e .

テストを開発して実行する

上記で説明したように、lightning-thunder リポジトリを複製し、編集可能なパッケージとしてインストールした後、開発要件をインストールすることで、Thunder を開発するための環境をセットアップできます。

pip install -r requirements/devel.txt

ここでテストを実行します。

pytest thunder/tests

Thunder は非常に徹底的にテストされているため、しばらく時間がかかることが予想されます。

こんにちは世界

以下は、Thunder で PyTorch コードをコンパイルして実行できる方法の簡単な例です。

 import torch
import thunder


def foo ( a , b ):
    return a + b


jfoo = thunder . jit ( foo )

a = torch . full (( 2 , 2 ), 1 )
b = torch . full (( 2 , 2 ), 3 )

result = jfoo ( a , b )

print ( result )

# prints
# tensor(
#  [[4, 4]
#   [4, 4]])

コンパイルされた関数jfoo 、元の関数と同じように PyTorch テンソルを取得して返します。そのため、Thunder によってコンパイルされたモジュールと関数は、より大きな PyTorch プログラムの一部として使用できます。

鉄道模型

Thunder は初期段階にあるため、まだ本番稼働には使用しないでください。

ただし、Mistral、Llama 2、Gemma、Falcon など、LitGPT でサポートされる LLM の事前トレーニングおよび微調整では、すでに優れたパフォーマンスを提供できます。

LitGPT と Thunder を一緒に実行する方法については、LitGPT の統合を確認してください。

Inside Thunder: コア機能の概要

Python 呼び出し可能モジュールまたは PyTorch モジュールを指定すると、Thunder は次のような最適化されたプログラムを生成できます。

• 前方パスと後方パスを計算します
• 操作を効率的な融合領域に統合します
• 最適化されたカーネルに計算をディスパッチします
• 計算をマシン間で最適に分散します。

そのために、Thunder には以下が同梱されています。

• PyTorch およびカスタム操作を対象とした Python プログラムを取得するための JIT
• 操作を縮小された操作セットのトレースとして表現するマルチレベル中間表現 (IR)
• 計算グラフのトレースに対する拡張可能な変換セット ( grad 、 fusions 、分散型 ( ddp 、 fsdpなど)、関数型 ( vmap 、 vjp 、 jvpなど))
• 拡張可能なエグゼキュータのコレクションにオペレーションをディスパッチする方法

Thunder はすべて Python で書かれています。そのトレースも、変換のすべての段階で有効な Python として表現されます。これにより、前例のないレベルの内省と拡張性が可能になります。

Thunder は、GPU などのアクセラレータ用のコードを直接生成しません。ユーザー プログラムを取得して変換するため、次のような高速エグゼキューターを使用してデバイス コードを最適に選択または生成できます。

• トーチ.コンパイル
• nvFuser
• クドン
• 頂点
• 変圧器エンジン
• PyTorch の熱心な
• PyCUDA、Numba、CuPy を介したカスタム CUDA カーネル
• OpenAI Triton で書かれたカスタム カーネル

Thunder でコンパイルされたモジュールと関数は、標準の PyTorch と完全に相互運用でき、PyTorch の autograd をサポートします。また、Thunder は torch.compile と連携して、最先端の最適化を活用します。

ドキュメント

オンラインドキュメントが利用可能です。ローカルでドキュメントを構築するには、次を使用できます

make docs

そして、ブラウザでdocs/build/index.htmlにある生成されたドキュメントを参照します。

ぜひ参加してください！

フィードバックやご貢献に感謝いたします。機能のリクエストや質問がある場合、またはコードや構成ファイルの提供を希望する場合は、遠慮なく GitHub Issue Tracker を使用してください。

経験やハードウェアのレベルに関係なく、個人の寄稿者を歓迎します。皆さんの貢献は貴重であり、この協力的で協力的な環境で皆さんが何を達成できるか楽しみにしています。

ライセンス

Lightning Thunder は、Apache 2.0 ライセンスに基づいてリリースされています。詳細については、LICENSE ファイルを参照してください。