詳解PyTorch在MPS后端的新特性

我們推出了一個(gè)新的系列，對(duì)PytorchConference2023 的博客進(jìn)行中文編譯，會(huì)陸續(xù)在公眾號(hào)發(fā)表。

大家好，我叫Kulinseth，我在蘋果的MPS團(tuán)隊(duì)工作，今天我將討論P(yáng)yTorch中MPS后端的改進(jìn)。接下來，我將介紹MPS后端進(jìn)入Beta Stage的新功能。我們添加了一些新功能，如支持分析器、自定義內(nèi)核和MPS開發(fā)者API，這些都是MPS后端的新特性。

Beta Stage
New features:
-Profiler
-Custom Kernel
-Events & MemoryAPI
Performance

之后我們還會(huì)介紹自去年發(fā)布以來MPS后端的一些性能改進(jìn)?，F(xiàn)在，讓我們從beta stage開始?；仡櫼幌?，MPS后端是在去年的PyTorch 1.12中開始的旅程，當(dāng)時(shí)我們?cè)贛ac平臺(tái)上推出了支持GPU加速的PyTorch。我們對(duì)內(nèi)存使用和新張量進(jìn)行了多次優(yōu)化。在PyTorch 2.0中，MPS backend對(duì)于beta stage來說已經(jīng)合格。這意味著我們支持了前60個(gè)最常用的運(yùn)算符，測(cè)試覆蓋面大大提高；隨著多個(gè)常用模型采用了這個(gè)后端作為macOS的默認(rèn)后端，network覆蓋面也得到了擴(kuò)展。但這些并不是我們所做的全部改進(jìn)。

在最新的PyTorch構(gòu)建中，支持了一些新功能，并且我們?cè)诔掷m(xù)不斷地進(jìn)行改進(jìn)，例如支持分析功能、自定義內(nèi)核以及一些開發(fā)者API。

Community engagement:
index_fill /histogram/ copysign/log_sigmoid / xlogy/ pixel_shuffle / hypot/
fmax / fmin / roll / hardsigmoid / logit / nansum / remainder/group_norm/mean_var/median/
repeat_interleave/cumsum/signbit/nansum/frac/masked_select

開發(fā)者們不僅在extend網(wǎng)絡(luò)中采用了PyTorch MPS后端，還貢獻(xiàn)了代碼，將許多新的操作符添加到我們的代碼庫中，例如group_norm、histogram、pixel_shuffle等等。

os signposts
- Operation executions
- Copies between CPU and GPU
- Fallbacks to the CPU
Metal System Trace
Command line tool

現(xiàn)在讓我們來討論一些添加到MPS后端的新功能。首先是profiler支持，這是通過使用IOS中的OS signposts功能實(shí)現(xiàn)的。它可以突出顯示在MPS后端上執(zhí)行的操作，以及在CPU和GPU之間切換的情況以及一些回退到CPU的操作。要使用profiler，我們有一個(gè)簡(jiǎn)單的示例，我會(huì)進(jìn)行演示并介紹一些需要啟用該功能的API。它已經(jīng)整合到了 Metal System Trace 中，并且還有一個(gè)命令行工具供開發(fā)者使用。

importtorch
fromtorchimportnn

model=nn.Sequential(
nn.Linear(784,256),
nn.Softshrink(),
nn.Linear(256,256),
nn.Softshrink(),
nn.Linear(256,256),
nn.Softshrink(),
nn.Linear(256,10)
).to("mps")

torch.mps.profiler.start(mode="interval",wait_until_completed=True)
#Yourmodelcodegoeshere
torch.mps.profiler.stop()

現(xiàn)在讓我們來看一個(gè)使用Linear和Softshrink的Sequential模型組成的簡(jiǎn)單樣本網(wǎng)絡(luò)。這只是一個(gè)簡(jiǎn)單的例子。你可以直接在PyTorch中將其實(shí)現(xiàn)，但我將使用它來說明我們可以如何做。我們可以使用MPS分析工具中啟用的開始和停止API，并采用不同的模式來捕獲標(biāo)識(shí)信息。

結(jié)果是，您可以獲得一個(gè)使用所有系統(tǒng)標(biāo)識(shí)信息的系統(tǒng)跟蹤，可以使用一個(gè)稱為Metal System Trace的工具進(jìn)行可視化。它除了包含其他大量的信息之外，還包括我們作為PyTorch的一部分啟用的標(biāo)識(shí)，以及在時(shí)間線上顯示的其他內(nèi)容。

在這里它突出顯示了Blitcall，您可以看到回退到CPU的情況，以及實(shí)際在MPS上執(zhí)行的操作。這使得您可以開始檢查您的網(wǎng)絡(luò)。正如您在這里所看到的，Softshrink在我們捕獲的時(shí)候，正回退到CPU。

此外，對(duì)于希望快速查看應(yīng)用程序花費(fèi)最多時(shí)間的操作的開發(fā)人員，我們還提供了一個(gè)命令行工具的功能。如圖所示，通過使用環(huán)境變量，您可以輸出有關(guān)每個(gè)層的信息，例如數(shù)據(jù)類型。并且它允許您快速審查他們的應(yīng)用程序。現(xiàn)在，繼續(xù)我們之前的示例，我們看到Softshrink操作在回退到CPU，這在GPU時(shí)間線上留下了一個(gè)很大的間隙。為了改善性能，其中一種方法是添加一些自定義內(nèi)核支持。

編寫自定義操作有三個(gè)步驟。首先在Object2C中實(shí)現(xiàn)操作以便在metal中查看。然后創(chuàng)建Python綁定并構(gòu)建您的擴(kuò)展。在構(gòu)建擴(kuò)展之后，您可以將該操作導(dǎo)入到您的應(yīng)用程序中并開始使用它。所以讓我們從操作實(shí)現(xiàn)開始。代碼很多，但我會(huì)從頭開始解釋。

#include 

torch::Tensor mps_softshrink(const torch::Tensor& input, float lambda = 0.5) {
    // Get a reference of the MPSStreamMTLCommandBuffer and dispatch_queue_t
    id commandBuffer = torch::get_command_buffer();
    dispatch_queue_t serialQueue = torch::get_dispatch_queue();
    
    dispatch_sync(serialQueue, ^{
        // Create the encoder
        id computeEncoder = [commandBuffer computeCommandEncoder];
        
        // Encode the pipeline state object and its parameters
        [computeEncoder setComputePipelineState:softShrinkPsO];
        
        torch::synchronize();
    });
}

首先導(dǎo)入torch擴(kuò)展頭文件，這其中包含撰寫C++擴(kuò)展所需的所有PyTorch部分。這里有一些我們已經(jīng)公開的API，以實(shí)現(xiàn)自定義功能。這個(gè)"get command buffer MPS backend API"是用來獲得對(duì)MPS流命令緩沖區(qū)的引用的。這個(gè)命令緩沖區(qū)與我們?cè)诤蠖擞脕砭幋a工作的命令緩沖區(qū)是相同的。您所做的工作與我們正在進(jìn)行的工作是相同的。它的優(yōu)先級(jí)很高，這使得您可以使用像"commit and continue"這樣的優(yōu)化來減少CPU方面的開銷，這個(gè)在去年的演講中討論過。我們有這個(gè)"getDispatchQueue API"來獲取對(duì)串行隊(duì)列的引用。使用獲取到的命令緩沖區(qū)創(chuàng)建一個(gè)編碼器，它允許您定義自定義GPU內(nèi)核。您使用調(diào)度隊(duì)列來對(duì)內(nèi)核進(jìn)行編碼，以確保來自多個(gè)線程的提交被序列化。在編碼完成所有工作后，使用"synchronize API"直到命令緩沖區(qū)完成。或者，如果您不需要序列化，可以使用"commit API" torch::commit。這允許您在內(nèi)部繼續(xù)進(jìn)行操作。

#include 

torch::Tensor mps_softshrink(const torch::Tensor& input, float lambda = 0.5) {
    // Function implementation goes here
    // ...
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("mps_softshrink", &mps_softshrink, "Apply MPS Softshrink operation");
}

// Compiling the extension
import torch.utils.cpp_extension

compiled_lib = torch.utils.cpp_extension.load(
    name='CustomSoftshrink',
    sources=['CustomSoftshrink.mm'],
    extra_cflags=['-std=c++17']
)

在自定義內(nèi)核支持的第二步中，您可以使用"Pybind11"來綁定Objective-C。以類似的方式將函數(shù)集成到Python中。通過使用CPP擴(kuò)展，您可以構(gòu)建自定義軟共享庫，該庫可以包含在您的應(yīng)用程序中。

from my_build import compiled_lib
from torch import nn

class MPSSoftshrink(nn.Module):
    def __init__(self, lambda_=0.5):
        super(MPSSoftshrink, self).__init__()
        self.lambda_ = lambda_

    def forward(self, input):
        return compiled_lib.mps_softshrink(input, self.lambda_)

model = nn.Sequential(
    nn.Linear(784, 256),
    MPSSoftshrink(),
    nn.Linear(256, 256),
    MPSSoftshrink(),
    nn.Linear(256, 256),
    MPSSoftshrink(),
    nn.Linear(256, 10)
).to("mps")

最后一步，自定義構(gòu)建庫已經(jīng)準(zhǔn)備好在您的應(yīng)用程序中使用。我們已經(jīng)取代了之前速度較慢且降級(jí)到CPU的Softshrink。這是您定制的MPS收縮庫?，F(xiàn)在，在新增的自定義內(nèi)核支持下效率更高。所有通過回退到CPU創(chuàng)建的副本和中間張量都已經(jīng)消失，模型運(yùn)行速度更快。

import torch.mps

# 創(chuàng)建開始事件并記錄
start_event = torch.mps.Event(enable_timing=True)
start_event.record()

# 在GPU上進(jìn)行一些訓(xùn)練操作
# ...

# 創(chuàng)建結(jié)束事件并記錄
end_event = torch.mps.Event(enable_timing=True)
end_event.record()

# 計(jì)算持續(xù)時(shí)間
duration = start_event.elapsed_time(end_event)

# 設(shè)置內(nèi)存分配的比例，限制進(jìn)程在 MPS 設(shè)備上的內(nèi)存分配
torch.mps.set_per_process_memory_fraction(0)

還有一些附加的API，可以在記錄、等待和流逝時(shí)間等事件上進(jìn)行事件管理和創(chuàng)建自定義計(jì)時(shí)操作。對(duì)于MPS分配器的API，如設(shè)置每個(gè)進(jìn)程的內(nèi)存分?jǐn)?shù)，使開發(fā)人員能夠更加細(xì)粒度地控制后端內(nèi)存操作。最后，總結(jié)一下這次演講。讓我們來看一些性能結(jié)果。如您所見，MPS后端已經(jīng)得到了顯著優(yōu)化。