大家好，我是Tim。

自從GPT模型誕生以來，其參數(shù)規(guī)模就在不停的擴(kuò)大。但模型并非簡單的直接變大，需要在數(shù)據(jù)、調(diào)度、并行計(jì)算、算法和機(jī)器資源上做相應(yīng)的改變。

今天就來總結(jié)下，什么是大模型，模型變大的難在哪里以及對(duì)于CV/NLP或者搜推廣場景上有什么應(yīng)對(duì)策略。

什么是大模型？

大模型，顧名思義主打的就是“大”。主要包括幾個(gè)方面：

1. 數(shù)據(jù)規(guī)模大，通過大量的數(shù)據(jù)提高模型的泛化能力和性能。
2. 大規(guī)模并行計(jì)算能力，隨著計(jì)算硬件的不斷進(jìn)步，如GPU和TPU的普及，大規(guī)模并行計(jì)算能力的提升使得訓(xùn)練和推理大模型成為可能。
3. 更“大”模型復(fù)雜性：大模型具備更深層次、更復(fù)雜的網(wǎng)絡(luò)結(jié)構(gòu)，可以捕捉更豐富的特征和關(guān)系，提高了模型的表達(dá)能力。

大模型主要分為兩類：一個(gè)是稀疏大模型，另一個(gè)是稠密大模型。

1. 稀疏大模型：稀疏大模型是指模型中存在大量稀疏參數(shù)的情況，一般是搜索、推薦、廣告類任務(wù)。它的特點(diǎn)是海量樣本及大規(guī)模稀疏參數(shù)（sparse embeddings），適合使用 CPU/GPU 參數(shù)服務(wù)器模式（PS）進(jìn)行訓(xùn)練。
2. 稠密大模型：稠密大模型是指模型中的參數(shù)大多數(shù)都是非零值，沒有明顯的稀疏性特征，一般是CV、NLP 任務(wù)。它的特點(diǎn)是常規(guī)樣本數(shù)據(jù)及大規(guī)模稠密參數(shù)，它適合用純 GPU 集合通信模式（Collective）進(jìn)行訓(xùn)練。

對(duì)于搜推廣類的稀疏大模型來說，一般包含稀疏特征的嵌入（embedding）和稠密模型兩個(gè)部分。

其中，稀疏特征的嵌入計(jì)算是稀疏大模型的關(guān)鍵，而稠密模型部分一般往往較小，可以放到一個(gè)GPU內(nèi)，因此可以進(jìn)行data并行以及all reduce通訊。

在訓(xùn)練中，需要在特征嵌入表（embedding table）上需要進(jìn)行復(fù)雜的查找、排列等操作，然后生成張量再做稠密模型的計(jì)算。

而特征嵌入表往往會(huì)占用非常大的存儲(chǔ)空間，需要很多臺(tái)GPU服務(wù)器才能完整存放，這就是典型的tensor并行。

在這樣的場景下，就會(huì)導(dǎo)致典型的alltoall的通訊模式，而alltoall通訊會(huì)帶來嚴(yán)重的incast通訊（多打一），進(jìn)而帶來網(wǎng)絡(luò)擁塞，給網(wǎng)絡(luò)架構(gòu)、擁塞控制協(xié)議、負(fù)載均衡算法等都提出了很高的要求。

對(duì)于CV、NLP 任務(wù)來說，由于模型參數(shù)非常多，遠(yuǎn)遠(yuǎn)超過了單個(gè)GPU顯存所能容納的空間（NVIDIA最新的A100也就是80GB顯存）。

所以往往既需要對(duì)模型某一層的tensor并行，也需要不同層之間的pipeline并行，才能放得下整個(gè)大模型。

在計(jì)算過程中，既有單機(jī)內(nèi)通訊，也有不同機(jī)器間的通訊，具體的通訊模式取決于模型的切分和放置方法。

為了加速訓(xùn)練過程，往往完整的大模型之間也會(huì)采用data并行，每一個(gè)完整的大模型會(huì)被投喂不同的訓(xùn)練數(shù)據(jù)，這就會(huì)導(dǎo)致大家熟悉的allreduce通信模式。

總結(jié)起來，稠密大模型和稀疏大模型在模型特征上有著明顯的差異，對(duì)計(jì)算/存儲(chǔ)/通信資源的需求也存在明顯的不同。

要達(dá)到GPU算力資源的最大化利用和最好的加速效果，需要結(jié)合模型特征和實(shí)現(xiàn)方式對(duì)GPU服務(wù)器架構(gòu)、網(wǎng)絡(luò)架構(gòu)、訓(xùn)練數(shù)據(jù)存儲(chǔ)和拉取、分布式訓(xùn)練框架進(jìn)行全局的考量和設(shè)計(jì)。

模型大了難在哪里？

大模型帶來的挑戰(zhàn)主要有兩點(diǎn)：海量樣本、參數(shù)（萬億級(jí)別）和較長的收斂時(shí)間。

別看只有這區(qū)區(qū)兩點(diǎn)，它會(huì)衍生出很多要解決的問題。

1. 大模型訓(xùn)練需要更大的算力。

大模型訓(xùn)練所需的總算力，其實(shí)很簡單，6 * 模型的參數(shù)量 * 訓(xùn)練數(shù)據(jù)的 token 數(shù)就是所有訓(xùn)練數(shù)據(jù)過一遍所需的算力。這里的 6 就是每個(gè) token 在模型正向傳播和反向傳播的時(shí)候所需的乘法、加法計(jì)算次數(shù)。

一堆矩陣相乘，簡單來想就是左邊若干個(gè)神經(jīng)元，右邊若干個(gè)神經(jīng)元，組成一個(gè)完全二分圖。選出其中任意一個(gè)左邊的神經(jīng)元 L 和右邊的神經(jīng)元 R。

• 正向傳播的時(shí)候： L把它的輸出乘上 L和 R 之間的權(quán)重 w，發(fā)給 R；R不可能只連一個(gè)神經(jīng)元吧，總要把多個(gè) L的加到一起，這就是 reduce，需要一次加法。
• R把它收到的梯度乘上 L和 R 之間的權(quán)重 w，發(fā)給 L；L也不可能只連一個(gè) R，需要把梯度 reduce 一下，做個(gè)加法；別忘了權(quán)重 w 需要更新，那就要計(jì)算 w 的梯度，把 R 收到的梯度乘上 L正向傳播的輸出（activation）；一個(gè) batch 一般有多個(gè) sample，權(quán)重 w 的更新需要把這些 sample 的梯度加到一起。
• 一共 3 次乘法，3 次加法，不管 Transformer 多復(fù)雜，矩陣計(jì)算就是這么簡單，其他的向量計(jì)算、softmax 之類的都不是占算力的主要因素，估算的時(shí)候可以忽略。

有了模型訓(xùn)練所需的總算力，除以每個(gè) GPU 的理論算力，再除以 GPU 的有效算力利用比例，就得到了所需的 GPU-hours。

如果訓(xùn)練一個(gè)通用大語言模型的基座需要半年或幾個(gè)月的時(shí)間，同時(shí)還占用非常多的機(jī)器資源，這就使得大模型的訓(xùn)練非常“貴”，使得大模型的訓(xùn)練成為了個(gè)別大企業(yè)的專用。

例如，如果你有100個(gè)實(shí)驗(yàn)想試試，而模型訓(xùn)練需要半年，那你只能在其中選擇優(yōu)先級(jí)高進(jìn)行實(shí)驗(yàn)。

這也就是為什么現(xiàn)在大家的大語言模型基座都是追蹤最新開源的模型。因?yàn)樽约焊阃ㄓ么竽Ｐ突环矫娓惆胩斐杀旧虾牟黄穑硪环矫婵赡苓€沒啥效果。

2. 大模型訓(xùn)練需要更多的顯存內(nèi)存資源。

深度學(xué)習(xí)訓(xùn)練需要的內(nèi)存包括模型參數(shù)、反向傳播的梯度、優(yōu)化器所用的內(nèi)存、正向傳播的中間狀態(tài)（activation），顯存占用 = 模型參數(shù)大小 + Batch Size * 優(yōu)化器參數(shù)與中間變量的大小。

• 優(yōu)化器所用的內(nèi)存的計(jì)算其實(shí)也很簡單，如果用最經(jīng)典的 Adam 優(yōu)化器，它需要用 32 位浮點(diǎn)來計(jì)算。即使我們使用mixed-precision進(jìn)行計(jì)算，每個(gè)參數(shù)需要也要存 4 字節(jié)的 32 位版本（正向傳播時(shí)用 16 位版本，優(yōu)化時(shí)用 32 位版本），還需要存 4 字節(jié)的 momentum 和 4 字節(jié)的 variance，一共 12 字節(jié)。如果是用類似 SGD 的優(yōu)化器，可以不存 variance，只需要 8 字節(jié)。
• 正向傳播的中間狀態(tài)（activation）是反向傳播時(shí)計(jì)算梯度必需的，而且跟 batch size 成正比。Batch size 越大，每次讀取模型參數(shù)內(nèi)存能做的計(jì)算就越多，這樣對(duì) GPU 內(nèi)存帶寬的壓力就越小。劃重點(diǎn)：正向傳播的中間狀態(tài)數(shù)量是跟 batch size 成正比的。

當(dāng)然也有節(jié)省內(nèi)存資源的辦法，例如算力換內(nèi)存，時(shí)間換內(nèi)存等。

算力換內(nèi)存的把戲，就是不要存儲(chǔ)那么多梯度和每一層的正向傳播的中間狀態(tài)，而是在計(jì)算到某一層的時(shí)候再臨時(shí)從頭開始重算正向傳播的中間狀態(tài)。

如果每一層都這么干，那么就只要 2 個(gè)字節(jié)來存這一層的梯度，但是計(jì)算中間狀態(tài)的算力開銷會(huì)很大。

因此實(shí)際中一般是把整個(gè) Transformer 分成若干組，一組有若干層，只保存每組第一層的中間狀態(tài)，后面的層就從該組第一層開始重新計(jì)算，這樣就平衡了算力和內(nèi)存的開銷。

時(shí)間換內(nèi)存的把戲，按順序執(zhí)行Mini-Batch數(shù)據(jù)的前向計(jì)算梯度，同時(shí)對(duì)梯度進(jìn)行累積，累積的結(jié)果在最后一個(gè)Mini-Batch計(jì)算后求平均更新模型變量。

此外，還可以多級(jí)緩存，GPU 內(nèi)存放不下可以換出到 CPU 內(nèi)存。

例如，對(duì)于 LLaMA-2 70B 模型，模型參數(shù)需要 140 GB，反向傳播的梯度需要 140 GB，優(yōu)化器的狀態(tài)（如果用 Adam）需要 840 GB。

3. 對(duì)數(shù)據(jù)的數(shù)量和數(shù)據(jù)的質(zhì)量要求極高

對(duì)于海量數(shù)據(jù)樣本來說，并不是都喂進(jìn)去就效果好，哪些數(shù)據(jù)有價(jià)值，哪些沒價(jià)值。由于數(shù)據(jù)量的增加，分辨數(shù)據(jù)的價(jià)值也帶來很大的困難。

此外，大量的數(shù)據(jù)存儲(chǔ)在哪里，一般可以存儲(chǔ)在HDFS或S3。但怎么保證存取能不把機(jī)器塞滿且能快速調(diào)取，對(duì)于相同數(shù)據(jù)的模型多次訓(xùn)練，是否可以通過cache來加速模型訓(xùn)練的時(shí)間。

4. Transformer的“不可能三角”特性

圖片

Transformer在訓(xùn)練并行性、推理效率和競爭性能之間很難取得平衡被稱為“不可能三角”。

Transformers 由于其固有的每步 O(N) 復(fù)雜度和內(nèi)存限制的鍵值緩存，在推理過程中表現(xiàn)出次優(yōu)效率。這種低效率使它們的實(shí)際部署變得復(fù)雜，特別是對(duì)于長序列。

5. 深度學(xué)習(xí)框架上不是很友好

因?yàn)檫@個(gè)領(lǐng)域最近幾年才開始熱門，而之前的框架pytorch、tensorflow等是早就出現(xiàn)的，當(dāng)時(shí)并沒有針對(duì)大模型的分布式訓(xùn)練的需求場景做深入的抽象設(shè)計(jì)和優(yōu)化。

所以這個(gè)領(lǐng)域需要通過大數(shù)據(jù)框架到深度學(xué)習(xí)框架的端到端打通，形成這樣的一套新的編程范式和對(duì)應(yīng)的計(jì)算框架來解決掉。

下面我們以CV、NLP場景和搜推廣場景進(jìn)行分別詳細(xì)說明。

CV和NLP場景：

對(duì)CV和NLP場景來說，其特點(diǎn)主要有：

• 模型一般復(fù)雜，單機(jī)性能要求高。業(yè)界主要使用高性能的GPU進(jìn)行計(jì)算，并采用All-reduce的通信拓?fù)溥M(jìn)行參數(shù)的同步更新。
• 模型大（DenseNet部分），比如NLP領(lǐng)域，GPT-3這樣的模型高達(dá)1750億參數(shù)，顯存占用高達(dá)2.8 TB，單機(jī)內(nèi)存無法容納。

當(dāng)面對(duì)GPT-3這種Dense部分大的模型，Allreduce 單卡內(nèi)存無法容納，我們需要采用模型并行(model parallelism)的方式將計(jì)算圖劃分到不同的設(shè)備上構(gòu)建有向無環(huán)圖(DAG)進(jìn)行分布式訓(xùn)練，其中Gpipe, Megatron, Oneflow和Whale都提出模型并行的相關(guān)解決方案。

相比于數(shù)據(jù)并行每個(gè)worker只有一部分?jǐn)?shù)據(jù)，模型并行下每個(gè)node使用所有數(shù)據(jù)。

下面我們簡單說明幾種模型并行的方法：

• Tensor Parallelism，主要是將一層Layer中的矩陣計(jì)算分別拆分到不同的機(jī)器上進(jìn)行運(yùn)算，比如1D Megatron。
• Pipeline Parallelism，會(huì)將模型的layers拆分到不同的機(jī)器上，則一次forward就需要跨過不同的機(jī)器串行地進(jìn)行運(yùn)算，而流行并行通過將batch size切分為更小的mirco batch，減少數(shù)據(jù)依賴，從而將整個(gè)計(jì)算過程異步起來，最大化資源利用率。

CTR推廣搜場景：

對(duì)于CTR大模型場景來說，其具有模型小，詞表大的特點(diǎn)。

1. 模型中的Dense部分，一般很小，往往一臺(tái)機(jī)器的內(nèi)存就可以容納。但是其特征量級(jí)可能高達(dá)成百上千億，造成Sparse部分或者Embedding table高達(dá)TB級(jí)別，使得單機(jī)無法容納。
2. 一個(gè)Batch的embedding lookup量級(jí)大，造成查詢耗時(shí)大。由于特征數(shù)量多，一個(gè)Batch可能包含幾十萬個(gè)ID類特征，TF原生的embedding lookup查詢耗時(shí)大，造成訓(xùn)練和inference性能低。
3. 數(shù)據(jù)具有大規(guī)模稀疏的特點(diǎn)。不同于CV和NLP場景，數(shù)據(jù)是稠密的圖像和文本，搜廣推的數(shù)據(jù)非常稀疏的，第一這來源于很多數(shù)據(jù)無法對(duì)所有用戶和場景有效采集到，第二是因?yàn)榻Ｊ褂玫奶卣髁考?jí)大造成的高維稀疏性，這會(huì)影響了數(shù)據(jù)的存儲(chǔ)格式和計(jì)算效率。

因此，解決CTR大模型的這種稠密參數(shù)較大的模型，關(guān)鍵是將Sparse參數(shù)由單機(jī)存儲(chǔ)改造為分布式存儲(chǔ)，并主要通過數(shù)據(jù)并行提高吞吐。

下面我們說明下對(duì)訓(xùn)練框架優(yōu)化點(diǎn)。核心的兩點(diǎn)，一個(gè)在于分布式通信拓?fù)涞脑O(shè)計(jì)，還有一個(gè)在于Embedding Lookup的性能優(yōu)化。

1. 稀疏參數(shù)，借助參數(shù)服務(wù)器（Param Server），將 embedding 存儲(chǔ)和更新負(fù)擔(dān)轉(zhuǎn)嫁到PS。稀疏參數(shù) Partition 存放，每個(gè) Worker 只有部分分片，梯度更新時(shí)進(jìn)行 AlltoAll，想辦法解決 稀疏tensor 的存儲(chǔ)、通信成本。
2. 稠密參數(shù)，借助于allreduce，將稠密參數(shù) Replication 存放，每個(gè) Worker 都有副本，梯度更新時(shí)進(jìn)行 allreduce。allreduce 和 alltoall 都會(huì)使用 nccl 進(jìn)行同步通信，效率較高。hb 進(jìn)行 alltoall 時(shí)，通信的是稀疏梯度，而不是完整的參數(shù)，通信量上和 ps 是一致的，但是通信效率會(huì)更高。

在實(shí)現(xiàn)上，可以通過替換TF原生算子進(jìn)行Sparse參數(shù)的讀取過程（核心算子是GatherV2算子）。

該算子的作用是從Embedding表中讀取ID列表索引對(duì)應(yīng)的Embedding數(shù)據(jù)并返回，本質(zhì)上是一個(gè)Hash查詢的過程；

通過替換該算子，并同時(shí)修改上下游節(jié)點(diǎn)的Input/Output，將從本地Embedding表中查詢，改造為從分布式KV中查詢。

總結(jié)

可以說，訓(xùn)練大模型，不僅需要充足的計(jì)算資源和數(shù)據(jù)，還需要深厚經(jīng)驗(yàn)和技能，還需要一定的耐心和定力，就像“煉丹”一樣。

每次煉丹師的出手都有著巨大時(shí)間和經(jīng)濟(jì)成本，如何在最小成本下找到最優(yōu)解，就是我們一直在探索的。