詳解PyTorch在MPS后端的新特性

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

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

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

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

在最新的PyTorch構建中，支持了一些新功能，并且我們在持續不斷地進行改進，例如支持分析功能、自定義內核以及一些開發者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

開發者們不僅在extend網絡中采用了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

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

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()

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

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

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

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

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

#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();
    });
}

首先導入torch擴展頭文件，這其中包含撰寫C++擴展所需的所有PyTorch部分。這里有一些我們已經公開的API，以實現自定義功能。這個"get command buffer MPS backend API"是用來獲得對MPS流命令緩沖區的引用的。這個命令緩沖區與我們在后端用來編碼工作的命令緩沖區是相同的。您所做的工作與我們正在進行的工作是相同的。它的優先級很高，這使得您可以使用像"commit and continue"這樣的優化來減少CPU方面的開銷，這個在去年的演講中討論過。我們有這個"getDispatchQueue API"來獲取對串行隊列的引用。使用獲取到的命令緩沖區創建一個編碼器，它允許您定義自定義GPU內核。您使用調度隊列來對內核進行編碼，以確保來自多個線程的提交被序列化。在編碼完成所有工作后，使用"synchronize API"直到命令緩沖區完成?；蛘?，如果您不需要序列化，可以使用"commit API" torch::commit。這允許您在內部繼續進行操作。

#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']
)

在自定義內核支持的第二步中，您可以使用"Pybind11"來綁定Objective-C。以類似的方式將函數集成到Python中。通過使用CPP擴展，您可以構建自定義軟共享庫，該庫可以包含在您的應用程序中。

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")

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

import torch.mps

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

# 在GPU上進行一些訓練操作
# ...

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

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

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

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