PyTorch 1.8 發布，支持 AMD，優化大規模訓練

內容導讀

北京時間 3 月 4 日，PyTorch 官方博客發布 1.8
版本，據官方介紹，新版本主要包括編譯器和分布式訓練更新，同時新增了部分移動端教程，

整體來看，本次版本更新涵蓋 1.7 版本發布以來，共計 3,000 多次 commit，包括編譯、代碼優化、科學計算前端 API 以及通過 pytorch.org 提供的二進制檔案支持 AMD ROCm，

同時 PyTorch 1.8 還為管道和模型并行的大規模訓練，進行了功能改進和梯度壓縮，

其中一些重點更新包括：

1. 通過 torch.fx ，可進行 python 之間的函式轉換；

2. 增加/調整 API，以支持 FFT( torch.fft )、線性代數函式(torch.linalg)，支持復雜張量自動求導（autograd），并提升了計算 hessian 和 jacobian 的性能表現；

3. 對分布式訓練進行了重大更新和改進，包括：改進 NCCL 可靠性，支持管道并行，RPC 分析，支持添加梯度壓縮的通訊 hooks，

在 PyTorch 1.8 中，官方對一些 PyTorch 庫也進行了重大更新，包括 TorchCSPRNG、TorchVision、TorchText 和 TorchAudio，

新增及更新 API

新增及更新 API 包括：與 NumPy 兼容的額外 API，及在推理和訓練時方面，提高代碼性能的額外 API，

以下為 PyTorch 1.8 主要更新功能的簡介，

[穩定版] Torch.fft 將支持高性能 NumPy 中的 FFT

PyTorch 1.8 中發布了 torch.fft 模塊，該模塊在實作 NumPy np.ft 功能的同時，還支持硬體加速和 autograd，

[測驗版] torch.linalg 將支持 NumPy 中的線性代數函

torch.linalg 以 NumPy np.linalg 為原型，為常見的線性代數運算提供與 NumPy 類似的支持，包括 Cholesky 分解、行列式、特征值等，

[測驗版] 利用 FX 進行 Pthon 代碼轉換

FX 使得開發者可以通過 transform(input_module : nn.Module) -> nn.Module 進行 Python 代碼轉換，輸入 Module 實體，得到轉換后的 Module 實體，

分布式訓練

為了提高 NCCL 穩定性，PyTorch 1.8 將支持穩定的異步錯誤/超時處理；支持 RPC 分析，此外，還增加了對管道并行的支持，并可以通過 DDP 中的通訊鉤子進行梯度壓縮，

詳情如下：

[測驗版] 管道并行

提供一個易于使用的 PyTorch API，可將管道并行作為訓練回圈的一部分，

[測驗版] DDP 通訊鉤子

DDP 通訊鉤子(communication hook)是一個通用介面，通過在 DistributedDataParallel 中覆寫 vanilla allreduce 來控制 worker 之間的梯度通訊，

PyTorch 1.8 新增部分內置的通訊鉤子，如 PowerSGD，用戶可以按需呼叫，此外，通訊鉤子介面還支持用戶自定義通訊策略，

分布式訓練的附加原型功能

除了在穩定版和測驗版中新增的分布式訓練功能外，Nightly 版本中也相應增加了部分功能，

具體如下：

[Prototype] ZeroRedundancyOptimizer
[Prototype] 行程組 NCCL 發送/接收
[Prototype] RPC 中用 TensorPipe 支持 CUDA
[Prototype] 遠程模塊

PyTorch 移動端

PyTorch 1.8 中為新用戶提供了多個移動端教程，旨在幫助新用戶更迅速地將 PyTorch 模型部署在移動端，

同時為老用戶提供開發工具，讓其更得心應手地用 PyTorch 進行移動端開發，

PyTorch 移動端新增教程包括：

• iOS 端用 DeepLabV3 進行影像分割

• Android 端用 DeepLabV3 進行影像分割

iOS 端用 DeepLabV3 進行影像分割

新增 demo APP 還包括影像分割、目標檢測、機器翻譯、智能問答等(iOS & Android)，

除了在 CPU 上對 MobileNetV3 等模型進行性能提升外，官方還對 Android GPU 后臺原型進行了改造，以覆寫更多機型、進行更快速的推理，

另外，PyTorch 1.8 還推出了 PyTorch Mobile Lite 解釋器功能，允許用戶減少運行時二進制檔案的大小，

[Prototype] PyTorch Mobile Lite Interpreter

PyTorch Lite Interpreter 是 PyTorch runtime 的精簡版，可在空間受限的設備中執行 PyTorch 程式，并減少二進制檔案對存盤空間的占用，

與當前版本中的設備運行時相比，這一功能可減少 70% 的二進制檔案大小，

性能優化

PyTorch 1.8 中新增對 benchmark utils 的支持，使用戶能夠更輕松地監控模型性能，另外還新開放了一個自動量化 API，

[測驗版] Benchmark utils

Benchmark utils 允許用戶進行精確的性能測量，并提供組合工具，幫助制定基準和進行后期處理，

代碼示例：

from torch.utils.benchmark import Timer

results = []
for num_threads in [1, 2, 4]:
    timer = Timer(
        stmt="torch.add(x, y, out=out)",
        setup="""
            n = 1024
            x = torch.ones((n, n))
            y = torch.ones((n, 1))
            out = torch.empty((n, n))
        """,
        num_threads=num_threads,
    )
    results.append(timer.blocked_autorange(min_run_time=5))
    print(
        f"{num_threads} thread{'s' if num_threads > 1 else ' ':<4}"
        f"{results[-1].median * 1e6:>4.0f} us   " +
        (f"({results[0].median / results[-1].median:.1f}x)" if num_threads > 1 else '')
    )

1 thread     376 us   
2 threads    189 us   (2.0x)
4 threads     99 us   (3.8x)

[Prototype] FX Graph Mode Quantization

FX Graph Mode Quantization 是 PyTorch 中新增的自動量化 API，它通過增加函式支持和自動化量化程序，改進 Eager Mode Quantization，

硬體支持

[測驗版]強化 PyTorch Dispatcher 的能力，改善 C++ 中后端開發體驗

PyTorch 1.8 支持用戶在 pytorch/pytorch repo 之外創建新的樹外(out-of-tree)設備，并與本地 PyTorch 設備保持同步，

[測驗版] AMD GPU 二進制檔案現已推出

PyTorch 1.8 新增對 ROCm wheel 的支持，用戶只需根據標準 PyTorch 安裝選擇器，安裝選項選擇 ROCm，然后執行命令，即可輕松上手使用 AMD GPU，

教程及檔案詳情，請訪問以下鏈接查看：

https://pytorch.org/blog/pytorch-1.8-released/