MindSpore技術理解(上)
引言
深度學習研究和應用在近幾十年得到了爆炸式的發展,掀起了人工智能的第三次浪潮,並且在圖像識別、語音識別與合成、無人駕駛、機器視覺等方面取得了巨大的成功。這也對算法的應用以及依賴的框架有了更高級的要求。深度學習框架的不斷發展使得在大型數據集上訓練神經網絡模型時,可以方便地使用大量的計算資源。
深度學習是使用多層結構,從原始數據中自動學習並提取高層次特征的一類機器學習算法。通常,從原始數據中提取高層次、抽象的特征是非常困難的。目前有兩種主流的深度學習框架:一種是在執行之前構造一個靜態圖,定義所有操作和網絡結構,典型代表是TensorFlow,這種方法以犧牲易用性為代價,來提高訓練期間的性能;另一種是立即執行的動態圖計算,典型代表是PyTorch。通過比較可以發現,動態圖更靈活、更易調試,但會犧牲性能。因此,現有深度學習框架,難以同時滿足易開發、高效執行的要求。
簡介
MindSpore作為新一代深度學習框架,是源於全產業的最佳實踐,最佳匹配昇騰處理器算力,支持終端、邊緣、雲全場景靈活部署,開創全新的AI編程范式,降低AI開發門檻。MindSpore是一種全新的深度學習計算框架,旨在實現易開發、高效執行、全場景覆蓋三大目標。為了實現易開發的目標,MindSpore采用基於源碼轉換(Source Code Transformation,SCT)的自動微分(Automatic Differentiation,AD)機制,該機制可以用控制流表示復雜的組合。函數被轉換成函數中間表達(Intermediate Representation,IR),中間表達構造出一個能夠在不同設備上解析和執行的計算圖。在執行前,計算圖上應用了多種軟硬件協同優化技術,以提升端、邊、雲等不同場景下的性能和效率。MindSpore支持動態圖,更易於檢查運行模式。由於采用了基於源碼轉換的自動微分機制,所以動態圖和靜態圖之間的模式切換非常簡單。為了在大型數據集上有效訓練大模型,通過高級手動配置策略,MindSpore可以支持數據並行、模型並行和混合並行訓練,具有很強的靈活性。此外,MindSpore還有“自動並行”能力,它通過在龐大的策略空間中進行高效搜索來找到一種快速的並行策略。MindSpore框架的具體優勢,請查看詳細介紹。
深度學習研究和應用在近幾十年得到了爆炸式的發展,在圖像識別[1]、語音識別與合成[2]、游戲[3]、 語言建模與分析[4]等方面取得了巨大的成功。深度學習框架[5–9]的不斷發展使得在大型數據集上訓練 神經網絡模型時,可以方便地使用大量的計算資源。 目前有兩種主流的深度學習框架:一種是在執行之前構造一個靜態圖,定義所有操作和網絡結構, 典型代表是 TensorFlow[5],這種方法以犧牲易用性為代價,來提高訓練期間的性能;另一種是立即 執行的動態圖計算,典型代表是 PyTorch[6]。通過比較可以發現,動態圖更靈活、更易調試,但會犧 牲性能。因此,現有深度學習框架難以同時滿足易開發、高效執行的要求。 本文提出了一種新的深度學習框架——MindSpore,該框架旨在實現三個目標:易開發、高效執行 和全場景覆蓋。MindSpore 由四個主要組件組成:MindExpression(ME)、GraphEngine(GE)、 MindData(MD)和 MindArmour(MA)。表 1 總結了 MindSpor的技術貢獻。
1. ME 為用戶提供了 Python 編程范式,具有以下三個突出特點。
− 自動微分:采用基於源碼轉換的自動微分機制,在訓練或推理階段,將一段 Python 代碼轉 換為數據流圖。因此,用戶可以方便地使用 Python 原生控制邏輯來構建復雜的神經網絡模型。
− 自動並行(Automatic Parallelization):由於大規模模型和數據集的不斷增加,跨分布式設 備並行化深度神經網絡(Deep Neural Network,DNN)訓練已經成為一種常見做法。然而, 當前框架(如 TensorFlow[5]、Caffe[10]和 MXNet[7])用於並行化訓練的策略仍然簡單直接, 而且通常是次優的。MindSpore 並行化訓練任務,透明且高效。“透明”是指用戶只需更改一行配置,提交一個版本的 Python 代碼,就可以在多個設備上運行這一版本的 Python 代碼 進行訓練。“高效”是指該算法以最小的代價選擇並行策略,降低了計算和通信開銷。
− 動態圖:MindSpore 支持動態圖,無須引入額外的自動微分機制(如算子重載微分機制), 從而大大增加了動態圖和靜態圖的兼容性。
2. GE 負責硬件相關的資源管理和優化,將平台特有的信息傳遞給 ME。GE 接收來自 ME 的數據流圖,並將該圖中的算子調度到目標設備上執行。GE 將數據流圖分解為優化后的子圖,並將它們調度到不同的設備上。GE 將每個設備抽象為一個執行引擎(Execution Engine),並提供執行引擎插件機制,用來支持各種不同的設備,而不影響其他組件。
3. MD 負責數據處理,並提供工具來幫助開發者調試和優化模型。通過自動數據加速技術實現了高性能的流水線,以進行數據處理。各種自動增強策略的出現,使用戶不必再尋找合適的數據增強策略。訓練看板將多種數據集成在一個頁面,方便用戶查看訓練過程。分析器可以打開執行黑匣子,收集執行時間和內存使用的相關數據,從而有針對性地進行性能優化。
4. MA 負責提供工具,以幫助開發者防御對抗性攻擊,實現隱私保護的機器學習。在形式方面,MA 提供了以下功能:生成對抗代碼、評估模型在特定對抗環境中的性能、開發出更健壯的模=型。MA 還支持豐富的隱私保護能力,如差分隱私[11]、機密人工智能計算[12]、可信協同學習[13、 14]等。
本文的其余內容如下:第 2 章介紹了 MindSpore 的架構;第 3 章通過闡述自動微分、自動並行和動態圖支持的設計細節,介紹了 MindSpore 的核心模塊——ME;第 4 章介紹了 GE 的工作過程;第 5 章和第 6 章分別介紹了 MD 和 MA 的相關細節;第 7 章介紹了 MindSpore 支持的端雲協同架構;為了說明高效執行這一特性,第 8 章闡述了使用 ResNet50 作為基准來評估“自動並行”特性,還介 紹了 MindSpore 的訓練和推理性能;最后,第 9 章對我們的工作進行了總結,並介紹了未來的重點 研究方向。
2 MindSpore 概述
2.1 MindSpore 架構如圖 1 所示,MindSpore 由 ME 和 GE 組成。
ME 提供 Python 接口和自動微分功能,可分為前端和后端。前端用於定義用戶級應用軟件編程接口(Application Programming Interface,API),用於構建和訓練神經網絡。由於 MindSpore 采用基於源碼轉換策略的自動微分機制,因此用戶可以 Python 方式編程。一種典型的情況是,用戶可以用常規的 Python 方式構建復雜的模型和流程控制,如 if、else、while 等。 ME后端是高性能執行和自動微分的核心,自動微分基於源碼轉換方式。如果用戶選擇以 Pynative 模式運行,則逐個下發算子運行。如果模型以圖模式運行,則后端會使用流水線根據 Python 代碼生成計算圖,通過解析 Python 代碼,生成抽象語法樹(Abstract Syntax Tree,AST),然后將其轉換 為 A-Normal-Form(ANF)圖[15]。ANF 是圖形化的,而不是語法化的,因此從算法上操作起來要容 易得多[16]。如果用戶需要訓練神經網絡,流水線自動生成反向計算節點,並添加到 ANF 圖中。流水線在構造完整圖之后進行許多優化(例如內存復用、算子融合、常數消除等)。如果用戶需要在分布式環境中訓練模型,流水線將應用自動並行提供的優化策略(參見 3.2 節)。后端的虛擬機通過會話 管理計算圖,調用 GE 來運行圖,並控制圖的生命周期。 GE 負責控制從 ME 接收的數據流圖的執行,GE 也是算子和不同硬件資源的管理器。在 Pynative 模式,GE 一次只能從 ME 接收一個算子,並立即將其下發運行。在圖形模式下,GE 從 ME 接收一 個計算圖,圖編譯器根據不同的執行引擎,將計算圖分解為若干個子圖,並對不同的執行引擎采用 不同的優化函數。這一過程充分利用了底層的硬件資源。然后,編譯器靜態分配必要的資源,並構建將交付到不同硬件資源的內核。任務執行器是內核執行前的最后一步,它通過調用由將運行內核 的相應硬件提供的執行引擎,啟動許多任務,以運行接收到的內核。 MD 中的數據處理在訓練過程中完成數據流水線,包括數據加載、數據論證、數據轉換等。MD 也提 供簡單的 API,支持 CV/NLP/GNN 等全場景的數據處理能力。同時,在這個過程中,如何提高數據處理能力,以匹配人工智能芯片的算力,是保證人工智能芯片發揮極致性能的關鍵。MD 中的 MindInsight 有四個模塊:訓練看板、溯源、分析器和調試器。分析器可以打開執行黑匣子,以收集執行時間和內存使用的相關數據。調試器是圖執行模式下的調試工具,用戶可以在訓練過程中,查看圖的內部結構和節點的輸入/輸出。MindInsight 對訓練過程生成的日志文件進行分析。開發者可以輕松地可視化訓練過程,在圖形用戶 界面(Graphical User Interface,GUI)上與 MindInsight 對比不同的路徑。 MA 幫助用戶開發更健壯的模型,保護用戶在訓練和推理數據中的隱私。MA 中的對抗性攻擊防御有 三個主要模塊:攻擊、防御和評估。攻擊模塊是一個在黑盒和白盒攻擊場景下生成對抗樣本的工具。 防御模塊從攻擊模塊接收對抗樣本,並在訓練過程中用這些樣本來提高模型的魯棒性。評估模塊提供了多種評估指標,允許開發者更容易可視化其模型的魯棒性。MA 為了實現隱私保護機器學習,實現了一系列差分隱私感知優化器,這些優化器在訓練過程中自動將噪聲添加到生成的梯度中。
2.2 編程范式 MindSpore
為用戶提供 Python 編程范式。借助基於源碼轉換的自動微分,用戶可以使用原生 Python 控制語法和其他一些高級 API,如元組(Tuple)、列表(List)和 Lambda 表達。為避免用戶混淆,MindSpore 引入了盡可能少的接口和概念。在單機平台上訓練簡單的神經網絡時, 用戶只需要了解以下五個組件:
張量(Tensor):是一個包含相同數據類型元素的多維矩陣。與其他一些訓練框架不同,MindSpore 中沒有標量(Scalar)變量(Variable)的概念。要計算一個張量的梯度,該張量的 requires_grad 自動微分屬性應該設置為 True。可以用 Numpy 初始化張量,也可以將張量的值轉換為一個 Numpy 對象。
數據集(Dataset):是一個單獨的異步流水線,該流水線為在訓練中將張量無延遲地饋入網絡的剩 余部分做准備。
算子(Operator):是構成神經網絡的基本計算單元。MindSpore 支持大多數常用的神經網絡算子 (如卷積、批標准化、激活等)和數學算子(如加法、乘法等)。此外,MindSpore 支持自定義算子, 允許用戶添加新的算子來適配特定的硬件平台,或者合並現有算子來生成復合算子。
單元(Cell):是張量和算子的集合,是所有神經網絡單元的基本類。一個單元也可以包含其他單元, 允許它們嵌套在同一個樹狀結構中。用戶通過在單元中定義 construct 函數來表達神經網絡的計算邏輯。construct 函數中的計算將在每個調用中執行。
模型(Model):是 MindSpore 中的一個高級 API,它封裝了一些低級 API,從而使用戶能夠更加方便地使用推理和訓練功能。如果用戶熟悉低級 API,並希望對計算過程建立細粒度的控制,則不一定需要使用模型。 從用戶的角度來看,用 MindSpore 編寫程序的核心是構建一個對應於神經網絡的單元。這個過程首 先從輸入張量開始,可以是常量張量或參數張量。然后,用戶使用 MindSpore 提供的不同算子構造 一個單元。最后,用戶可以使用模型封裝這個單元來訓練神經網絡,也可以直接將輸入數據傳遞給 單元來執行推理任務。
代碼 1 展示了 Python 中 MindSpore 程序。
該程序展示了定義以及訓練 LeNet[17]神經網絡的過程。 代碼前 6 行,導入 MindSpore 合適的庫。第 7 行到第 25 行定義了與 LeNet 神經網絡相對應的 LeNet5 單元。_init_函數實例化 LeNet 使用的所有算子。construct 函數定義了 LeNet 的計算邏輯。第 26 行和第 27 行從 Mnisit 數據集讀取數據,並生成一個迭代器 ds 用作訓練的輸入。第 28 行將 LeNet5 類實例化為 network。用 SoftmaxCrossEntropyWithLogits 函數來計算損失(loss)(第 29 行),並 用 momentum 來優化參數(第 30 行到第 31 行),loss 和 optimizer 用來創建模型(Model)對象。 最后,用 epoch 來控制迭代次數,調用模型的訓練方法,並在每個 eval_step 對模型進行評估。
代碼 1 LeNet5 的 MindSpore 實現
1 import mindspore.nn as nn
2 from mindspore.ops import operations as P
3 from mindspore.network.optim import Momentum
4 from mindspore.train import Model
5 from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits
6 import mindspore.dataset as de
7 class LeNet5(nn.Cell):
8 def __init__(self):
9 super(LeNet5, self).__init__()
10 self.conv1 = nn.Conv2d(1, 6, 5, pad_mode='valid')
11 self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid')
12 self.fc1 = nn.Dense(16 * 5 * 5, 120)
13 self.fc2 = nn.Dense(120, 84)
14 self.fc3 = nn.Dense(84, 10)
15 self.relu = nn.ReLU()
16 self.max_pool2d = nn.MaxPool2d(kernel_size=2
17 self.flatten = P.Flatten()
18 def construct(self, x):
19 x = self.max_pool2d(self.relu(self.conv1(x)))
20 x = self.max_pool2d(self.relu(self.conv2(x)))
21 x = self.flatten(x)
22 x = self.relu(self.fc1(x))
23 x = self.relu(self.fc2(x))
24 x = self.fc3(x)
25 return x
26 ds = de.MnistDataset(dataset_dir="./MNIST_Data")
27 ds = ds.batch(batch_size=64)
28 network = LeNet5()
29 loss = SoftmaxCrossEntropyWithLogits()
30 optimizer = nn.Momentum(network.trainable_params(),
31 learning_rate=0.1, momentum=0.9)
32 model = Model(network, loss, optimizer)
33 model.train(epoch=10, train_dataset=ds)
3 MindExpression
3.1 基於源碼轉換的自動微分
目前主流的深度學習框架有三種自動微分技術:
l 基於靜態計算圖的轉換:在編譯時將網絡轉換為靜態數據流圖,然后將鏈式規則轉換為數據流圖,實現自動微分。
l 基於動態計算圖的轉換:以算子重載的方式記錄前向執行時網絡的操作軌跡,然后將鏈式規則,應用到動態生成的數據流圖中,實現自動微分。
l 基於源碼的轉換:該技術是從函數式編程框架演化而來,對中間表達(程序在編譯過程中的表 達形式),以即時(Iust-In-Time,JIT)編譯的形式進行自動微分變換,支持復雜的流程控制場 景、高階函數和閉包。基於源碼轉化的自動微分如圖 2 所示。 TensorFlow 早期采用靜態計算圖,而 PyTorch 采用動態計算圖。靜態圖可以利用靜態編譯技術優化網絡性能,但是組建或調試網絡非常復雜。使用動態圖非常方便,但很難在性能上達到極限優化。
MindSpore 開發了一種新的策略,即基於源碼轉換的自動微分。
一方面,它支持流程控制的自動微分,因此構建像 PyTorch 這樣的模型非常方便。
另一方面,MindSpore 可以對神經網絡進行靜態編譯優化,從而獲得良好的性能。
MindSpore 自動微分的實現可以理解為對程序本身進行符號微分,因為 MindSpore IR 是函數式的中間表達,它與基本代數中的復合函數有直觀的對應關系,只要已知基礎函數的求導公式,就能推導出由任意基礎函數組成的復合函數的求導公式。MindSpore IR 中每個原語操作可以對應為基礎代數中的基礎函數,這些基礎函數可以構建更復雜的流程控制。
3.2 自動並行
隨着深度學習的發展,為了實現更高的准確率和更豐富的應用場景,訓練數據集和深度神經網絡模型的規模日益增大。特別是自然語言處理(Natural Language Processing,NLP)領域,數據集的范圍從 200MB 到 541TB 不等。模型的尺寸也從 BERT[18]的 3.4 億個參數,Transformer-x l[19]的 8 億 個參數,GPT-2 [19]的 150 億個參數,到最近 NVIDIA Megatron-LM[20]的 80 多億個參數。因此,在大型數據集上訓練大型模型,不僅需要深度學習框架支持數據並行和模型並行,還需要支持混合並行。
當前的主流框架(如 TensorFlow[5]、Caffe[10]和 MXNet[7])大多通過手動切分深度神經網絡模型來實現模型並行,但手動切分模型難度非常大,對開發者的要求非常高,需要有豐富經驗的專家來操作。
實現混合並行(數據並行和模型並行同時進行)又極大增加了開發的復雜度。最近的研究成果[21–25] 提出了簡化混合並行的方法,但這些方法在幾個方面都存在局限性。
首先,這些方法在整個模型中,固定了對張量維數的切分策略,這可能導致切分策略是次優的,因為模型的不同部分適用不同的策 略。
第二,[22,24]不適合用於許多用於語言建模和行人重新識別(Re-identification,ReID)的深度神經網絡,這些網絡往往是非線性網絡。
第三,[21]將划分策略搜索問題表述為混合整數程序,並使用現有的求解程序來求解,當處理大型模型時,速度會變得非常慢。此外,文獻[21,24,25]旨在僅優 化通信或內存開銷,可能不會減少訓練時間。
MindSpore 的目標是在模型訓練過程中允許並行轉換。為此,在並行化策略搜索中引入了張量重排 布(Tensor Redistribution,TR),這使輸出張量的設備布局在輸入到后續算子之前能夠被轉換,如 3.2 圖 3 中紅色矩形所示。但是,對於復雜大型模型的搜索並行策略,在考慮張量重排布時,需要克服的挑戰主要有兩個。
首先,既然張量重排布將通信算子(例如 AllGather)引入數據流圖,那么如何像普通算子一樣自動地對通信算子求導呢?對於每個相應的前向算子,都需要獲取反向算子,用來更新可訓練參數。目前的框架需要專家在反向階段手動添加 SEND 和 RECV 源語來傳遞梯度, 這對模型開發者來說是一項具有挑戰性的工作,尤其是在模型比較復雜的情況下。
其次,隨着張量重排布對策略空間的極大擴展,如何為復雜的大型模型高效地找到一個好的策略?在功能和效率方面,該算法需要為具有非線性結構的大模型快速找到一種策略。在性能方面,算法返回的策略應該會縮短端到端的訓練時間。需要對運行成本進行仔細建模,這一過程也增加了人工成本。
MindSpore 針對上述兩個挑戰,推出了解決方案。為了實現通信算子的自動微分,MindSpore 定義 了通信算子的反向算子。例如,AllGatheris的反向算子為ReduceScatter,SEND的反向算子為RECV 后跟一個 ADD。定義這些反向算子十分重要,因為 Auto-diff 過程可以一次性地區分整個前向圖,而無須跳過任何算子,這也是為什么 Auto-diff 是 Auto-parallel 后面一步的原因。
針對第二個挑戰,在 同時考慮計算和通信開銷的情況下,建立一個代價模型來選擇一個好策略。為了快速地為復雜大圖找到一個好策略,提出了幾種方法:
一種是支持多圖操作的算法,將原始圖轉換成線性圖;
一種是策略分離機制,在保證返回解的精度的同時,有效地縮小搜索空間。例如,ResNet50 在 8 台設備上搜索並行策略的時間在 1s 內,而返回的解決方案確實縮短了訓練時間。例如,當模型較大(類的 數量超過 128K)時,返回的解決方案與原始數據並行策略相比減少了大約 55%的訓練時間。
代碼 2 半自動並行配置中的並行過渡
1 class Submodel(nn.Cell):
2 def _init_(self, shape):
3 self.bn = BatchNorm(set_strategy={[4, 1]})
4 self.matmul = MatMul(set_strategy={[1, 1], [1, 4]})
5 self.W = Parameter(Tensor(shape), require_grad=True)
6 def construct(self, X):
7 Y = self.bn(X)
8 Z = self.matmul(y, self.W)
9 return Z
MindSpore 較靈活,它支持用戶指定的高級策略配置,稱之為半自動並行(semi-auto-parallel)。在 代碼 2 和圖 3 中,展示了一個從數據到模型的並行轉換的例子。該子模型的結構為 BatchNorm 算子后跟一個 MatMul 算子,廣泛應用於 ResNet、ReID 等分類任務。在 BatchNorm 算子中,X 按行拆 分為四部分,數據可以並行,效率非常高。在 MatMul 算子中,可學習參數的權重 W 被分成四部分,模型可以並行,由於參數數量較多,這部分的模型並行更有效。由於 BatchNormi 的輸出布局與 MatMul 的輸入布局不同,所以框架插入了一個張量重排布(該例中為 AllGather 和 ConCat),這一過程對用戶是透明的。用戶也不必關注哪個設備運行了模型的哪個部分,框架會自動安排。
然而,不同的模型結構在每個算子中具有不同大小的參數,如圖 4 所示,並且它們使用於不同的切分策略。在圖 4(3)中,將第一算子配置為模型並行,將后續算子配置為數據並行,也需要插入張量重排布,這樣可以獲得更好的性能。
在訓練新模型時,多次配置 set_strategy,耗時耗力。在這種情況下,如果配置了自動並行 (auto-parallel),則不需要指定 set_strategy,該算法將找到一個有效的策略。例如,當 ResNet 中 的分類數量超過 130K 時,算法返回的策略導致在 50ms 內訓練一個迭代。相比之下,原始數據並 行訓練一次迭代超過 111ms。更詳細的評估,見 8.1 節。
3.3 動態圖
由於編譯器能獲得靜態圖的全局信息,所以靜態圖在大多數情況下都表現出更好的運行性能。而動態圖可以保證更好的易用性,使用戶能夠更加方便地構建和修改模型。為了同時支持靜態圖和動態圖,大多數先進的訓練框架,需要維護兩種自動微分機制,即基於 Tape 的自動微分機制和基於圖的自動微分機制。從開發者的角度來看,維護兩套自動微分機制成本較高;從用戶的角度來看,在靜態模式和動態模式之間切換也相當復雜。
MindSpore 采用了基於源碼轉換的自動微分機制,同時支持靜態圖和動態圖,高效易用。在 MindSpore 中,稱動態圖為 Pynative 模式,因為代碼使用 Python 解釋器在這種模式下運行。如代 碼 3 所示,從靜態圖模式切換到動態圖模式只需要一行代碼,反之亦然。此外,為了提高 Pynative 模式的運行效率,MindSpore 支持 staging 機制,如代碼第 4 行所示。在函數(fc_relu)前添加 ms_function 裝飾器,該函數將以靜態圖模式編譯和運行。
代碼 3 LeNet5 的 MindSpore 實現
1 from mindspore import context
2 import numpy as np
3 class LeNet5(nn.Cell):
4 @ms_function
5 def fc_relu(self, x):
6 x = self.relu(self.fc2(x))
7 x = self.fc3(x)
8 def construct(self, x):
9 x = self.max_pool2d(self.relu(self.conv1(x)))
10 x = self.max_pool2d(self.relu(self.conv2(x)))
11 x = self.flatten(x)
12 x = self.relu(self.fc1(x))
13 x = fc_relu(x)
14 return x
15 data=np.ones((batch_size,3,224,224),np.float32)*0.01
16 net = LeNet5()
17 # switch to Pynative mode
18 context.set_context(mode=context.PYNATIVE_MODE)
19 pynative_out = net(data)
20 # switch back to static graph mode
21 context.set_context(mode=context.GRAPH_MODE)
22 graph_out = net(data)
靜態圖和動態圖的架構如 3.3 圖 5 所示。在動態圖模式中,框架遍歷模型的所有算子,為每個算子 生成計算圖,並將計算圖傳遞給 GE 進行前向計算。在完成前向計算后,框架為模型生成前向和反向計算圖,並將它們發送到 GE 中執行。由於 MindSpore 使用基於源碼轉換的自動微分機制,因此,在生成反向圖時,可以省略用戶為檢查其模型而設置的所有代碼,如 pdb 和 print。