lightning thunder下載 - lightning thunder源碼下載

lightning thunder

其他源碼

Initial release

下載

讓 PyTorch 模型快如閃電。

Lightning.ai • 效能 • 入門 • 安裝 • 範例 • Thunder 內部 • 參與其中！ • 文件

CI測試一般檢查

歡迎來到⚡閃電雷霆

Thunder 讓 PyTorch 模型快如閃電。

Thunder 是 PyTorch 的來源到來源編譯器。它透過同時組合和使用不同的硬體執行器（例如，nvFuser、torch.compile、cuDNN 和 TransformerEngine FP8）來提高 PyTorch 程式的速度。

它支援單 GPU 和多 GPU 配置。 Thunder 的目標是可用、可理解和可擴展。

筆記

閃電雷霆處於阿爾法階段。請隨意參與，但預計一路上會遇到一些坎坷。

單GPU效能

透過最佳化的複合效應和使用一流的執行器，Thunder 可以比標準的非編譯 PyTorch 程式碼（「PyTorch eager」）實現顯著的加速。下圖顯示了 LitGPT 中實現的 Llama 2 7B 的預訓練吞吐量。

如上圖所示，與使用 nvFuser、torch.compile、cuDNN 和 TransformerEngine FP8 等執行器組合的 H100 上的 eager 程式碼相比，Thunder 的訓練吞吐量提高了 40%。

多GPU效能

Thunder也支援DDP和FSDP等分散式策略，用於在多個GPU上訓練模型。下圖顯示了在沒有 FP8 混合精度的情況下針對 Llama 2 7B 測得的歸一化吞吐量；對 FSDP 的支持正在進行中。

開始使用

開始使用 Thunder 的最簡單方法是使用我們的從零到 Thunder 教學工作室，無需額外的安裝或設定。

安裝迅雷

Thunder 處於 alpha 階段，最新的開發正在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 儲存庫並以 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 還處於早期階段，尚未用於生產作業。

然而，它已經可以為 LitGPT 支援的 LLM 預訓練和微調提供出色的性能，例如 Mistral、Llama 2、Gemma、Falcon 等。

查看 LitGPT 集成，了解如何一起運行 LitGPT 和 Thunder。

Thunder 內部：核心功能簡介

給定一個 Python 可呼叫或 PyTorch 模組，Thunder 可以產生一個最佳化的程式：

計算其前向和後向傳遞
將作業合併到高效率的融合區域
將計算分派給優化的內核
在機器之間最佳地分配計算

為此，Thunder 附帶了：

用於取得針對 PyTorch 和自訂操作的 Python 程式的 JIT
多級中間表示 (IR)，將操作表示為縮減操作集的踪跡
計算圖軌跡上的一組可擴展轉換，例如grad 、 fusions 、分佈式（如ddp 、 fsdp ）、函數（如vmap 、 vjp 、 jvp ）
一種將操作分派給可擴展執行器集合的方法

Thunder 完全用 Python 寫。甚至它的痕跡在轉換的所有階段都被表示為有效的 Python。這使得內省和可擴展性達到前所未有的水平。

Thunder 不會直接為 GPU 等加速器產生程式碼。它會取得並轉換使用者程序，以便可以使用快速執行器以最佳方式選擇或產生設備代碼，例如：