大家好,我是Tim。
自從GPT模型誕生以來,其參數規模就在不停的擴大。但模型并非簡單的直接變大,需要在數據、調度、并行計算、算法和機器資源上做相應的改變。
今天就來總結下,什么是大模型,模型變大的難在哪里以及對于CV/NLP或者搜推廣場景上有什么應對策略。
什么是大模型?
大模型,顧名思義主打的就是“大”。主要包括幾個方面:
- 數據規模大,通過大量的數據提高模型的泛化能力和性能。
- 大規模并行計算能力,隨著計算硬件的不斷進步,如GPU和TPU的普及,大規模并行計算能力的提升使得訓練和推理大模型成為可能。
- 更“大”模型復雜性:大模型具備更深層次、更復雜的網絡結構,可以捕捉更豐富的特征和關系,提高了模型的表達能力。
大模型主要分為兩類:一個是稀疏大模型,另一個是稠密大模型。
- 稀疏大模型:稀疏大模型是指模型中存在大量稀疏參數的情況,一般是搜索、推薦、廣告類任務。它的特點是海量樣本及大規模稀疏參數(sparse embeddings),適合使用 CPU/GPU 參數服務器模式(PS)進行訓練。
- 稠密大模型:稠密大模型是指模型中的參數大多數都是非零值,沒有明顯的稀疏性特征,一般是CV、NLP 任務。它的特點是常規樣本數據及大規模稠密參數,它適合用純 GPU 集合通信模式(Collective)進行訓練。
對于搜推廣類的稀疏大模型來說,一般包含稀疏特征的嵌入(embedding)和稠密模型兩個部分。
其中,稀疏特征的嵌入計算是稀疏大模型的關鍵,而稠密模型部分一般往往較小,可以放到一個GPU內,因此可以進行data并行以及all reduce通訊。
在訓練中,需要在特征嵌入表(embedding table)上需要進行復雜的查找、排列等操作,然后生成張量再做稠密模型的計算。
而特征嵌入表往往會占用非常大的存儲空間,需要很多臺GPU服務器才能完整存放,這就是典型的tensor并行。
在這樣的場景下,就會導致典型的alltoall的通訊模式,而alltoall通訊會帶來嚴重的incast通訊(多打一),進而帶來網絡擁塞,給網絡架構、擁塞控制協議、負載均衡算法等都提出了很高的要求。
對于CV、NLP 任務來說,由于模型參數非常多,遠遠超過了單個GPU顯存所能容納的空間(NVIDIA最新的A100也就是80GB顯存)。
所以往往既需要對模型某一層的tensor并行,也需要不同層之間的pipeline并行,才能放得下整個大模型。
在計算過程中,既有單機內通訊,也有不同機器間的通訊,具體的通訊模式取決于模型的切分和放置方法。
為了加速訓練過程,往往完整的大模型之間也會采用data并行,每一個完整的大模型會被投喂不同的訓練數據,這就會導致大家熟悉的allreduce通信模式。
總結起來,稠密大模型和稀疏大模型在模型特征上有著明顯的差異,對計算/存儲/通信資源的需求也存在明顯的不同。
要達到GPU算力資源的最大化利用和最好的加速效果,需要結合模型特征和實現方式對GPU服務器架構、網絡架構、訓練數據存儲和拉取、分布式訓練框架進行全局的考量和設計。
模型大了難在哪里?
大模型帶來的挑戰主要有兩點:海量樣本、參數(萬億級別)和較長的收斂時間。
別看只有這區區兩點,它會衍生出很多要解決的問題。
1. 大模型訓練需要更大的算力。
大模型訓練所需的總算力,其實很簡單,6 * 模型的參數量 * 訓練數據的 token 數就是所有訓練數據過一遍所需的算力。這里的 6 就是每個 token 在模型正向傳播和反向傳播的時候所需的乘法、加法計算次數。
一堆矩陣相乘,簡單來想就是左邊若干個神經元,右邊若干個神經元,組成一個完全二分圖。選出其中任意一個左邊的神經元 L 和右邊的神經元 R。
- 正向傳播的時候: L把它的輸出乘上 L和 R 之間的權重 w,發給 R;R不可能只連一個神經元吧,總要把多個 L的加到一起,這就是 reduce,需要一次加法。
- R把它收到的梯度乘上 L和 R 之間的權重 w,發給 L;L也不可能只連一個 R,需要把梯度 reduce 一下,做個加法;別忘了權重 w 需要更新,那就要計算 w 的梯度,把 R 收到的梯度乘上 L正向傳播的輸出(activation);一個 batch 一般有多個 sample,權重 w 的更新需要把這些 sample 的梯度加到一起。
- 一共 3 次乘法,3 次加法,不管 Transformer 多復雜,矩陣計算就是這么簡單,其他的向量計算、softmax 之類的都不是占算力的主要因素,估算的時候可以忽略。
有了模型訓練所需的總算力,除以每個 GPU 的理論算力,再除以 GPU 的有效算力利用比例,就得到了所需的 GPU-hours。
如果訓練一個通用大語言模型的基座需要半年或幾個月的時間,同時還占用非常多的機器資源,這就使得大模型的訓練非常“貴”,使得大模型的訓練成為了個別大企業的專用。
例如,如果你有100個實驗想試試,而模型訓練需要半年,那你只能在其中選擇優先級高進行實驗。
這也就是為什么現在大家的大語言模型基座都是追蹤最新開源的模型。因為自己搞通用大模型基座,一方面搞半天成本上耗不起,另一方面可能還沒啥效果。
2. 大模型訓練需要更多的顯存內存資源。
深度學習訓練需要的內存包括模型參數、反向傳播的梯度、優化器所用的內存、正向傳播的中間狀態(activation),顯存占用 = 模型參數大小 + Batch Size * 優化器參數與中間變量的大小。
- 優化器所用的內存的計算其實也很簡單,如果用最經典的 Adam 優化器,它需要用 32 位浮點來計算。即使我們使用mixed-precision進行計算,每個參數需要也要存 4 字節的 32 位版本(正向傳播時用 16 位版本,優化時用 32 位版本),還需要存 4 字節的 momentum 和 4 字節的 variance,一共 12 字節。如果是用類似 SGD 的優化器,可以不存 variance,只需要 8 字節。
- 正向傳播的中間狀態(activation)是反向傳播時計算梯度必需的,而且跟 batch size 成正比。Batch size 越大,每次讀取模型參數內存能做的計算就越多,這樣對 GPU 內存帶寬的壓力就越小。劃重點:正向傳播的中間狀態數量是跟 batch size 成正比的。
當然也有節省內存資源的辦法,例如算力換內存,時間換內存等。
算力換內存的把戲,就是不要存儲那么多梯度和每一層的正向傳播的中間狀態,而是在計算到某一層的時候再臨時從頭開始重算正向傳播的中間狀態。
如果每一層都這么干,那么就只要 2 個字節來存這一層的梯度,但是計算中間狀態的算力開銷會很大。
因此實際中一般是把整個 Transformer 分成若干組,一組有若干層,只保存每組第一層的中間狀態,后面的層就從該組第一層開始重新計算,這樣就平衡了算力和內存的開銷。
時間換內存的把戲,按順序執行Mini-Batch數據的前向計算梯度,同時對梯度進行累積,累積的結果在最后一個Mini-Batch計算后求平均更新模型變量。
此外,還可以多級緩存,GPU 內存放不下可以換出到 CPU 內存。
例如,對于 LLaMA-2 70B 模型,模型參數需要 140 GB,反向傳播的梯度需要 140 GB,優化器的狀態(如果用 Adam)需要 840 GB。
3. 對數據的數量和數據的質量要求極高
對于海量數據樣本來說,并不是都喂進去就效果好,哪些數據有價值,哪些沒價值。由于數據量的增加,分辨數據的價值也帶來很大的困難。
此外,大量的數據存儲在哪里,一般可以存儲在HDFS或S3。但怎么保證存取能不把機器塞滿且能快速調取,對于相同數據的模型多次訓練,是否可以通過cache來加速模型訓練的時間。
4. Transformer的“不可能三角”特性
圖片
Transformer在訓練并行性、推理效率和競爭性能之間很難取得平衡被稱為“不可能三角”。
Transformers 由于其固有的每步 O(N) 復雜度和內存限制的鍵值緩存,在推理過程中表現出次優效率。這種低效率使它們的實際部署變得復雜,特別是對于長序列。
5. 深度學習框架上不是很友好
因為這個領域最近幾年才開始熱門,而之前的框架pytorch、tensorflow等是早就出現的,當時并沒有針對大模型的分布式訓練的需求場景做深入的抽象設計和優化。
所以這個領域需要通過大數據框架到深度學習框架的端到端打通,形成這樣的一套新的編程范式和對應的計算框架來解決掉。
下面我們以CV、NLP場景和搜推廣場景進行分別詳細說明。
CV和NLP場景:
對CV和NLP場景來說,其特點主要有:
- 模型一般復雜,單機性能要求高。業界主要使用高性能的GPU進行計算,并采用All-reduce的通信拓撲進行參數的同步更新。
- 模型大(Dens.NET部分),比如NLP領域,GPT-3這樣的模型高達1750億參數,顯存占用高達2.8 TB,單機內存無法容納。
當面對GPT-3這種Dense部分大的模型,Allreduce 單卡內存無法容納,我們需要采用模型并行(model parallelism)的方式將計算圖劃分到不同的設備上構建有向無環圖(DAG)進行分布式訓練,其中Gpipe, Megatron, Oneflow和Whale都提出模型并行的相關解決方案。
相比于數據并行每個worker只有一部分數據,模型并行下每個node使用所有數據。
下面我們簡單說明幾種模型并行的方法:
- Tensor Parallelism,主要是將一層Layer中的矩陣計算分別拆分到不同的機器上進行運算,比如1D Megatron。
- Pipeline Parallelism,會將模型的layers拆分到不同的機器上,則一次forward就需要跨過不同的機器串行地進行運算,而流行并行通過將batch size切分為更小的mirco batch,減少數據依賴,從而將整個計算過程異步起來,最大化資源利用率。
CTR推廣搜場景:
對于CTR大模型場景來說,其具有模型小,詞表大的特點。
- 模型中的Dense部分,一般很小,往往一臺機器的內存就可以容納。但是其特征量級可能高達成百上千億,造成Sparse部分或者Embedding table高達TB級別,使得單機無法容納。
- 一個Batch的embedding lookup量級大,造成查詢耗時大。由于特征數量多,一個Batch可能包含幾十萬個ID類特征,TF原生的embedding lookup查詢耗時大,造成訓練和inference性能低。
- 數據具有大規模稀疏的特點。不同于CV和NLP場景,數據是稠密的圖像和文本,搜廣推的數據非常稀疏的,第一這來源于很多數據無法對所有用戶和場景有效采集到,第二是因為建模使用的特征量級大造成的高維稀疏性,這會影響了數據的存儲格式和計算效率。
因此,解決CTR大模型的這種稠密參數較大的模型,關鍵是將Sparse參數由單機存儲改造為分布式存儲,并主要通過數據并行提高吞吐。
下面我們說明下對訓練框架優化點。核心的兩點,一個在于分布式通信拓撲的設計,還有一個在于Embedding Lookup的性能優化。
- 稀疏參數,借助參數服務器(Param Server),將 embedding 存儲和更新負擔轉嫁到PS。稀疏參數 Partition 存放,每個 Worker 只有部分分片,梯度更新時進行 AlltoAll,想辦法解決 稀疏tensor 的存儲、通信成本。
- 稠密參數,借助于allreduce,將稠密參數 Replication 存放,每個 Worker 都有副本,梯度更新時進行 allreduce。allreduce 和 alltoall 都會使用 nccl 進行同步通信,效率較高。hb 進行 alltoall 時,通信的是稀疏梯度,而不是完整的參數,通信量上和 ps 是一致的,但是通信效率會更高。
在實現上,可以通過替換TF原生算子進行Sparse參數的讀取過程(核心算子是GatherV2算子)。
該算子的作用是從Embedding表中讀取ID列表索引對應的Embedding數據并返回,本質上是一個Hash查詢的過程;
通過替換該算子,并同時修改上下游節點的Input/Output,將從本地Embedding表中查詢,改造為從分布式KV中查詢。
總結
可以說,訓練大模型,不僅需要充足的計算資源和數據,還需要深厚經驗和技能,還需要一定的耐心和定力,就像“煉丹”一樣。
每次煉丹師的出手都有著巨大時間和經濟成本,如何在最小成本下找到最優解,就是我們一直在探索的。