點(diǎn)藍(lán)色字關(guān)注“機(jī)器學(xué)習(xí)算法工程師”

設(shè)為星標(biāo)，干貨直達(dá)！

(資料圖片)

PyTorch 2.0引入了**torch.compile()**來加速模型，這篇文章我們將介紹如何使用**torch.compile()**來加速Hugging Face和TIMM庫(kù)的模型。

torch.compile() 使得嘗試不同的編譯器后端變得容易，從而使用單行裝飾器 torch.compile() 使 PyTorch 代碼更快。它可以直接在 nn.Module 上工作，作為 torch.jit.script() 的直接替代品，但不需要您進(jìn)行任何源代碼更改。我們希望這一行代碼更改能夠?yàn)槟呀?jīng)運(yùn)行的絕大多數(shù)模型提供 30%-2 倍的訓(xùn)練時(shí)間加速。

opt_module=torch.compile(module)

torch.compile 支持任意 PyTorch 代碼、控制流、變異，并帶有對(duì)動(dòng)態(tài)形狀的實(shí)驗(yàn)性支持。我們對(duì)這一發(fā)展感到非常興奮，我們將其稱為 PyTorch 2.0。

這個(gè)版本對(duì)我們來說不同的是，我們已經(jīng)對(duì)一些最流行的開源 PyTorch 模型進(jìn)行了基準(zhǔn)測(cè)試，并獲得了 30% 到 2 倍的大幅加速（見https://github.com/pytorch/torchdynamo/issues/681） 。

這里沒有技巧，我們已經(jīng) pip 安裝了流行的庫(kù)，比如https://github.com/huggingface/transformers, https://github.com/huggingface/accelerate 和 https://github.com/rwightman/pytorch-image-models等流行的庫(kù)，然后對(duì)它們運(yùn)行 torch.compile() 就可以了。

很難同時(shí)獲得性能和便利性，但這就是核心團(tuán)隊(duì)發(fā)現(xiàn) PyTorch 2.0 如此令人興奮的原因。Hugging Face 團(tuán)隊(duì)也很興奮，用他們的話說：

TIMM 的主要維護(hù)者 Ross Wightman：“PT 2.0 開箱即用，適用于推理和訓(xùn)練工作負(fù)載的大多數(shù) timm 模型，無需更改代碼?！?/p>

Sylvain Gugger 是 transformers 和 accelerate 的主要維護(hù)者：“只需添加一行代碼，PyTorch 2.0 就可以在訓(xùn)練 Transformers 模型時(shí)提供 1.5 到 2.x 的加速。這是引入混合精度訓(xùn)練以來最激動(dòng)人心的事情！”

本教程將向您展示如何使用這些加速，這樣您就可以像我們一樣對(duì) PyTorch 2.0 感到興奮。

對(duì)于 GPU（新一代 GPU 的性能會(huì)大大提高）：

pip3installnumpy--pretorch--force-reinstall--extra-index-urlhttps://download.pytorch.org/whl/nightly/cu117

對(duì)于CPU：

pip3install--pretorch--extra-index-urlhttps://download.pytorch.org/whl/nightly/cpu

當(dāng)安裝好后，你可以通過以下方式來進(jìn)行驗(yàn)證：

gitclonehttps://github.com/pytorch/pytorchcdtools/dynamopythonverify_dynamo.py

另外一種安裝方式是采用docker，我們還在 PyTorch nightly 二進(jìn)制文件中提供了所有必需的依賴項(xiàng)，您可以使用它們下載：

dockerpullghcr.io/pytorch/pytorch-nightly

對(duì)于臨時(shí)實(shí)驗(yàn)，只需確保您的容器可以訪問所有 GPU：

dockerrun--gpusall-itghcr.io/pytorch/pytorch-nightly:latest/bin/bash

讓我們從一個(gè)簡(jiǎn)單的例子開始，一步步把事情復(fù)雜化。請(qǐng)注意，您的 GPU 越新，您可能會(huì)看到更顯著的加速。

importtorchdeffn(x,y):a=torch.sin(x).cuda()b=torch.sin(y).cuda()returna+bnew_fn=torch.compile(fn,backend="inductor")input_tensor=torch.randn(10000).to(device="cuda:0")a=new_fn()

這個(gè)例子實(shí)際上不會(huì)運(yùn)行得更快，但它具有教育意義。

以 torch.cos() 和 torch.sin() 為特色的示例，它們是逐點(diǎn)操作的示例，因?yàn)樗鼈冊(cè)谙蛄可现饌€(gè)元素地進(jìn)行操作。你可能真正想要使用的一個(gè)更著名的逐點(diǎn)運(yùn)算是類似 torch.relu() 的東西。eager模式下的逐點(diǎn)操作不是最優(yōu)的，因?yàn)槊總€(gè)操作都需要從內(nèi)存中讀取一個(gè)張量，進(jìn)行一些更改，然后寫回這些更改。

PyTorch 2.0 為您所做的最重要的優(yōu)化是融合。

回到我們的示例，我們可以將 2 次讀取和 2 次寫入變成 1 次讀取和 1 次寫入，這對(duì)于較新的 GPU 來說尤其重要，因?yàn)槠款i是內(nèi)存帶寬（您可以多快地向 GPU 發(fā)送數(shù)據(jù)）而不是計(jì)算（您的速度有多快） GPU 可以處理浮點(diǎn)運(yùn)算）。

PyTorch 2.0 為您做的第二個(gè)最重要的優(yōu)化是 CUDA graphs。CUDA graphs有助于消除從 python 程序啟動(dòng)單個(gè)內(nèi)核的開銷。

torch.compile() 支持許多不同的后端，但我們特別興奮的一個(gè)是生成 Triton 內(nèi)核（https://github.com/openai/triton，用 Python 編寫的，但性能優(yōu)于絕大多數(shù)手寫的 CUDA 內(nèi)核）的 Inductor。假設(shè)我們上面的示例名為 trig.py，我們實(shí)際上可以通過運(yùn)行來檢查代碼生成的 triton 內(nèi)核：

TORCHINDUCTOR_TRACE=1pythontrig.py

@pointwise(size_hints=[16384],filename=__file__,meta={"signature":{0:"*fp32",1:"*fp32",2:"i32"},"device":0,"constants":{},"configs":[instance_descriptor(divisible_by_16=(0,1,2),equal_to_1=())]})@triton.jitdefkernel(in_ptr0,out_ptr0,xnumel,XBLOCK:tl.constexpr):xnumel=10000xoffset=tl.program_id(0)*XBLOCKxindex=xoffset+tl.reshape(tl.arange(0,XBLOCK),[XBLOCK])xmask=xindex
你可以驗(yàn)證融合這兩個(gè) sins 確實(shí)發(fā)生了，因?yàn)檫@兩個(gè) sin 操作發(fā)生在一個(gè)單一的 Triton 內(nèi)核中，并且臨時(shí)變量保存在寄存器中，可以非?？焖俚卦L問。
下一步，讓我們嘗試一個(gè)真實(shí)的模型，比如來自 PyTorch hub 的 resnet50。
importtorchmodel=torch.hub.load("pytorch/vision:v0.10.0","resnet18",pretrained=True)opt_model=torch.compile(model,backend="inductor")model(torch.randn(1,3,64,64))
如果您實(shí)際運(yùn)行，您可能會(huì)驚訝于第一次運(yùn)行很慢，那是因?yàn)檎诰幾g模型。后續(xù)運(yùn)行會(huì)更快，因此在開始對(duì)模型進(jìn)行基準(zhǔn)測(cè)試之前預(yù)熱模型是常見的做法。
您可能已經(jīng)注意到我們?nèi)绾卧诖颂幨褂谩癷nductor”顯式傳遞編譯器的名稱，但它不是唯一可用的后端，您可以在 torch._dynamo.list_backends() 中運(yùn)行以查看可用后端的完整列表。為了好玩，您應(yīng)該嘗試 aot_cudagraphs 或 nvfuser。
現(xiàn)在讓我們做一些更有趣的事情，我們的社區(qū)經(jīng)常使用來自 transformers （https://github.com/huggingface/transformers） 或 TIMM （https://github.com/rwightman/pytorch-image-models）的預(yù)訓(xùn)練模型和我們的設(shè)計(jì)之一PyTorch 2.0 的目標(biāo)是任何新的編譯器堆棧都需要開箱即用，可以與人們實(shí)際運(yùn)行的絕大多數(shù)模型一起工作。因此，我們將直接從 Hugging Face hub 下載預(yù)訓(xùn)練模型并對(duì)其進(jìn)行優(yōu)化。
importtorchfromtransformersimportBertTokenizer,BertModel#Copypastedfromherehttps://huggingface.co/bert-base-uncasedtokenizer=BertTokenizer.from_pretrained("bert-base-uncased")model=BertModel.from_pretrained("bert-base-uncased").to(device="cuda:0")model=torch.compile(model)#Thisistheonlylineofcodethatwechangedtext="Replacemebyanytextyou"dlike."encoded_input=tokenizer(text,return_tensors="pt").to(device="cuda:0")output=model(**encoded_input)
如果您從模型和 encoded_input 中刪除 to(device="cuda:0") ，那么 PyTorch 2.0 將生成 C++ 內(nèi)核，這些內(nèi)核將針對(duì)在您的 CPU 上運(yùn)行進(jìn)行優(yōu)化。你可以檢查 Triton 或 C++ 內(nèi)核的 BERT，它們顯然比我們上面的三角函數(shù)示例更復(fù)雜，但如果你了解 PyTorch，你也可以類似地瀏覽它并理解。
相同的代碼也可以https://github.com/huggingface/accelerate 和 DDP 一起使用。
同樣讓我們嘗試一個(gè) TIMM 示例：
importtimmimporttorchmodel=timm.create_model("resnext101_32x8d",pretrained=True,num_classes=2)opt_model=torch.compile(model,backend="inductor")opt_model(torch.randn(64,3,7,7))
我們使用 PyTorch 的目標(biāo)是構(gòu)建一個(gè)廣度優(yōu)先的編譯器，該編譯器將加速人們?cè)陂_源中運(yùn)行的絕大多數(shù)實(shí)際模型。Hugging Face Hub 最終成為我們非常有價(jià)值的基準(zhǔn)測(cè)試工具，確保我們所做的任何優(yōu)化實(shí)際上都有助于加速人們想要運(yùn)行的模型。
本文翻譯自https://pytorch.org/blog/Accelerating-Hugging-Face-and-TIMM-models/
 

關(guān)鍵詞：         
       絕大多數(shù) 
              
       基準(zhǔn)測(cè)試 
              
       的目標(biāo)是