本案例中，阿里巴巴PAI團隊為了使PAI用戶無需經過繁瑣的模型改寫即可實現混合精度訓練、充分利用Tesla V100 GPU Tensor核心的強大計算力縮短模型訓練時間、提升模型迭代速度，阿里巴巴團隊在PAI平臺上開發了自動混合精度訓練Feature，Feature已在PAI平臺上線開放對內使用，顯著節省了用戶使用Tensor核心硬件加速單元的負擔，并且以通用性的方式獲得了顯著的作業訓練性能加速，比較充分地釋放出底層Tesla V100的計算潛力，且不影響模型的收斂性。

PAI（Platform Artificial Intelligence) 作為阿里巴巴集團的機器學習平臺，一直致力于底層引擎技術、編譯優化及在離線預測優化的技術創新，同時通過PAI Studio、PAI DSW 、EAS、算法市場等相關生態產品，打造完整、豐富的生態體系，將AI技術更好的賦能于集團內外各類用戶。

PAI平臺管控了阿里巴巴集團內外大量的CPU、GPU、RDMA網卡等硬件資源，在確保單個作業執行性能的同時，如何確保這些硬件資源的整體使用效率，是平臺側重點關注的問題。因為確保通過平臺層面的技術創新，讓阿里 “用更少的資源，跑得更快，支持更多業務同時迭代” 是中臺技術賦能業務創新的典型體現。

NVIDIA Tesla V100 GPU是NVIDIA于2017年推出的新一代GPU，采用Volta架構，引入了針對混合精度計算的Tensor核心，硬件層面顯著提升了峰值計算能力，相較于深度學習常用的浮點32位的計算，理論最高加速比可以達到8X。

NVIDIA Tesla V100 GPU推出之后，阿里巴巴集團率先采用了Tesla V100。然而，一款新硬件的推出對上層軟件棧也提出了新的挑戰和要求，想要充分發揮新的底層硬件的性能效率，就需要軟件層進行大量的適配工作（包括訓練性能以及模型訓練方式的適配，以確保獲得理想的time-to-accuracy的加速效果）。

為了使PAI用戶無需經過繁瑣的模型改寫即可實現混合精度訓練、充分利用NVIDIA V100 GPU Tensor核心的強大計算力縮短模型訓練時間、提升模型迭代速度，阿里團隊在PAI平臺上開發了自動混合精度訓練Feature。

該Feature已在PAI平臺上線開放對內使用，顯著節省了用戶使用Tensor核心硬件加速單元的負擔，并且以通用性的方式獲得了顯著的作業訓練性能加速，比較充分地釋放出底層新硬件的計算潛力，同時在近期向社區提交了相關的Pull Request因為有效發揮混合精度加速新硬件的效率確實存在不少陷阱和問題。

作為基礎設施平臺方，阿里巴巴PAI團隊希望能夠從平臺層面將這些問題盡可能以透明、通用的方式進行解決，從而解決上層用戶的使用負擔，讓用戶透明地享受到Telsa V100新GPU的加速紅利。

已經正常使用PAI自動混合精度訓練的用戶，均能夠在無需改動一行代碼的前提下獲得顯著的訓練加速，且不影響模型的收斂性。

1. 混合精度訓練簡介

混合精度訓練是指在訓練過程中在不同的計算中使用不同的數值精度，從而充分挖掘GPU硬件上每一個晶體管的極致性能，目前在PAI上支持的混合精度訓練主要是指FP16和FP32兩個數值精度。在標準Tensorflow訓練任務中，variable及tensor的表示精度目前均默認為FP32的，那么為什么要引入FP16的數值精度呢？

對于神經網絡訓練過程而言，影響速度的兩個關鍵因素是，計算和訪存（這實際上也是由馮諾依曼體系結構的本質決定的）：

從計算而言，NVIDIA在2017年推出了Volta架構，其中重要的模塊即為Tensor核心（如下圖），在Tensor核心的加速下，V100在混合精度下的吞吐量比在FP32下有8X的加速。Tensor核心主要對matrix-multiply-and-accumulate類計算進行了加速。基于Tensor核心阿里巴巴PAI團隊可以對 MatMul、 Convolution等計算原語進行加速，而二者恰好是神經網絡中計算量較大的op。

從訪存而言，若Tensor的數值精度由FP32變為FP16，那么也可以得到理論上2X的訪存加速。 因此阿里巴巴PAI團隊引入了混合精度對神經網絡的訓練進行加速。

Volta V100 Tensor核心計算

2. Motivation

易用性是PAI平臺在追求高性能、高效率的同時同樣重視的另一個關鍵目標，因為其用戶，即AI Developer，希望專注于模型和算法本身，不希望被各種復雜的性能優化問題分散精力，這要求其性能優化能夠作為turn-key solution，盡可能對用戶透明，避免對用戶模型代碼的入侵。混合精度訓練加速優化同樣如此。

目前一般的混合精度訓練方法需要用戶自己將模型代碼修改為混合精度版本，因此會對模型代碼產生較深的入侵，特別是復雜模型，增加了用戶負擔。這個過程包含多個步驟：

1) 將計算部分轉換為FP16；

2) 將不支持FP16和不應使用FP16計算的算子再逐一轉換回FP32；

3) Variable使用FP32存儲，這可以通過將Variable定義為FP32，在計算時再cast乘FP16，即tf.cast(tf.get_variable(..., dtype=tf.float32), tf.float16)來做到。

但是這種做法存在很多局限。首先，很多variable的定義被封裝在TF API中，如LSTMCell的weights和bias，模型代碼層面不可見，這種情況則需要通過custom_getter實現；此外，有些情況下混合精度的訓練過程還需要加入Loss Scaling策略才能確保收斂。整個過程易出錯，一旦出現收斂性問題就比較難調試，而且難以區分轉換出錯還是模型本身不適合混合精度訓練。

此外，阿里團隊還發現手工混合精度轉換存在一些性能問題，比如：實測大部分模型手工簡單轉換為混合精度模型之后的加速遠低于預期，甚至出現大幅性能下降（如果是對混合精度訓練細節非常了解的優化專家來進行手工精細調優，確實會取得更好的性能，但要求所有的算法工程師都具備相應的能力其實是一個不合理的要求）。

分析下來，發現根源問題在于模型代碼層面的轉換是粗粒度的。部分算子其實天然不適合使用FP16，例如 Embedding Lookup，還有部分算子的FP16版本性能較差，例如 BiasAddGrad。模型代碼層面的粗粒度轉換很難對轉換過程進行精細控制從而將這些FP16性能很差的算子排除在外。

當然，模型代碼層面也可以做到細粒度的轉換，但是這樣做一方面增加模型代碼的入侵深度，使得人工工作量隨著模型復雜度急劇上升；另外這么做也需要一定的性能調優經驗，且避免不了多次迭代嘗試，這會給AI Developer帶來很大的額外負擔；最后還有一個克服不了的問題，那就是在模型代碼層面無法觸及反向計算圖，所以像前面提到的BiasAddGrad算子在模型代碼層面就控制不了。

基于以上兩點原因，阿里團隊在計算圖優化環節實現了自動的、精細的混合精度轉換，對用戶完全透明，用戶跟正常訓練FP32模型一樣，無需修改模型代碼；在計算圖優化環節，通過相應的性能優化避免混合精度轉換過程中伴生的性能不利因素，其實現的Loss Scaling策略對用戶更友好、適用面更廣泛。

3. 自動混合精度化的實現

阿里巴巴PAI團隊的自動混合精度轉換在PAI平臺深度定制的Tensorflow框架PAI-TensorFlow的計算圖優化引擎中實現。雖然相比于模型代碼層面的轉換，計算圖層面的轉換能夠做到自動化、對用戶透明，但是計算圖層面的轉換也有其缺點、有著相應的制約和性能開銷，本節主要介紹阿里巴巴PAI團隊在計算圖混合精度轉換的同時如何克服這些約束、降低性能開銷。

3.1 White-list

為了避免將不支持FP16、不適合使用FP16、FP16性能較差的算子轉換為FP16導致運行錯誤或者性能大幅下降，阿里巴巴PAI團隊維護了一個white-list，只有在這個white-list中的算子才有可能被轉換為FP16。這個white-list的更新，一方面，會基于一些heuristic的規則來進行指定，此外，阿里巴巴PAI團隊也會引入一個在線反饋tuning的機制，來對算子的FP16-friendliness數據進行自動采集，以期盡可能減少人力負擔。

3.2 Cost Model

但不是white-list中的所有算子都應該轉換為FP16，計算圖中每個節點需不需要轉換，除了和這個計算節點本身的算子類型有關，還和這個節點周圍其他節點等因素緊密相關，所以需要一個Cost Model來預測每個節點轉換為FP16之后對整體性能的影響，從而作出轉換與否的判斷。

舉個例子，前面提到的embedding_lookup算子，雖然也支持FP16，但是通過阿里巴巴PAI團隊的cost model可以輕易地自動將計算圖中的 embedding_lookup節點排除在外。如下圖所示，上側是原始計算圖中的 embedding_lookup節點，其訪存開銷是 2*N*D*4Bytes，其中系數2表示一次讀和一次寫，4Bytes是因為讀寫都是FP32數據；下側則是將 embedding_lookup節點轉換為FP16之后的示意圖，轉換后的訪存開銷則變為2*N*D*2+V*D*4+V*D*2Bytes，其中第一項表示embedding_lookup本身的讀寫量減半，而增加的后兩項則是因為在embedding_lookup之前需要將embedding variable轉換為FP16. 通常情況下詞表大小V遠遠大于batch size N, 因此轉換后的總的訪存開銷遠遠大于轉換前，通過cost model就能夠判斷出這個轉換不應該發生。

3.3 FP16 Propagation

在計算圖層面看模型的視角和在TF代碼層面看模型的視角不太一樣，在計算圖層面看到的只是一個個計算節點和他們之間的連接關系，缺失了更宏觀的結構信息，比如哪些是前向計算、哪些是后向計算、以及哪些是Optimizer的計算（雖然在某種程度上可以通過節點的名字來判斷，但這依賴于計算圖構造過程中的命名規則，更多是為了debug目的，不足以作為本質的判斷依據，一方面因為用戶有自定義命名的自由度；另一方面某些計算圖優化Pass在添加計算節點時則不受計算圖構造的命名規則的約束）。

這個局限性給計算圖層面的自動混合精度轉換造成了一定的障礙，因為PAI團隊并不希望Optimizer本身的計算也被轉換為FP16，原因是因為部分模型需要loss scaling才能防止gradients underflow ，Loss Scaling包含Scaling與Unscaling兩個步驟，降低在兩個動作之間的梯度值轉換為FP16后Underflow的風險，而Optimizer本身的計算位于Unscaling之后，如果Optimizer本身的計算也被轉換為FP16，那么就存在很大的underflow風險。正如前面所說，在計算圖層面并不能完全確認哪些計算是optimizer本身的。

為了解決這個問題，阿里巴巴PAI團隊提出并實現了FP16 Propagation方案。基本思想是，計算圖中節點的FP16轉換順序并不是隨意的，而是從MatMul/Conv2D節點開始向下游傳播，那么只要在gradient unscaling的位置(使團隊自己添加的Mul節點）做一個特殊標記，FP16的傳播遇到這個標記便停止，從而保證gradient unscaling之后的節點都不會被轉換為FP16。

3.4 Cast Elimination & Cast Fusion

計算圖層面的自動細粒度還存在一個不可避免的問題是會產生大量的Cast轉換節點，因為將每個節點轉換為FP16必然會在輸入端插入FP32轉FP16的Cast, 輸出端插入FP16轉FP32的Cast，這些Cast會帶來很大的額外開銷，因此阿里巴巴PAI團隊通過Cast Elimination 和 Cast Fusion兩個方法來降低所引入的這部分開銷。

Cast Elimination就是將兩個連續的相反的Cast節點互相抵消，如下圖所示。這種情況在FP16轉換后大量存在，如果沒有Cast Elimination，這些Cast節點帶來的額外開銷大部分情況下足以抵消混合精度訓練所帶來的性能收益。

雖然Cast Elimination已經能夠消除大部分Cast節點，但并不能完全消除，對于剩余的Cast節點可以進一步通過Cast Fusion使其和前后的節點融合，從而避免帶來額外的訪存開銷和Kernel Launch開銷。

4. 自動Loss Scaling的實現

4.1 loss scaling介紹

SSD訓練過程中的gradients分布

由于數值精度的下降，模型的訓練過程可能會有一定的精度下降或者不收斂情況。由于FP16的表示范圍有限，數值較小的gradients可能會出現underflow，導致數值直接被截斷為0，如上圖，圖中所示為一個經典的SSD圖像檢測模型在訓練過程中的gradients分布圖，圖中紅線以左都會因為underflow的問題在FP16下變為0，而同時觀察到FP16有很大一部分表征能力并沒有用到，因此可以將gradients放大一定的倍數進行表達。對于此問題可以通過loss scaling策略進行解決，即在loss上乘以一個scaling factor，在gradients進行apply之前進行unscale即可。

下圖所示即為loss scaling的計算流程示意圖，在loss位置，將loss乘以scaling factor，進行后向計算，在得到FP32 gradients進行unscaling操作，保證gradients數值的準確性。其中對于scaling factor，阿里巴巴PAI團隊實際是引入了額外的超參數，若其值較大，則在后向計算中，gradients可能出現overflow的情況，若其值較小，則依然不能很好的解決FP16下underflow的問題。對于它的控制，NVIDIA提出了auto scaling策略，其主要想法是，在不溢出的情況下，可以使用一個盡量大的scaling factor。

Loss scaling實現流程框架圖

4.2 Loss scaling的自動實現

對于loss scaling策略，用戶可以用類似

grads = [grad / scale for grad in tf.gradients(loss * scale, params)]

的方式進行手動實現，但手動實現存在以下兩方面的劣勢：

用戶的code可能有多種方式實現梯度的計算，例如可能直接調用tf.gradients() ，也可能調用optimizer.compute_gradients()，甚至也可能是調用optimizer.minimize()進行計算。在不同的策略下，用戶需要找到gradients的實際計算位置，并對其進行loss scaling計算操作；

若使用constant loss scaling，則用戶需要根據經驗指定scaling factor, 則有可能引入了額外的調參過程；而用戶若希望使用auto loss scaling的策略，則用戶需要理解auto loss scaling的完整策略，并手動實現，這個過程使得用戶需要對code進行一個不小的改動量，帶來的overhead是比較大的，且存在一定的出錯概率。

目前社區tensorflow和NVIDIA所給的示例中給出了另外一種實現方式，即對用戶的optimizer進行一層混合精度的wrapper封裝，并重新實現了compute_gradients()/apply_gradients()/minimize()等函數。但這種實現方式存在以下問題：

用戶code對于gradients的計算，不一定是通過調用optimizer.compute_gradients()實現的，也可能是直接調用tf.gradients()，在這種情況下用戶就需要改動自己的code以適用于混合精度的loss scaling訓練策略，這對用戶的overhead是比較大的；

wrapper封裝方式無法解耦與其他optimizer wrapper的耦合依賴；例如，在分布式訓練時，用戶一般會用分布式optimizer對于標準optimizer進行一層封裝，例如opt = tf.train.SyncReplicasOptimizer(opt)，此時用戶可能存在不知道封裝哪一個optimizer的困惑，并且目前社區的code是沒有對于分布式訓練wrapper與混合精度訓練wrapper進行解耦的，社區目前的版本與tf.train.SyncReplicasOptimizer耦合使用時會出現死鎖情況；

更細節地，wrapper方式不能實現對于colocate_gradients_with_ops等一些device placement控制功能的支持。

因此，為了解放用戶，阿里巴巴PAI團隊采取了另外一種實現方式，即對tf.gradients()進行了decorator，在這種方式下，可以支持optimizer.compute_gradients()/optimizer.minimize()/tf.gradients()任一方式下的梯度計算，同時是對梯度反向計算的原子操作tf.gradients()進行了decorator，因此它與其他的optimizer wrapper不存在耦合情況。細節地，用戶在使用混合精度時，若希望啟動loss scaling只需要將optimizer相關的部分都放在with gradients.mixed_precision_scope下即可，例如如下所示。

4.3 模型收斂性驗證

對于混合精度訓練，阿里巴巴PAI團隊基于不同類型的model進行了收斂性驗證。實驗結果可以驗證部分model在混合精度訓練下也可保持與FP32訓練結果相當的模型精度，部分model出現了精度下降或訓練不收斂的情況，但均可通過loss scaling的策略恢復模型精度，保證其與FP32下一致的模型精度。

4.3.1 Classification CNNs

對于分類CNN網絡，阿里巴巴PAI團隊基于ResNet50在ImageNet上進行了收斂性分析，下圖所示為訓練過程的test accuracy曲線，對比實驗包括:

FP32 baseline;

混合精度訓練(MP)+no loss scaling;

混合精度訓練(MP)+constant scaling (scaling factor=64);

混合精度訓練(MP)+auto scaling。 由對比可見，即使不加loss scaling策略，混合精度訓練的精度與FP32下基本無差別。

ResNet50訓練曲線對比

4.3.2 Detection CNNs

對于detection CNNs，阿里巴巴PAI團隊在Faster RCNN和SSD上進行了收斂性驗證：

Faster RCNN on PASCAL VOC2007

對于此model，在混合精度訓練下，不加loss scaling的情況下，mAP有1.65的drop，但在constant scaling及auto scaling的策略下，此精度均可recover，且精度略高于FP32情況下的模型精度

SSD on PASCAL VOC2007+2012

SSD模型在混合精度訓練情況下，不加loss scaling時出現了不收斂的情況，但在constant scaling及auto scaling情況下，其收斂精度均可與FP32 baseline相當。

4.3.3 NMT

對于NMT任務，阿里巴巴PAI團隊使用了small NMT模型在WMT German-English任務上進行了混合精度訓練，在auto scaling策略下，其所得BLEU值與FP32情況下訓練結果相當。

4.3.4 GAN

對于GAN的收斂性驗證，阿里巴巴PAI團隊在基于公開model Prograssive GAN (PGAN)和阿里巴巴PAI內部推進的一個AI字體生成業務model進行了驗證。

下圖所示從左至右分別為兩個模型在FP32、FP16(no loss scaling)以及FP16(auto scaling)情況下的生成圖片，其主觀結果無明顯差異。

4.3.5 Wide & Deep Learning

對于搜索、推薦、廣告領域常用的WDL模型，阿里巴巴PAI團隊基于Tensor flow官方案例 進行了收斂性驗證，實驗結果如下。即使在不加loss scaling的情況下，團隊觀察到混合精度的訓練結果與FP32也是相當的。

5. 線下性能評測

5.1 Language Model

首先選取了Tensorflow官方的Language Model進行評測：

評測模型提供了FP16訓練的Option --use_fp16，這屬于前面提到的模型代碼層面的手工轉換，正好可以和PAI自動混合精度訓練做一些性能對比。

收斂性：無需Loss Scaling, AutoMixedPrecision收斂性與FP32訓練無差異。

性能：如下圖所示(單位words/sec)，PAI自動混合進度訓練的性能明顯超過手工轉換的FP16模型的性能，相對于V100 FP32性能基準，加速達到2.3X。

5.2 NMT Transformer

同時阿里巴巴PAI團隊針對阿里達摩院機器翻譯團隊的In-house Production Model也進行了混合精度的加速效果驗證。

收斂性：加上Loss Scaling, AutoMixedPrecision收斂性與FP32訓練無差異

性能：如下圖所示(單位step/sec), 這個模型有兩種配置，hidden size = 512的配置相對于V100 FP32加速比為1.6X， hidden size = 1024的配置加速比為2.1X。

6. 線上性能對比

首先要說明的一點是，雖然已經有多個業務的多個模型在線上使用PAI混合精度訓練，但并不是每個模型的每個配置阿里巴巴PAI團隊都能拿到和非混合精度訓練的性能對比，因為用戶在順利使用PAI混合精度訓練后并沒有再使用非混合精度訓練的需求，有時候用戶跑非混合精度訓練純粹是為了幫阿里巴巴PAI團隊做性能對比，團隊也不愿意過多地浪費用戶寶貴的時間和集群計算資源來做性能對比。所以下面會比較細節地列舉兩個項目各自其中一個模型的性能對比，再給出一些代表性的模型性能對比。

6.1 高德某導航類相關項目

該項目是PAI自動混合精度訓練的第一個用戶，該項目自研的模型，既包括多個卷積層也包括多個全連接層。

該項目最新開發的模型使用PAI混合精度訓練，模型收斂性沒有收到影響。

使用PAI混合精度訓練和V100非混合精度訓練的性能對比，如下表所示（單位sec/step）

對于小尺寸模型，在相同的V100 GPU上，PAI混合精度訓練相比于非混合精度訓練加速2.3X。

對于正常尺寸的模型，如果不用混合精度訓練，在V100 GPU上則因為OOM根本無法訓練，這也體現了PAI混合精度訓練除了性能上的加速，也能降低顯存的消耗。

6.2 速賣通 (AliExpess)某商品數據挖掘項目

AliExpess商品數據挖掘項目是PAI自動混合精度訓練的第二個用戶，使用PAI混合精度訓練，同樣對模型收斂性沒有任何影響。

和前面的高德導航類項目新開發的模型直接在V100集群使用PAI混合精度訓練不同的是，該項目的模型在使用PAI混合精度訓練之前已經在使用P100集群訓練，切換到PAI混合精度訓練后獲得了非常顯著的加速，一方面是V100 GPU相對于P100 GPU的性能提升，另一方面是在V100 GPU上PAI混合精度訓練帶來的性能提升。訓練時間對比如下圖所示

對于小尺寸模型，在相同的V100 GPU上，PAI混合精度訓練相比于非混合精度訓練加速2.3X。

對于正常尺寸的模型，如果不用混合精度訓練，在V100 GPU上則因為OOM根本無法訓練，這也體現了PAI混合精度訓練除了性能上的加速，也能降低顯存的消耗。

6.3 PAI自動混合精度代表性模型加速示例

因為自動混合精度功能在設計上就充分考慮到了通用性及透明性，所以業務多樣性支持得比較出色，下面展示的部分數據包括PAI平臺對內業務CNN、Transformer、BERT以及FCN占模型結構主體的不同模型，業務場景橫跨了NLP、圖像、電商數據挖掘、高德交通導航等。