大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

語言: CN / TW / HK

來源|機器之心

大模型多了去了,告訴我怎麼加速?自 2018 年 BERT 誕生,到 GPT-3、ViT 等擁有數以億計的引數規模的模型不斷湧現,AI 模型引數量的爆發式增長已不足為奇,讓煉丹師無暇顧及甚至感到麻木。

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

與此同時,大模型對計算和記憶體資源提出了巨大的挑戰。訓練成本急劇上升,比如用一塊非常先進的 NVIDIA A100 GPU 訓練千億引數模型的 GPT-3,需要用時 100 多年。

大模型對視訊記憶體的需求增長遠超過 GPU 視訊記憶體增長的速度,根據 OpenAI 的報告,模型大小的增長速度是每 3.5 月翻一倍,而 GPU 視訊記憶體則需要 18 個月才能翻倍。受限於 GPU 視訊記憶體,單個 GPU 無法再容納大規模模型引數。

因此,業內不得不將計算擴充套件到多個 GPU 裝置上,分散式訓練則成為廣大開發者的必然選擇。

但是分散式訓練的門檻太高,即便擁有充足的計算資源,也可能因為搞不定分散式訓練而望洋興嘆。分散式並行程式設計通常需要工程師瞭解計算機系統和架構的專業知識,再加上相關的實踐經驗,這進一步增加了探索前沿演算法和新模型的難度。凡此種種,使得大模型成為部分科技巨頭的特權。如何加速模型訓練效率,讓更多工程師可以使用、研究大模型成為當務之急。

問題是,市面上那麼多支援分散式訓練的模型庫,選哪個最合適?

近期,由一流科技團隊研發的以高效性起家的國產開源深度學習框架 OneFlow 上線了 LiBai(李白)模型庫,這個新生代模型庫覆蓋了 Hugging Face、Megatron-LM、DeepSpeed、FairSeq 這些所有主流 Transformer 庫的優點,分散式訓練效能一如既往地力壓群雄,更重要的是,全域性視角程式設計最大程度降低了分散式使用門檻,讓大模型訓練飛入尋常百姓家。

LiBai 模型庫地址:\ http://github.com/Oneflow-Inc/libai

那麼,它具體是怎麼做到的?以下還將從訓練效能、易用性等方面對上述分散式訓練工具做對比,為你在下一次做分散式訓練時的工具選擇提供參考指南。

1

一鍵自動分散式訓練,效能超越 Megatron-LM 和 DeepSpeed

作為一個簡單高效的分散式模型訓練工具箱,具體而言,LiBai 有如下六大特點:

  • 支援單卡程式碼平滑地擴充套件到分散式。LiBai 內建的模型與 PyTorch 保持一致風格,大大降低學習和使用成本,只需要簡單配置,就可以便捷地擴充套件至任意規模的並行。這意味著,你可以在單卡上增加新功能,進行模型除錯,跑通程式碼後再絲滑地遷移到分散式上進行訓練。如果完全不想配置分散式訓練,或是覺得手動配置的分散式訓練太慢,那可以試用分散式託管特性,只需安裝自動並行的包(http://libai.readthedocs.io/en/latest/tutorials/basics/Auto_Parallel.html ),並在 LiBai 裡配置一行 graph.auto_parallel=True,就可以專注於模型本身,在完全不用操心分散式的同時獲得較快的訓練速度。
  • 相容 Hugging Face。OneFlow 和 PyTorch 在 API 層次高度相容,可以通過簡單的程式碼修改就可以匯入 Hugging Face 模型,只須 import oneflow as torch ,基於 LiBai 的資料並行、自動混合精度、Activation Checkpoint、ZeRO 等機制進行一個大規模模型的訓練。如果把模型的個別層次替換為 LiBai 內建的 layers ,就可以使用 3D 並行來訓練一個大模型。
  • 模組化設計。在 LiBai 的實現中,不僅為模型構建提供可複用的基礎計算模組,也針對資料載入、訓練邏輯、指標計算等做了抽象和模組化處理,方便使用者根據自己的需求重寫,然後作為外掛整合到 LiBai 的訓練系統中進行訓練。
  • 開箱即用。大模型訓練通常需要依賴一些技術,LiBai 提供了混合精度訓練、梯度重計算、梯度累加、ZeRO 等特性,可以輕鬆與資料並行、模型並行、流水並行組合使用。
  • 快速復現實驗。OneFlow 團隊參考了 Detectron2 LazyConfig(http://github.com/facebookresearch/detectron2/blob/main/docs/tutorials/lazyconfigs.md ) 來構建 LiBai 的配置系統,相比於傳統的 argparse 和 yacs-based 配置方式,LiBai 的配置系統更加靈活,使用 Python 語法完成整體構建,所以新增新的引數和模組非常方便,只需要 import 對應的模組即可完成新模組的新增。同時,訓練配置還可以序列化成 yaml 檔案進行儲存,方便直接在檔案中進行關鍵字搜尋來查詢配置項,如果使用者想要復現之前的實驗的結果,也直接傳入儲存的 config.yaml 作為訓練配置,保留非常多指令碼的檔案既不利於檢視有效修改,在復現實驗的同時也容易弄混實驗配置。
  • 高效效能。通過和 Megatron-LM 進行嚴格的 kernel 對齊,實現了多種 kernel fusion 操作,同時得益於 OneFlow 靜態圖的設計,不管是單卡效能還是各種組合並行的效率,LiBai 都優於英偉達深度優化的 Megatron-LM 和微軟的 DeepSpeed。

得益於 OneFlow SBP 原生支援各種並行技術,LiBai 實現了演算法描述和並行系統的解耦,僅用 3 萬多行程式碼,它就實現了英偉達 Megatron-LM 和微軟 DeepSpeed 兩大流行方案共計 10 萬行程式碼才能涵蓋的功能 。

資料最具說服力,以下的所有實驗資料均在相同的硬體環境、相同的第三方依賴(CUDA、 cuDNN 等)、使用相同的引數和網路結構下,全面對比了 LiBai 和 Megatron-LM 在不同模型下的效能表現(所有效能結果均公開且可復現, http://libai.readthedocs.io/en/latest/tutorials/get_started/Benchmark.html )。未來,OneFlow 團隊會公佈更大規模叢集上 LiBai 的表現。

  • Megatron-LM 固定 commit:http://github.com/NVIDIA/Megatron-LM/commit/e156d2fea7fc5c98e645f7742eb86b643956d840
  • LiBai commit: http://github.com/Oneflow-Inc/libai/commit/9fc504c457da4fd1e92d854c60b7271c89a55222
  • OneFlow commit: http://github.com/Oneflow-Inc/oneflow/commit/55b822e4d3c88757d11077d7546981309125c73f

資料並行

注:以下每組引數的含義:

DP 資料並行、MP 模型並行、PP 流水並行、2D 並行、3D 並行

fp16:開啟混合精度訓練 (amp)

nl: num layers (當 Pipeline parallel size = 8 時,為了讓每個 stage 有相對數量的 layer 進行計算,我們將 num layers 從 24 調整為 48)

ac: enable activation checkpointing mb: micro-batch size per gpu gb: global batch size total dxmxp,其中: d = 資料並行度(data-parallel-size) m = 模型並行度(\ tensor-model-parallel-size) p = 流水並行度(\ pipeline-model-parallel-size)

1n1g 表示單機單卡,1n8g 表示單機 8 卡, 2n8g 表示 2 機每機 8 卡共 16 卡, 4n8g 表示 4 機共 32 卡 grad_acc_num_step = global_batch_size / (micro_batch_size * data_parallel_size) 展示的結果為 Throughout

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

(注:本組 num layers = 24,開啟 amp,1n1g micro-batch size = 24, 其餘組 micro-batch size = 16)

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

(注:本組 num layers = 24,開啟 amp,1n1g micro-batch size = 6, 其餘組 micro-batch size = 4)

模型並行

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

(注:本組 num layers = 24,開啟 amp, 開啟 activation checkpointing,

micro-batch size = 128, global batch size = 1024, grad acc step = 8)

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

(注:本組 num layers = 24,開啟 amp)

流水並行

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

(注:前兩組 num layers = 24,grad acc step = 8, 最後一組 num layers = 48, grad acc step = 16,均開啟 amp,開啟 activation checkpointing)

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

(注:前兩組 num layers = 24,grad acc step = 8, 最後一組 num layers = 48, grad acc step = 16,均開啟 amp,開啟 activation checkpointing)

2-D 並行

資料 + 模型並行

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

(注:本組均為 num layers = 24,均開啟 amp,開啟 activation checkpointing, micro-batch size = 128,grad acc step = 8)

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

(注:本組均為 num layers = 24,均開啟 amp,開啟 activation checkpointing, micro-batch size = 32,grad acc step = 8)

資料 + 流水並行

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

(注:本組均為 num layers = 24,均開啟 amp,開啟 activation checkpointing, micro-batch size = 128,grad acc step = 8)

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

(注:本組均為 num layers = 24,均開啟 amp,開啟 activation checkpointing, micro-batch size = 32,grad acc step = 8)

3-D 並行

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

(注:本組均為 num layers = 24,均開啟 amp,開啟 activation checkpointing,grad acc step = 8)

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

(注:本組均為 num layers = 24,均開啟 amp,開啟 activation checkpointing,grad acc step = 8)

從上述效能對比資料可以看到,在嚴格對齊的實驗環境下,在 Bert、GPT-2 模型上,LiBai 的訓練速度全方位超過 Megatron-LM。

2

人有我優,人無我有:LiBai vs 其他訓練方案

如前所述,為了解決大模型訓練的難題,業內已經有 Hugging Face、DeepSpeed 和 Megatron-LM、FairSeq 等流行方案,OneFlow 是否還有必要再研發一個 LiBai 模型庫?

接下來,讓我們具體比較下上述模型庫的優勢和不足,也許你就會做出判斷。

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

HuggingFace:提供了全面的 SOTA Transformer 模型 ,便於使用預訓練好的模型進行微調,提供強大的社群和生態,方便開發者使用預訓練好的模型。但它只支援資料並行的方式,如果模型超過單個 GPU 視訊記憶體容量就不適用了,而且想用它做從零開始訓練速度也比較受限。

FairSeq:主要針對序列模型,在 NLP 和 CV 大一統的趨勢下,缺少對 CV 模型的支援。

Megatron-LM:基於 PyTorch 實現了資料並行、模型並行和流水並行,效能高,可以用來真正地訓練超大規模模型。

不過它做了大量的定製化,對於不熟悉分散式訓練的演算法工程師而言學習和使用門檻太高,基本上只能被分散式專家所複用。另外,Megatron-LM 提供的模型也遠遠少於 Hugging Face,使得想使用 PyTorch 復現大模型的工程師都必須要等待其他分散式高手基於 Megatron-LM 實現了自己想用的模型才行。

DeepSpeed:基於 PyTorch 的模型視訊記憶體優化相關的深度定製庫,提供分散式訓練、混合精度訓練、ZeRO 等技術,可以有效節約記憶體的開銷,使得在資料並行下也可以有效地訓練大模型。但是,DeepSpeed 還不支援模型並行,當模型某些層的引數大到超過了單個 GPU 的視訊記憶體,或是使用 DeepSpeed 的切分方式導致的通訊效率不是最優時,最好還是使用模型並行(Tensor 並行、流水並行),此時,只能結合 Megetron-LM 加上侵入原始碼的改造來滿足需求。

除了作為 PyTorch 生態中實現大模型訓練之鼻祖的 Megatron-LM 和 DeepSpeed ,來自國內外的多家知名機構也研發和推出瞭如 FairSeq 等若干大模型訓練庫,需要特別指出的是,這些庫的分散式核心功能無一例外都是基於 Megatron-LM、DeepSpeed 完成的。

LiBai 模型庫的與眾不同之處在於,它不是以上任一分散式訓練工具的簡單升級或包裝,而是基於 OneFlow 的分散式和圖編譯器能力構建的大規模預訓練模型開發套件。唯有如此,LiBai 不僅在效能上無出其右,在分散式易用性上更是不遑多讓:

  • 相容性。可以有效和目前基於 PyTorch 實現的 SOTA 模型相容,方便使用者快速遷移模型。
  • 高效性。無論是單卡還是多卡,使用者使用 LiBai 都可以提高訓練效率。
  • 易用性。LiBai 具有優秀的擴充套件性,可以很方便地根據需求修改模型,增加新功能,更快地完成原型功能的開發。以幾乎無感知、零學習成本的方式幫助使用者大幅降低分散式深度學習訓練的門檻,使用者在使用 LiBai 開發新模型和新功能時,只要會單張 GPU 程式設計就能自動擴充套件到大規模 GPU 的叢集,無須為分散式訓練重寫程式碼,從而提高開發的效率。

看完上述對比,相信 LiBai 也會成為 AI 工程師做分散式訓練的絕佳選擇,你覺得呢?

3

LiBai 支援所有常見並行訓練策略

分散式訓練大模型是個複雜問題,涉及到資料並行(data parallel),模型並行(tensor/model parallel),流水並行(pipeline parallel)等多種並行策略,LiBai 模型庫支援這三種常見的並行策略以及這些並行策略的任意組合(並行策略的基本概念:http://docs.oneflow.org/master/parallelism/01_introduction.html)。

自行實現這些並行策略讓人十分頭疼,比如以前為了使用自動混合精度訓練,需要學習配置 Apex;為了支援資料載入流水線,需要學習配置 DALI;為了使用 ZeRO 減少視訊記憶體佔用,需要學習配置 DeepSpeed …… 但用 LiBai 就完全不用擔心這類問題,它內建了多種並行策略且具備良好的可擴充套件性。

以下是 LiBai 中各類並行方法的例項。

萬能並行的實現方式

藉助 OneFlow 的 SBP 介面,使用者可以很方便地根據自身的需求,依照 GPU 的分組排布情況對網路中的輸入或者權重進行切分,以實現資料或張量並行。

在 LiBai 的 layers 模組(libai.layers)下,已內建一系列可自適應不同並行策略的網路層,包括常用的 Linear、MLP、Transformer 模組等,使用 LiBai 的 layers 搭建的神經網路, 只需調整配置檔案中關於分散式配置的超參,就可以輕鬆實現純資料並行、純張量並行以及資料 & 張量混合並行的訓練策略。

關於分散式配置的格式如下:

```

configs/common/train.py

Distributed arguments

dist=dict( data_parallel_size=1, tensor_parallel_size=1, pipeline_parallel_size=1, ) ```

通過 data_parallel_size 與 tensor_parallel_size 來控制輸入資料與模型權重在不同 GPU 組上的切分方式,當用戶使用 LiBai 的內建 layers 模組搭建好神經網路後,可以在自己的訓練配置檔案中修改分散式超參, 以實現不同的並行訓練策略,上圖所有值都取為 1 表示在單卡上執行。假設使用者擁有一臺 8 卡機器,下面介紹一下如何通過修改此配置檔案實現資料並行、張量並行以及流水並行訓練。

具體操作可參考 LiBai 分散式配置文件:http://libai.readthedocs.io/en/latest/tutorials/basics/Distributed_Configuration.html

純資料並行 & 純模型並行

當用戶要在 8 卡上進行純資料(或模型)並行訓練, 只需要在訓練配置檔案中對分散式超參進行覆寫即可:

純資料並行

```

your config.py

from libai.config import get_config train = get_config("common/train.py").train

train.dist.data_parallel_size = 8 ```

訓練時,在不同的 rank 上會複製一份相同的模型,每個 rank 會分別處理一部分的輸入資料, 以實現資料並行訓練。

  • 純模型並行

```

your config.py

from libai.config import get_config train = get_config("common/train.py").train

train.dist.tensor_parallel_size = 8 ```

在這種情況下, 模型會自動在 8 個 GPU 上進行切分, 每個 GPU 僅包含整體模型結構的一部分, 以實現模型並行訓練。

資料 & 模型混合並行訓練

當用戶要在 8 卡上進行資料與模型混合並行訓練, 只需要在訓練配置檔案中對分散式超參進行以下簡單改動:

```

your config.py

from libai.config import get_config train = get_config("common/train.py").train

train.dist.data_parallel_size = 2 train.dist.tensor_parallel_size = 4 ```

這種情況下, LiBai 會自動對 GPU 進行分組, 我們以 [0, 1, 2, 3, 4, 5, 6, 7] 對 8 個 GPU 進行編號,當設定了 data_parallel_size=2 以及 tensor_parallel_size=4 後,在執行時,會自動將 8 個 GPU 進行分組,可以表示為 [[0, 1, 2, 3], [4, 5, 6, 7]], 其中[0, 1, 2, 3] 為一組,[4, 5, 6, 7]為一組,執行時,會在組之間進行資料並行訓練,在組內進行模型並行訓練。

流水並行的配置

流水並行的核心概念可以簡單總結為:將網路分為多個階段(stage), 不同的 stage 被分發到不同的 GPU 上, 每個 stage 的計算結果傳遞給下一個 stage 進行計算,最終按接力的方式完成訓練。關於流水並行的具體內容可參考:\ http://docs.oneflow.org/master/parallelism/01_introduction.html#_6。

樸素流水並行配置

在 LiBai 下可以通過設定 placement 引數,將網路的不同層分配到不同的 GPU 上,placement 引數的值可以通過 libai.utils.distributed 下的 get_layer_placement()介面輕鬆配置,LiBai 會自動根據配置檔案(config)中的分散式配置,來做 stage 的切分,將不同的 placement 自動分配到不同的 stage 上,所以只需要為網路的每一層配置好 placement,再結合分散式配置,便可以輕鬆實現流水並行配置。

在大部分網路中,往往用一層 Linear 層作為網路的頭部(head), 產生網路的最終結果用作分類或者其他任務, 所以以 Linear 層為例, 簡要介紹 LiBai 中最簡單的流水並行配置方法:

``` from libai.layers import Linear

self.head = Linear(hidden_size, num_classes) ```

配置網路模組(module)的 placement

在 LiBai 中可以通過兩種方式將一層網路分配到對應的 placement 上:

1、通過 to_global 介面結合 get_layer_placement()來手動指定 placement, 這裡通過設定 get_layer_placement(-1)來將 head 層配置到最後一組接力的 placement 上。

``` from libai.layers import Linear import libai.utils.distributed as dist

self.head = Linear(hidden_size, num_classes).to_global(placement=dist.get_layer_placement(-1)) ```

2、(Recommended) 在 libai.layers 中實現的 module 自帶 layer_idx 引數, 可以直接設定 layer_idx 引數來指定這一層的 placement

``` from libai.layers import Linear

self.head = Linear(hidden_size, num_classes, layer_idx=-1) ```

配置輸入資料的 placement

在配置好了網路中模組的 placement 後, 還需要指定輸入資料的 placement, 因為只有當輸入和網路在同一個 stage 的時候才可以進行計算, 最直觀的方式就是為輸入和網路配置相同的 placement, 可以結合 to_global 與 get_layer_placement()實現:

``` class MyModule(nn.Module): def init(self, ... *, layer_idx): ... self.layer_idx = layer_idx ...

def forward(self, input_data):
    input_data = input_data.to_global(placement=dist.get_layer_placement(self.layer_idx))
    ...

```

結合配置檔案輕鬆實現樸素流水並行

在配置好網路中不同層的 placement 以及輸入的 placement 後,在執行流水並行前,使用者只需要調整配置檔案(config)即可,需要提前知道網路中的層數,並且調整配置檔案中的 pipeline_num_layers:

```

set the number of pipeline stages to be 2

train.dist.pipeline_parallel_size = 2

set model layers for pipeline

train.dist.pipeline_num_layers = hidden_layers ```

1F1B 是在 PipeDream(\ http://arxiv.org/pdf/1806.03377.pdf )中介紹的一種新的流水並行訓練方式,可以更好地節省視訊記憶體與利用資源。LiBai 也可以比較容易地支援這種 1F1B 的策略(http://github.com/Oneflow-Inc/libai/blob/main/docs/source/tutorials/advanced_tutorials/customize_dataloader.md

3D 並行的實現

掌握了資料 & 模型混合並行,以及流水並行以後,配置資料 + 模型 + 流水並行也只是綜合一下上述各種並行的改動即可。

```

your config.py

from libai.config import get_config train = get_config("common/train.py").train

train.dist.data_parallel_size = 2 train.dist.tensor_parallel_size = 2 train.dist.pipeline_parallel_size = 2

hidden_layers = 8 #網路的層數 train.dist.pipeline_num_layers = hidden_layers ```

還是以 8 卡作為例子,在設定 data_parallel_size,tensor_parallel_size, pipeline_parallel_size 都為 2 以後,在執行時,模型將根據使用者設定的 pinepine_num_layers 在 GPU 上自動進行劃分。

以上述配置為例,模型將在 [0, 1, 2, 3] 和[4, 5, 6, 7]號 GPU 上拆分為 2 個 stage。其中,stage0 會在 [0, 2] 和[1, 3]號 GPU 上資料並行;在 [0, 1] 和[2, 3]號 GPU 上模型並行;stage1 會在 [4, 6] 和[5, 7]號 GPU 上資料並行;在 [4, 5] 和[6, 7]號 GPU 上模型並行。

自定義並行訓練

根據上文的介紹,LiBai 在 libai/layers / 下提供了封裝好的模組供使用者呼叫。通過這些模組的組合,使用者可以拼湊出自己的並行網路。

當 LiBai 中的模組無法滿足使用者需求時,使用者也可以非常方便地自定義並行策略。不同於 PyTorch 下需要手工插入 scatter -> forward -> reduce 等一系列複雜的通訊操作,在 LiBai 中,使用者只需在初始化 tensor 時定義 sbp 和 placement,便可像寫單機執行的程式碼一樣跑起來自己的並行程式碼。(sbp 和 placement 的詳情可參考:\ http://docs.oneflow.org/master/parallelism/04_2d-sbp.html )。

舉例來說,在使用者進行 4 卡訓練時,網路的中間結果有一個 shape 為 (16, 8) 的 2D Parallel 的 tensor 在 GPU 上的劃分方式為如下圖, 在 LiBai 中。該 tensor 的 placement 分佈為 ranks=[[0, 1],[2, 3]],SBP 為 (S[0], S[1]) 或(S[1], S[0])。

``` [ |
X00 gpu0 | X01 gpu1


X10 gpu2 |  X11 gpu3
         |           ]

```

其中, Xij 的 shape 都為 (8, 4) 均勻的分佈在每張卡上, 如果你想對這個 tensor 加入一些隨機噪聲,那麼在 LiBai 中可以非常方便地加上如下程式碼:

LiBai 中封裝 dist.get_nd_sbp()是為了相容 1D parallel 的需求,同時 dist.get_layer_placement()是為了方便配置 pipeline parallel。大多數情況下,使用者可以直接參照以下程式碼:

```

test.py

import oneflow as flow from omegaconf import DictConfig from oneflow import nn

from libai.utils import distributed as dist

cfg = DictConfig( dict(data_parallel_size=2, tensor_parallel_size=2, pipeline_parallel_size=1)) dist.setup_dist_util(cfg)

class Noise(nn.Module): def init(self): super().init() self.noise_tensor = flow.randn( 16, 8, sbp=dist.get_nd_sbp([flow.sbp.split(0), flow.sbp.split(1)]), placement=dist.get_layer_placement(layer_idx=0) ) # 也可以換成以下的寫法 # self.noise_tensor = flow.randn( # 16, 8, # sbp=(flow.sbp.split(0), flow.sbp.split(1)), # placement=flow.placement("cuda", ranks=[[0, 1],[2, 3]]) # )

def forward(self, x):
    return x + self.noise_tensor

Noise = Noise()

x = flow.zeros( 16, 8, sbp=(flow.sbp.split(0), flow.sbp.split(1)), placement=flow.placement("cuda", ranks=[[0, 1],[2, 3]]) ) y = Noise(x)

print(f"rank: {flow.env.get_rank()}, global tensor: shape {y.shape} sbp {y.sbp} placement {y.placement}, local tensor shape: {y.to_local().shape}") ```

執行指令:

python3 -m oneflow.distributed.launch --nproc_per_node 4 test.py

以下顯示輸出,根據 shape 可以看到每個 rank 下 tensor 的分佈,以及在 global 視角下該 tensor 的資訊。

``` rank: 2, global tensor: shape oneflow.Size([16, 8]) sbp (oneflow.sbp.split(axis=0), oneflow.sbp.split(axis=1)) placement oneflow.placement(type="cuda", ranks=[[0, 1], [2, 3]]), local tensor shape: oneflow.Size([8, 4])

rank: 3, global tensor: shape oneflow.Size([16, 8]) sbp (oneflow.sbp.split(axis=0), oneflow.sbp.split(axis=1)) placement oneflow.placement(type="cuda", ranks=[[0, 1], [2, 3]]), local tensor shape: oneflow.Size([8, 4])

rank: 1, global tensor: shape oneflow.Size([16, 8]) sbp (oneflow.sbp.split(axis=0), oneflow.sbp.split(axis=1)) placement oneflow.placement(type="cuda", ranks=[[0, 1], [2, 3]]), local tensor shape: oneflow.Size([8, 4])

rank: 0, global tensor: shape oneflow.Size([16, 8]) sbp (oneflow.sbp.split(axis=0), oneflow.sbp.split(axis=1)) placement oneflow.placement(type="cuda", ranks=[[0, 1], [2, 3]]), local tensor shape: oneflow.Size([8, 4]) ```

4

未來計劃

大模型訓練難於上青天?效率超群、易用的“李白”模型庫來了

LiBai 目前已支援 BERT、GPT、ViT、Swin-Transformer、T5 等常見模型,以及 MoCoV3、MAE 等最新研究,開箱即用,並且可以很方便地在下游任務上進行微調。

此外,OneFlow 也會更好地相容 Hugging Face 的模型,接入其生態,同時再利用 OneFlow 自動並行功能,幫助使用者享受只寫單卡程式碼即自動擴充套件到分散式系統的一勞永逸的激爽體驗。

未來,在支援更多模型訓練的基礎上,OneFlow 也會持續完善推理和 Serving 相關的功能,從而打通訓練和部署的全流程,讓 OneFlow 成為使用者的一站式開發平臺。

  • LiBai 模型庫地址:http://github.com/Oneflow-Inc/libai
  • LiBai 文件地址:http://libai.readthedocs.io/en/latest
  • OneFlow 專案地址:http://github.com/Oneflow-Inc/oneflow

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