一種分散式深度學習程式設計新正規化：Global Tensor

語言: CN / TW / HK
時間 2022-07-27 05:15:17

撰文｜姚遲、許嘯宇、左益豪、程國良

Global Tensor 是指多機多裝置執行的 Tensor，它是實現全域性視角（Global View）程式設計的介面。

當前的並行程式，大都採用單程式多資料（SPMD）的方式來程式設計。並行執行同樣的程式，但是處理的是不同資料，以此實現資料的並行處理。以 PyTorch DistributedDataParallel（DDP） 為例，每個程序執行同樣的神經網路計算邏輯，但是每個程序載入資料集的不同分片。

單程式多資料（SPMD）程式設計的缺陷是多資料的通訊很繁瑣。在深度學習場景下，SPMD 程式設計需要在原計算程式碼中插入通訊操作，比如資料並行時對梯度彙總（AllReduce 操作），模型並行時需要 AllGather/ReduceScatter 操作。如果並行模式複雜，或者需要試驗新的並行模式，插入通訊操作就變得難以開發和維護。

全域性視角（Global View）程式設計提供了單程式單資料（SPSD）的程式設計視角。 與 SPMD 程式設計不同的是，Global View的資料是同一個邏輯資料，從程式設計介面層面看是單一資料，其實更簡潔自然。

當我們把一個單程序程式擴充套件到並行執行時，一個單程序資料被擴充套件成多程序資料，多個程序上的這些資料都對應原單程序程式中的同一個邏輯資料。這個邏輯資料在 OneFlow 中叫 Global Tensor。

程式設計時，Global Tensor 讓使用者可以用 SPSD 的介面來程式設計，即按照單機單裝置的邏輯視角來寫程式。然後 OneFlow 框架內部會自動地轉換成物理的 SPMD/MPMD 方式來做並行/分散式執行。

使用 Global Tensor，就可以採用比較自然的 Global View 視角，把多機多裝置看作一個裝置來程式設計，實現 SPSD 程式設計。

1

Global Tensor概述

在程式語言中，Global 的含義通常是程序內的全域性可見，比如全域性變數（Global Variable）。

但 Global Tensor 中 “Global” 的含義是程序間全域性可見，所以 Global Tensor 更為準確的的說法是 Global (on all processes) Tensor，即所有程序可見的 Tensor。

Global Tensor 在每個程序上都存在，在所有程序上被某運算元執行時，就自動完成了對該 Tensor 的多機多裝置執行。

當前常用的 Tensor，只在單個程序內可見，存在於一個裝置上，OneFlow 中把這種 Tensor 叫做 Local Tensor。Local 是相對 Global 而言的，所以 Local Tensor 可以認為是 Local (on one process) Tensor。

一種分散式深度學習程式設計新正規化：Global Tensor

OneFlow 的運算元大部分相容 Local Tensor 和 Global Tensor 的執行。Local Tensor 可以便捷地轉化為 Global Tensor。如此，單機單卡執行的程式碼可以平滑地轉換成多機多卡執行的程式碼。

使用 Global Tensor，可以非常便捷地進行多機多卡的模型開發，相比使用原始通訊運算元，可以成倍提高並行執行模型的開發效率。

2

建立Global Tensor

現在，嘗試在有 2 張 GPU 的主機上建立一個 Global Tensor。以 randn 運算元為例，建立一個 Python 檔案 test_randn_global.py，加入以下內容：

``` import oneflow as flow

Place a global tensor on cuda device of rank(process) 0 and 1

placement = flow.placement(type="cuda", ranks=[0, 1])

Each rank's local data is a part data as a result of spliting global data on dim 0

sbp = flow.sbp.split(dim=0)

Create a global tensor by randn

x = flow.randn(4, 5, placement=placement, sbp=sbp)

Print local data

print("Local data of global tensor:\n ", x.to_local().numpy())

Print global data

print("Global data of global tensor:\n ", x.numpy()) ```

在上述程式碼中有一些新出現的概念：

placement 表示 Global Tensor 分佈的物理裝置，引數 type 指定了物理裝置的型別，這裡使用“cuda” 表示 GPU 裝置，引數 ranks 指定了裝置 ID。對於沒有 2 張 GPU 的使用者，在這裡可以將 type 指定為 "cpu"，這樣可以使用 CPU 模擬多個裝置，下文的程式碼同樣適用。

sbp 表示 Global Tensor 分佈的方式，程式碼中的 sbp = flow.sbp.split(dim=0) 表示把 Global Tensor 在維度 0 均勻切分。

to_local() 可以從 Global Tensor 中獲取它在當前 rank 的 Local Tensor，因為 Global Tensor 在每個 rank 都內含了一個 Local Tensor 作為實際存在的本地分量。

然後配置下多程序啟動依賴的環境變數。這裡是兩卡執行，對應兩個程序啟動，所以需要開啟兩個 Terminal，分別配置如下環境變數：

Terminal 0

export MASTER_ADDR=127.0.0.1 MASTER_PORT=17789 WORLD_SIZE=2 RANK=0 LOCAL_RANK=0

Terminal 1

export MASTER_ADDR=127.0.0.1 MASTER_PORT=17789 WORLD_SIZE=2 RANK=1 LOCAL_RANK=1

以上環境變數的詳細解釋及藉助工具做分散式啟動，請參考\ http://zhuanlan.zhihu.com/p/543441584/edit#_2

最後，在兩個 Terminal 下分別啟動一下test_randn_global.py，觀察 Global Tensor 的建立結果：

python3 test_randn_global.py

這樣在 Terminal 0 即 rank 0 可以看到：

Local data of global tensor: [[-0.07157125 -0.92717147 1.5102768 1.4611115 1.014263 ] [-0.1511031 1.570759 0.9416077 0.6184639 2.4420679 ]] Global data of global tensor: [[-0.07157125 -0.92717147 1.5102768 1.4611115 1.014263 ] [-0.1511031 1.570759 0.9416077 0.6184639 2.4420679 ] [-0.38203463 0.453836 0.9136015 2.35773 -0.3279942 ] [-0.8570119 -0.91476554 -0.06646168 0.50022084 -0.4387695 ]]

在 Terminal 1 即 rank 1 可以看到：

Local data of global tensor: [[-0.38203463 0.453836 0.9136015 2.35773 -0.3279942 ] [-0.8570119 -0.91476554 -0.06646168 0.50022084 -0.4387695 ]] Global data of global tensor: [[-0.07157125 -0.92717147 1.5102768 1.4611115 1.014263 ] [-0.1511031 1.570759 0.9416077 0.6184639 2.4420679 ] [-0.38203463 0.453836 0.9136015 2.35773 -0.3279942 ] [-0.8570119 -0.91476554 -0.06646168 0.50022084 -0.4387695 ]]

可以發現兩個 rank 的 Local Tensor 在維度 0 拼接後，就是完整的 Global Tensor 的值。

3

由Local Tensor得到Global Tensor

可以先建立 Local Tensor，再利用 Tensor.to_global （http://oneflow.readthedocs.io/en/master/tensor.html#oneflow.Tensor.to_global ）方法，將 Local Tensor 轉為 Global Tensor。

建立如下程式，採用上文同樣的方式啟動：

``` import oneflow as flow

x = flow.randn(2, 5).cuda() print(x.is_local) # True print(x.is_global) # False placement = flow.placement(type="cuda", ranks=[0, 1]) sbp = flow.sbp.split(0) x_global = x.to_global(placement=placement, sbp=sbp) print(x_global.shape) # (4, 5) print(x.is_local) # True print(x_global.is_global) # True ```

該程式在 2 個 GPU 裝置上分別建立了 shape=(2,5) 的 Local Tensor，即 x。

然後定義 placement 為 rank 0 和 1 上的 CUDA 裝置，SBP 為 tensor 第 0 維的切分，原本 Local Tensor 經過 to_global 變換後，就得到一個名為 x_global 的 Global Tensor。

可以觀察到，x_global 的 shape 變為了 (4, 5)，這是 Global Tensor 的 shape（global shape）。

Global Tensor 與 Local Tensor 之間為總量與分量的關係。Local Tensor 是總量在本 rank 的分量。分量和總量的具體關係由 Placement 和 SBP 確定，比如這裡的關係是在 0 和 1 號 GPU 上，x_global 在第 0 維 split 而得到 x。

to_global 方法根據如上關係可以從 x.shape 推理出 x_global.shape ：把兩個 GPU 上的 Local Tensor x 在第 0 維拼接後得到 x_global。

Global Tensor 除了 shape，還有資料部分。一個 Global Tensor 的內部，在每個 rank 上都內含了一個 Local Tensor 作為其本地分量，這個 Local Tensor 就是 Global Tensor 在每個 rank 的物理資料。這符合期待，每個 rank 只需儲存一部分物理資料。

4

由Global Tensor得到Local Tensor

如果想得到 Global Tensor 的本地分量，可以通過 to_local 方法得到。例如：

``` import oneflow as flow

placement = flow.placement(type="cuda", ranks=[0, 1]) sbp = flow.sbp.split(0) x = flow.randn(4, 5, placement=placement, sbp=sbp) print(x.to_local()) ```

當執行 x.to_local() 時，兩個不同的 rank 將分別得到一個 shape 為 (2, 5) 的本地分量 tensor。

在 Terminal 0 即 rank 0 可以看到：

tensor([[-0.2730, 1.8042, 0.0721, -0.5024, -1.2583], [-0.3379, 0.9371, 0.7981, -0.5447, -0.5629]], dtype=oneflow.float32)

在 Terminal 1 即 rank 1 可以看到：

tensor([[ 0.6829, 0.4849, 2.1611, 1.4059, 0.0934], [-0.0301, -0.6942, -0.8094, -1.3050, -0.1778]], dtype=oneflow.float32)

to_local() 沒有任何引數，因為 Global Tensor 已經通過 placement 和 SBP 確定好了它的本地分量，所以直接取本地分量對應的 Local Tensor 就好。

5

由Global Tensor轉成另一個Global Tensor

分散式計算通常都需要在正常的計算邏輯之間插入通訊操作，而使用 OneFlow 時只需要做 Global Tensor 的資料分佈型別轉換。

相比普通的 Local Tensor，從型別上講，Global Tensor 的最大區別是帶有全域性資料分佈型別（Global Data Distribution Type）。 全域性資料分佈型別指定了 Global Tensor 在每個程序（Rank）的資料分佈情況，由 Placement 和 SBP 組成。

全域性資料分佈型別中的 Placement 指定了資料分佈的裝置集合:

  • 引數 type 指定了物理裝置的型別，cuda 表示 GPU 裝置記憶體, cpu 表示 CPU 裝置記憶體；
  • 引數 ranks 指定了程序 ID 集合，因為隱含了一個 Rank 對應一個物理裝置，所以 ranks 就是裝置 ID 集合; 實際上 ranks 是一個由 rank id 組成 nd-array，支援高維裝置排布。

詳情請參考 oneflow.placement（http://oneflow.readthedocs.io/en/master/tensor_attributes.html?highlight=placement#oneflow.placement ）。

全域性資料分佈型別中的 SBP 指定了全域性資料和區域性資料的關係:

  • S 即 split(dim)，區域性和全域性是切分關係， 表示在 dim 維度做了切分的資料分佈關係；
  • B 即 broadcast，區域性和全域性是廣播關係，表示做了廣播的資料分佈關係；
  • P 即 partial_sum，區域性和全域性是部分關係，表示做了 element-wise 累加的資料分佈關係。

詳情請參考 oneflow.sbp.sbp（http://oneflow.readthedocs.io/en/master/tensor_attributes.html?highlight=placement#oneflow.sbp.sbp ）。

資料重分佈（Re-distribution）是平行計算中經常要處理的，即變換資料分佈，比如把分片資料聚合到一起。在 MPI 程式設計正規化（SPMD）下，資料重分佈需要寫顯式的通訊操作，如 AllReduce、AllGather、ReduceScatter。在 OneFlow 的 Global View 程式設計正規化（SPSD) 下，資料重分佈可以通過 Global Tensor 的全域性資料分佈型別轉換完成。

全域性資料分佈型別的轉換類似常規程式語言中的（顯式）型別轉換。型別轉換時，只需指定要變換到的型別，裡面隱含的操作會被系統自動完成。比如 double 型別到 int 型別的轉換，去掉小數點部分的操作就是系統自動完成的。

同樣，只需指定 Global Tensor 要轉換的新全域性資料分佈型別，裡面隱含的通訊操作會被 OneFlow 自動完成。全域性資料分佈型別轉換的介面是 Tensor.to_global，to_global 有 placement 和 sbp 兩個引數，這兩個引數即期望轉換成的新全域性資料分佈型別。

全域性資料分佈型別轉換中隱含的主要操作是通訊的推理和執行，背後的實現機制是 OneFlow 的 Boxing，這是一種自動做資料 Re-distribution 的機制。

下面看一個例子，該例子可以把一個按 split 分佈的 Global Tensor 轉換為一個按 broadcast 分佈的 Global Tensor：

``` import oneflow as flow

x = flow.randn(2, 5).cuda() placement = flow.placement(type="cuda", ranks=[0, 1]) sbp = flow.sbp.split(0) x_global = x.to_global(placement=placement, sbp=sbp) print(x_global.shape) # (4, 5) print(x_global.to_local()) sbp_b = flow.sbp.broadcast x_global_b = x_global.to_global(placement=placement, sbp=sbp_b) print(x_global_b.shape) # (4, 5) print(x_global_b.to_local()) ```

可以看到，x_global 到 x_global_b 的全域性資料分佈型別變化就是 sbp 從 flow.sbp.split(0) 變成了 flow.sbp.broadcast。它們的 global shape 都是 (4, 5)，但是本地分量從一個分片變成了一個完整的資料，這個變化可以從對 to_local() 的列印結果觀察到。

這裡的 to_global 變換完成了對 local tensor 的歸併。通常來講，SPMD 程式設計模式要求使用者手寫一個 all-gather 集合通訊來完成。而在 OneFlow Global View 中，只需做一下型別轉換。

通過 Global Tensor 的型別轉換，就自動完成通訊操作的推理和執行。讓演算法開發者可以思考資料的分佈（Thinking in data distribution） ，而不是思考如何通訊（Thinking in data communication operation） ，實現了所想即所得，從而提高分散式程式的開發效率。

這裡補充介紹一下 Global Tensor 的 numpy() 方法。對於任意的 Global Tensor 如 x_global，x_global.numpy() 等價於 x_global.to_global(spb=flow.sbp.broadcast).to_local().numpy()，即內部隱含了一次將原 Global Tensor 轉成 SBP 為 flow.sbp.broadcast() 的 Global Tensor，然後進行一次 to_local 操作，最後對這個 Local Tensor 呼叫 numpy() 方法。所以 x_global.numpy() 得到的是一個完整的資料。

6

Global Tensor參與計算

這一節介紹 Global Tensor 如何參與實際計算。以 Global Tensor 參與矩陣乘法計算為例，構造如下程式：

``` import oneflow as flow

placement = flow.placement(type="cuda", ranks=[0, 1]) x = flow.randn(4, 5, placement=placement, sbp=flow.sbp.split(dim=0)) w = flow.randn(5, 8, placement=placement, sbp=flow.sbp.broadcast) y = flow.matmul(x, w) print(y.is_global) # True print(y.shape) # (4, 8) print(y.sbp) # (flow.sbp.split(dim=0)) print(y.to_local().numpy()) ```

以上程式建立了兩個 Global Tensor，分別是 x 和 w，它們參與 oneflow.matmul 計算得到 y。

OneFlow 中的大部分運算元都支援計算 Global Tensor。flow.matmul 執行 Global Tensor 時，在介面上並無特殊之處。可以認為 OneFlow 中的運算元都是多型的。即根據輸入，決定自己的行為：

  • 如果運算元的輸入是 Local Tensor，那麼運算元會按照普通的單機單裝置執行模式進行計算；
  • 如果運算元的輸入是 Global Tensor，那麼運算元會採用 Global View（多機多裝置）模式進行計算；

當用戶需要將單卡程式碼改為分散式程式碼時，運算元支援多型執行為使用者提供了極大的便利：只需要把輸入的 (Local) Tensor 轉換成 Global Tensor 。

類似於單裝置執行時要求輸入資料所在裝置相同，以上程式中， flow.matmul 這一運算元可以順利執行的前置條件是：輸入的 x 和 w 的 placement 相同。

程式中矩陣相乘的結果 y 同樣是一個 Global Tensor 。flow.matmul 對輸入 x 和 w 做計算時，會自動進行輸出資料的 placement 和 SBP 的推理，規則如下：

  • Placement：輸出和輸入的 placement 相同；
  • SBP：輸出的 SBP 的推理規則，因運算元型別而異，這個推理規則是 OneFlow 內建的，詳情可見: SBP Signature

此處，flow.sbp.split(0) 和 flow.sbp.broadcast 相乘的輸出資料會被推理成 flow.sbp.split(0)。x 在每個 rank 上是一個分片資料，w 是一個完整的資料，二者矩陣乘法得到的 y 是一個分片的資料。看到這裡，瞭解常見並行執行方式的使用者可以發現：這裡實現了一個數據並行的前向計算，x 是切片的資料，w 是完整的引數。

7

結語

上文介紹了：

  • Global View 提供的 SPSD 程式設計視角；
  • Global Tensor 的跨程序可見的執行特點；
  • Global Tensor 和 Local Tensor 的互轉；
  • 通過 Global Tensor 的全域性資料分佈型別轉換來實現分散式通訊；
  • OneFlow 運算元的多型特性支援了 Global Tensor 的執行。

至此，本文從 Global Tensor 的建立開始，最終完成了一個基於 Global Tensor 的資料平行計算流程。更多並行方式和 SBP 的推理邏輯，將在後續內容介紹。

擴充套件閱讀：OneFlow 多機多卡啟動和依賴的環境變數

OneFlow 的 Global Tensor 執行採用的是 多客戶端模式 (Multi-Client) ，每個裝置對應一個程序。n 機 m 卡 的環境，就對應 n * m 個程序。每個程序都有一個程序 rank 編號，Global Tensor 中的 placement 引數中的 ranks 對應的就是這個 rank 編號。

以 2 機 2 卡 為例， 0 號機器中兩張卡分別對應編號 0 和 1，第 1 號機器中兩張卡分別對應編號 2 和 3。此時 flow.placement(type="cuda", ranks=[2]) 可以唯一標識第 1 號機器中的第 0 卡。

一般地，對於 n 機 m 卡 的環境，flow.placement(type="cuda", ranks=[k]) 唯一標識第 k / n 號機器的第 k % m 張卡。

因為採用多客戶端模式，所以需要對應每個裝置都啟動一個程序。在 OneFlow 中，所有程序都只需要啟動相同的指令碼程式。不同程序之間通過不同的環境變數來區分程序編號和建立通訊連線。

環境變數說明：

  • MASTER_ADDR：多機訓練的第 0 號機器的 IP；
  • MASTER_PORT：多機訓練的第 0 號機器的監聽埠，不與已經佔用的埠衝突即可；
  • WORLD_SIZE：整個叢集中計算裝置的數目，因為目前還不支援各個機器上顯示卡數目不一致，因此 WORLD_SIZE 的數目實際上是 $ 機器數目 \times 每臺機器上的顯示卡數目$。建立 Global Tensor 中的示例是單機 2 卡的情況，因此 WORLD_SIZE=2；
  • RANK：叢集內所有機器下的程序編號；
  • LOCAL_RANK：單個機器內的程序編號。

RANK 和 LOCAL_RANK 的區別：

  • 在單機訓練（單機單卡或單機多卡）時，兩者相等；
  • 在多機訓練時，每臺機器上的 LOCAL_RANK 的上限，就是每臺機器上的計算裝置的數目；RANK 的上限，就是所有機器上所有計算裝置的總和，它們的編號均從 0 開始（因為編號從 0 開始，所以不包含上限）。

以 2 機 2 卡 為例，每張顯示卡的 LOCAL_RANK 與 RANK 對應情況如下：

一種分散式深度學習程式設計新正規化：Global Tensor

使用環境變數啟動雖然繁瑣，但是適用性廣，可以採用任意的方式來啟動程序。

另外為了方便使用，OneFlow 也提供了一個分散式啟動多程序且自動構建環境變數的工具 oneflow.distributed.launch（http://docs.oneflow.org/master/parallelism/04_launch.html ）。

（原文：http://docs.oneflow.org/master/cookies/global_tensor.html）

歡迎下載體驗 OneFlow v0.8.0 最新版本：\ http://github.com/Oneflow-Inc/oneflow/

「」