我們推出了一個新的系列，對PytorchConference2023 的博客進行中文編譯，會陸續(xù)在公眾號發(fā)表。

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

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

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

在最新的PyTorch構(gòu)建中，支持了一些新功能，并且我們在持續(xù)不斷地進行改進，例如支持分析功能、自定義內(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后端，還貢獻了代碼，將許多新的操作符添加到我們的代碼庫中，例如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功能實現(xiàn)的。它可以突出顯示在MPS后端上執(zhí)行的操作，以及在CPU和GPU之間切換的情況以及一些回退到CPU的操作。要使用profiler，我們有一個簡單的示例，我會進行演示并介紹一些需要啟用該功能的API。它已經(jīng)整合到了 Metal System Trace 中，并且還有一個命令行工具供開發(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)在讓我們來看一個使用Linear和Softshrink的Sequential模型組成的簡單樣本網(wǎng)絡(luò)。這只是一個簡單的例子。你可以直接在PyTorch中將其實現(xiàn)，但我將使用它來說明我們可以如何做。我們可以使用MPS分析工具中啟用的開始和停止API，并采用不同的模式來捕獲標識信息。

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

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

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

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

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

#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擴展，您可以構(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)準備好在您的應(yīng)用程序中使用。我們已經(jīng)取代了之前速度較慢且降級到CPU的Softshrink。這是您定制的MPS收縮庫。現(xiàn)在，在新增的自定義內(nèi)核支持下效率更高。所有通過回退到CPU創(chuàng)建的副本和中間張量都已經(jīng)消失，模型運行速度更快。

import torch.mps

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

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

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

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

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

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