震惊！我竟然在1080Ti上加载了一个35亿参数的模型（ZeRO, Zero Redundancy Optimizer）
背景
在最近幾年，雖然大規(guī)模預(yù)訓(xùn)練模型已經(jīng)越來(lái)越普遍，但是關(guān)于如何訓(xùn)練這些模型的內(nèi)容卻很少有人關(guān)注，一般都是一些財(cái)大氣粗的企業(yè)或?qū)嶒?yàn)室來(lái)訓(xùn)練大模型并發(fā)布，然后中小型企業(yè)以及高校來(lái)使用。即便如此也有一些門檻，受限于機(jī)器配置，可能效果更好的大模型并不能直接加載到顯卡中，或者是單機(jī)多卡希望可以通過(guò)分布式的方法進(jìn)行微調(diào)。
目前訓(xùn)練超大規(guī)模語(yǔ)言模型主要有兩條技術(shù)路線：TPU + XLA + TensorFlow 和 GPU + PyTorch + Megatron-LM + DeepSpeed。前者由Google主導(dǎo)，由于TPU和自家云平臺(tái)GCP深度綁定，對(duì)于非Google開(kāi)發(fā)者來(lái)說(shuō)， 只可遠(yuǎn)觀而不可把玩，后者背后則有NVIDIA、Meta、微軟等大廠加持，社區(qū)氛圍活躍，也更受到群眾歡迎。
這篇文章介紹的ZeRO方法，就是可以讓大模型變得更加親民，不管是預(yù)訓(xùn)練、微調(diào)還是部署，即使沒(méi)有A100、V100、P100之類的GPU服務(wù)器，也可以讓加載大模型不是夢(mèng)。
| 軟件 | DeepSpeed |
| 模型 | GPT2-3.5B-chinese |
本次實(shí)驗(yàn)的模型是由IDEA發(fā)布的 Wenzhong2.0-GPT2-3.5B-chinese，這是一個(gè)基于解碼器結(jié)構(gòu)的單向語(yǔ)言模型，使用100G中文常用數(shù)據(jù)，32個(gè)A100訓(xùn)練了28個(gè)小時(shí)，是目前最大的開(kāi)源GPT2中文大模型。接下來(lái)我們將在8塊1080Ti顯卡上微調(diào)這個(gè)大模型。
分布式模型訓(xùn)練
數(shù)據(jù)并行（Data Parallel, DP）
將模型復(fù)制到多張卡上，然后將數(shù)據(jù)切分到多個(gè)GPU上，并行訓(xùn)練，每一輪結(jié)束之后進(jìn)行參數(shù)同步。
數(shù)據(jù)并行是最常見(jiàn)的并行形式，因?yàn)樗芎?jiǎn)單。在數(shù)據(jù)并行訓(xùn)練中，數(shù)據(jù)集被分割成幾個(gè)碎片，每個(gè)碎片被分配到一個(gè)設(shè)備上。這相當(dāng)于沿批次維度對(duì)訓(xùn)練過(guò)程進(jìn)行并行化。每個(gè)設(shè)備將持有一個(gè)完整的模型副本，并在分配的數(shù)據(jù)集碎片上進(jìn)行訓(xùn)練。在反向傳播之后，模型的梯度將被全部減少，以便在不同設(shè)備上的模型參數(shù)能夠保持同步。
數(shù)據(jù)并行由于簡(jiǎn)單易實(shí)現(xiàn)，應(yīng)用最為廣泛，當(dāng)然這不表示它沒(méi)有”缺點(diǎn)“，每張卡都存儲(chǔ)一個(gè)模型，此時(shí)顯存就成了模型規(guī)模的天花板。如果我們能減少模型訓(xùn)練過(guò)程中的顯存占用，那不就可以訓(xùn)練更大的模型了？一個(gè)簡(jiǎn)單的觀察是，如果有2張卡，那么系統(tǒng)中就存在2份模型參數(shù)，如果有4張卡，那么系統(tǒng)中就存在4份模型參數(shù)，如果有N張卡，系統(tǒng)中就存在N份模型參數(shù)，但其實(shí)其中N-1份都是冗余的，我們有必要讓每張卡都存一個(gè)完整的模型嗎？系統(tǒng)中能否只有一個(gè)完整模型，每張卡都存 1/N 參數(shù)，卡數(shù)越多，每張卡的顯存占用越少，這樣越能訓(xùn)練更大規(guī)模的模型。
優(yōu)點(diǎn)：加快訓(xùn)練速度
缺點(diǎn)：模型大小上限為單個(gè)GPU顯存
模型并行（Model Parallel, MP）
在數(shù)據(jù)并行訓(xùn)練中，一個(gè)明顯的特點(diǎn)是每個(gè) GPU 持有整個(gè)模型權(quán)重的副本。這就帶來(lái)了冗余問(wèn)題。另一種并行模式是模型并行，即模型被分割并分布在一個(gè)設(shè)備陣列上。通常有兩種類型的并行：張量并行和流水線并行。張量并行是在一個(gè)操作中進(jìn)行并行計(jì)算，如矩陣-矩陣乘法。流水線并行是在各層之間進(jìn)行并行計(jì)算。因此，從另一個(gè)角度來(lái)看，張量并行可以被看作是層內(nèi)并行，流水線并行可以被看作是層間并行。
張量并行（Tensor Parallel, TP）
將模型參數(shù)切分成分塊矩陣，分配到不同的GPU上，參數(shù)更新時(shí)再進(jìn)行同步。
張量并行將一個(gè)大的張量進(jìn)行切分，然后分配到不同的GPU上，每個(gè)GPU只處理張量的一部分，并且只有在一些需要所有張量計(jì)算的時(shí)候才進(jìn)行聚合。
張量并行需要在不同設(shè)備之間傳遞結(jié)果，因此不建議在跨節(jié)點(diǎn)之間進(jìn)行張量并行，除非網(wǎng)絡(luò)速度非常快。
優(yōu)點(diǎn)：可以不受單張GPU顯存限制，訓(xùn)練更大的模型。
缺點(diǎn)：計(jì)算/通信效率低。
流水線并行（Pipeline Parallel, PP）
將模型按層進(jìn)行分組，然后分配到不同的GPU，在GPU之間進(jìn)行前向傳播和反向傳播。
原生的流水線并行是指將模型層組劃分到到多個(gè)GPU上，然后只需要簡(jiǎn)單地將數(shù)據(jù)從一個(gè)GPU移動(dòng)到另一個(gè)GPU。流水線并行的實(shí)現(xiàn)相對(duì)比較簡(jiǎn)單，只需要將不同的層通過(guò).to()方法切換即可。
流水線并行的核心思想是，模型按層分割成若干塊，每塊都交給一個(gè)設(shè)備。在前向傳遞過(guò)程中，每個(gè)設(shè)備將中間的激活傳遞給下一個(gè)階段。在后向傳遞過(guò)程中，每個(gè)設(shè)備將輸入張量的梯度傳回給前一個(gè)流水線階段。這允許設(shè)備同時(shí)進(jìn)行計(jì)算，并增加了訓(xùn)練的吞吐量。流水線并行訓(xùn)練的一個(gè)缺點(diǎn)是，會(huì)有一些設(shè)備參與計(jì)算的冒泡時(shí)間，導(dǎo)致計(jì)算資源的浪費(fèi)。
優(yōu)點(diǎn)：層內(nèi)計(jì)算/通信效率增加
缺點(diǎn)：存在空閑等待時(shí)間
數(shù)據(jù)并行+流水線并行
有了這幾種并行方式之后，組合起來(lái)才能發(fā)揮更好的效果，數(shù)據(jù)并行和流水線并行相對(duì)來(lái)說(shuō)比較簡(jiǎn)單，先將數(shù)據(jù)分成兩份，放到兩個(gè)節(jié)點(diǎn)上，然后在每個(gè)節(jié)點(diǎn)的多個(gè)GPU上再執(zhí)行流水線并行。
數(shù)據(jù)并行+流水線并行+模型并行
為了再進(jìn)一步的提升效率，可以將三種并行方式都結(jié)合起來(lái)，稱之為3D并行，如下圖所示。
站在三維直角坐標(biāo)系的角度看，x表示的是同一層內(nèi)的張量并行，y軸表示的是同一節(jié)點(diǎn)內(nèi)的流水線并行，z軸表示的是不同節(jié)點(diǎn)間的數(shù)據(jù)并行。
DeepSpeed
ZeRO（Zero Redundancy Optimizer）
類似于張量并行進(jìn)行切分，支持多種offload技術(shù)。
目標(biāo)：優(yōu)化存儲(chǔ)效率的同時(shí)還能保持較高的計(jì)算和通信效率。
為了能夠在比較普通的機(jī)器上也能微調(diào)大模型，我們首先需要分析一下模型訓(xùn)練過(guò)程中都有哪些部分需要消耗存儲(chǔ)空間。在進(jìn)行深度學(xué)習(xí)訓(xùn)練的時(shí)候，有4大部分的顯存開(kāi)銷，分別是模型參數(shù)（Parameters），模型參數(shù)的梯度（Gradients），優(yōu)化器狀態(tài)（Optimizer States）以及中間激活值（Intermediate Activations）。
優(yōu)化模型占用空間
在訓(xùn)練過(guò)程中當(dāng)然模型占用的空間是最大的，但是現(xiàn)有的方法中，不管是數(shù)據(jù)并行DP還是模型并行MP都不能很好的解決。數(shù)據(jù)并行有很好的計(jì)算/通信效率，但是由于模型復(fù)制了多份，導(dǎo)致空間利用率很差，而模型并行雖然內(nèi)存利用率高，但是由于對(duì)模型的進(jìn)行了很精細(xì)的拆分，導(dǎo)致計(jì)算/通信效率很低。除此之外，所有這些方法都靜態(tài)保存了整個(gè)訓(xùn)練過(guò)程中所需的所有模型參數(shù)，但實(shí)際上并不是整個(gè)訓(xùn)練期間都需要這些內(nèi)容。
ZeRO-DP
基于上述問(wèn)題，提出了ZeRO-DP技術(shù)，即ZeRO驅(qū)動(dòng)的數(shù)據(jù)并行，兼顧數(shù)據(jù)并行的計(jì)算/通信效率和模型并行的空間效率。首先ZeRO-DP會(huì)對(duì)模型狀態(tài)進(jìn)行分區(qū)，避免了復(fù)制模型導(dǎo)致的冗余，然后在訓(xùn)練期間使用動(dòng)態(tài)通信調(diào)度保留數(shù)據(jù)并行的計(jì)算粒度和通信量，也能維持一個(gè)類似的計(jì)算/通信效率。
ZeRO-DP有三個(gè)優(yōu)化階段：① 優(yōu)化器狀態(tài)分區(qū)、② 梯度分區(qū)、③ 參數(shù)分區(qū)。
模型參數(shù)（fp16）、模型梯度（fp16）和Adam狀態(tài)（fp32的模型參數(shù)備份，fp32的momentum和fp32的variance）。假設(shè)模型參數(shù)量 Φ ，則共需要 2Φ+2Φ+(4Φ+4Φ+4Φ)=4Φ+12Φ=16Φ 字節(jié)存儲(chǔ)。
優(yōu)化剩余空間
優(yōu)化了模型的空間利用率之后，接下來(lái)就要優(yōu)化激活值、臨時(shí)緩沖區(qū)和不可用空間碎片。
ZeRO-R
ZeRO與模型并行
既然ZeRO解決了數(shù)據(jù)并行中空間利用率低的問(wèn)題，那么再進(jìn)一步，怎么與模型并行進(jìn)行配合使用呢？用了ZeRO之后，模型并行就僅僅被用在訓(xùn)練超大模型，ZeRO-DP在減少每個(gè)設(shè)備的空間占用方面跟模型并行一樣有效，或者當(dāng)模型并行不能均勻的劃分模型時(shí)有效。
在分布式模型訓(xùn)練過(guò)程中，數(shù)據(jù)并行非常易于使用，但是模型并行的門檻相對(duì)比較高，需要AI工程師手動(dòng)調(diào)整模型，甚至還需要系統(tǒng)開(kāi)發(fā)人員自定義分布式算子，當(dāng)然也有一些已有的工作，比如Megatron-LM。
混合精度訓(xùn)練（Mixed-Precision Training）
FP16
正常模型所使用的參數(shù)是Float浮點(diǎn)數(shù)類型，長(zhǎng)度為4個(gè)字節(jié)，也就是32位。在一些不太需要高精度計(jì)算的應(yīng)用中，比如圖像處理和神經(jīng)網(wǎng)絡(luò)中，32位的空間有一些浪費(fèi)，因此就又出現(xiàn)了一種新的數(shù)據(jù)類型，半精度浮點(diǎn)數(shù)，使用16位（2字節(jié)）來(lái)存儲(chǔ)浮點(diǎn)值，簡(jiǎn)稱FP16。
FP16的格式包括三部分：① 1 bit的符號(hào)位，② 5 bit的指數(shù)位寬，③ 11 bit的尾數(shù)精度（10位顯式存儲(chǔ)，1位隱式存儲(chǔ)）。
但是在大型語(yǔ)言模型中直接使用FP16會(huì)有一些問(wèn)題，如下圖所示，這是huggingface訓(xùn)練一個(gè)104B的模型時(shí)的loss，可以發(fā)現(xiàn)非常的不穩(wěn)定，除此之外Facebook最近發(fā)布的175B的OPT模型也報(bào)告了相同的問(wèn)題。
問(wèn)題在于FP16的指數(shù)位寬只有5 bit，因此可能會(huì)出現(xiàn)溢出的問(wèn)題，它能表示的最大整數(shù)就是65504，也就是說(shuō)一旦權(quán)重超過(guò)這個(gè)值就會(huì)溢出。
BF16
為了解決FP16的問(wèn)題，Google的人工智能研究小組就開(kāi)發(fā)了一種新的浮點(diǎn)數(shù)格式——BF16（Brain Floating Point），用于降低存儲(chǔ)需求，提高機(jī)器學(xué)習(xí)算法的計(jì)算速度。
BF16的格式包括三部分：① 1 bit的符號(hào)位，② 8 bit的指數(shù)位寬，③ 8 bit的尾數(shù)精度（7位顯式存儲(chǔ)，1位隱式存儲(chǔ)）。
模型參數(shù)不管是使用BF16還是FP16，在優(yōu)化器中更新的權(quán)重參數(shù)必須是用FP32的形式保存的，因此16位的浮點(diǎn)數(shù)僅用于計(jì)算。
ZeRO-Offload
ZeRO說(shuō)到底是一種數(shù)據(jù)并行方案，可是很多人只有幾張甚至一張卡，顯存加起來(lái)都不夠，那怎么辦呢？在操作系統(tǒng)中，當(dāng)內(nèi)存不足時(shí)，可以選擇一些頁(yè)面進(jìn)行換入換出，為新的數(shù)據(jù)騰出空間。類比一下，既然是因?yàn)轱@存不足導(dǎo)致一張卡訓(xùn)練不了大模型，那么ZeRO-Offload的想法就是：顯存不足，內(nèi)存來(lái)補(bǔ)。
在一個(gè)典型的服務(wù)器上，CPU 可以輕松擁有幾百GB的內(nèi)存，而每個(gè) GPU 通常只有16或32GB的內(nèi)存。相比于昂貴的顯存，內(nèi)存比較廉價(jià)，之前的很多工作都是聚焦在內(nèi)存顯存的換入換出，并沒(méi)有用到CPU的計(jì)算能力，也沒(méi)有考慮到多卡的場(chǎng)景。ZeRO-Offload則是將訓(xùn)練階段的某些模型狀態(tài)從GPU和顯存卸載到CPU和內(nèi)存。
當(dāng)然ZeRO-Offload并不希望為了最小化顯存占用而犧牲計(jì)算效率， 否則的話還不如直接使用CPU和內(nèi)存，因?yàn)榧词箤⒉糠諫PU的計(jì)算和顯存卸載到CPU和內(nèi)存，肯定要涉及到GPU和CPU、顯存和內(nèi)存的通信，而通信成本一般是非常高的，此外GPU的計(jì)算效率比CPU的計(jì)算效率高了好幾個(gè)數(shù)量積，因此也不能讓CPU參與過(guò)多的計(jì)算。
在ZeRO-Offload的策略中，將模型訓(xùn)練過(guò)程看作數(shù)據(jù)流圖，用圓形節(jié)點(diǎn)表示模型狀態(tài)，比如參數(shù)、梯度和優(yōu)化器狀態(tài)，用矩形節(jié)點(diǎn)表示計(jì)算操作，比如前向傳播、反向傳播和參數(shù)更新，邊表示數(shù)據(jù)流向。下圖是某一層的一次迭代過(guò)程（iteration/step），使用了混合精讀訓(xùn)練，前向計(jì)算（FWD）需要用到上一次的激活值（activation）和本層的參數(shù)（parameter），反向傳播（BWD）也需要用到激活值和參數(shù)計(jì)算梯度，
如果用Adam優(yōu)化器進(jìn)行參數(shù)更新（Param update），邊添加權(quán)重，物理含義是數(shù)據(jù)量大小（單位是字節(jié)），假設(shè)模型參數(shù)量是 M ，在混合精度訓(xùn)練的前提下，邊的權(quán)重要么是2M（fp16），要么是4M（fp32）。
**現(xiàn)在要做的就是沿著邊把數(shù)據(jù)流圖切分為兩部分，分布對(duì)應(yīng)GPU和CPU，**計(jì)算節(jié)點(diǎn)（矩形節(jié)點(diǎn)）落在哪個(gè)設(shè)備，哪個(gè)設(shè)備就執(zhí)行計(jì)算，數(shù)據(jù)節(jié)點(diǎn)（圓形）落在哪個(gè)設(shè)備，哪個(gè)設(shè)備就負(fù)責(zé)存儲(chǔ)，將被切分的邊權(quán)重加起來(lái)，就是CPU和GPU的通信數(shù)據(jù)量。
ZeRO-Offload的切分思路是：圖中有四個(gè)計(jì)算類節(jié)點(diǎn)：FWD、BWD、Param update和float2half，前兩個(gè)計(jì)算復(fù)雜度大致是 O(MB) ， B 是batch size，后兩個(gè)計(jì)算復(fù)雜度是 O(M) 。為了不降低計(jì)算效率，將前兩個(gè)節(jié)點(diǎn)放在GPU，后兩個(gè)節(jié)點(diǎn)不但計(jì)算量小還需要和Adam狀態(tài)打交道，所以放在CPU上，Adam狀態(tài)自然也放在內(nèi)存中，為了簡(jiǎn)化數(shù)據(jù)圖，將前兩個(gè)節(jié)點(diǎn)融合成一個(gè)節(jié)點(diǎn)FWD-BWD Super Node，將后兩個(gè)節(jié)點(diǎn)融合成一個(gè)節(jié)點(diǎn)Update Super Node。如下圖右邊所示，沿著gradient 16和parameter 16兩條邊切分。
現(xiàn)在的計(jì)算流程是，在GPU上面進(jìn)行前向和后向計(jì)算，將梯度傳給CPU，進(jìn)行參數(shù)更新，再將更新后的參數(shù)傳給GPU。為了提高效率，可以將計(jì)算和通信并行起來(lái)，GPU在反向傳播階段，可以待梯度值填滿bucket后，一邊計(jì)算新的梯度一邊將bucket傳輸給CPU，當(dāng)反向傳播結(jié)束，CPU基本上已經(jīng)有最新的梯度值了，同樣的，CPU在參數(shù)更新時(shí)也同步將已經(jīng)計(jì)算好的參數(shù)傳給GPU，如下圖所示。
到目前為止還都是單卡的場(chǎng)景，在多卡場(chǎng)景中，ZeRO-Offload可以利用ZeRO-2，將優(yōu)化器狀態(tài)和梯度進(jìn)行切分，每張卡只保留1N\frac{1}{N}N1?，結(jié)合上ZeRO-Offload同樣是將這1N\frac{1}{N}N1?的優(yōu)化器狀態(tài)和梯度卸載到內(nèi)存，在CPU上進(jìn)行參數(shù)更新。在多卡場(chǎng)景，利用CPU多核并行計(jì)算，每張卡至少對(duì)應(yīng)一個(gè)CPU進(jìn)程，由這個(gè)進(jìn)程負(fù)責(zé)進(jìn)行局部參數(shù)更新。
并且CPU和GPU的通信量和 N 無(wú)關(guān)，因?yàn)閭鬏數(shù)氖莊p16 gradient和fp16 parameter，總的傳輸量是固定的，由于利用多核并行計(jì)算，每個(gè)CPU進(jìn)程只負(fù)責(zé) 1N 的計(jì)算，反而隨著卡數(shù)增加節(jié)省了CPU計(jì)算時(shí)間。
ZeRO-Infinity
從GPT-1到GPT-3，兩年時(shí)間內(nèi)模型參數(shù)0.1B增加到175B，而同期，NVIDIA的顯卡是從V100的32GB顯存增加A100的80GB，顯然，顯寸的提升速度遠(yuǎn)遠(yuǎn)趕不上模型模型增長(zhǎng)的速度，這就是內(nèi)存墻問(wèn)題。
所謂內(nèi)存墻，其實(shí)就是硬件之間的通信會(huì)遇到瓶頸，不管是GPU之間還是GPU和CPU之間，通信的效率遠(yuǎn)低于計(jì)算的效率。在過(guò)去20年中，硬件的峰值計(jì)算能力增加了90,000倍，但是內(nèi)存/硬件互連帶寬卻只是提高了30倍。
ZeRO-Infinity這篇論文就是用來(lái)解決內(nèi)存墻的問(wèn)題，但是它并不是打通了內(nèi)存墻，而是通過(guò)NVMe。NVMe是一種接口規(guī)范，最開(kāi)始是用在SSD固態(tài)硬盤上的，之前的SSD用的都是SATA接口規(guī)范，數(shù)據(jù)傳輸速率最大6Gbps，而NVMe就是一種新的技術(shù)，NVMe固態(tài)硬盤擁有高達(dá)20Gbps的理論傳輸速度。
ZeRO-Infinity就是依靠 CPU 甚至是 NVMe 磁盤來(lái)訓(xùn)練大型模型。主要的想法是，在不使用張量時(shí)，將其卸載回 CPU 內(nèi)存或 NVMe 磁盤。通過(guò)使用異構(gòu)系統(tǒng)架構(gòu)，有可能在一臺(tái)機(jī)器上容納一個(gè)巨大的模型。
實(shí)踐
目標(biāo)：微調(diào) Wenzhong2.0-GPT2-3.5B-chinese
實(shí)驗(yàn)
配置文件
#!/bin/bashset -x -eecho "START TIME: $(date)" MICRO_BATCH_SIZE=1 ROOT_DIR=$(pwd)ZERO_STAGE=3config_json="$ROOT_DIR/training_config.json"cat <<EOT >$config_json {"train_micro_batch_size_per_gpu": $MICRO_BATCH_SIZE,"steps_per_print": 1000,"gradient_clipping": 1,"zero_optimization": {"stage": ${ZERO_STAGE},"allgather_partitions": false,"allgather_bucket_size": 2e8,"overlap_comm": true,"reduce_scatter": true,"reduce_bucket_size": 2e8,"contiguous_gradients": true,"offload_optimizer": {"device": "cpu","pin_memory": true},"offload_param": {"device": "cpu","pin_memory": true},"stage3_max_live_parameters" : 2e8,"stage3_max_reuse_distance" : 2e8,"stage3_prefetch_bucket_size": 2e8,"stage3_param_persistence_threshold": 2e8,"sub_group_size" : 2e8,"round_robin_gradients": true},"bf16": {"enabled": true},"optimizer": {"type": "Adam","params": {"lr": 1e-5,"betas": [0.9,0.95],"eps": 1e-8,"weight_decay": 1e-2}},"scheduler": {"type": "WarmupLR","params":{"warmup_min_lr": 5e-6,"warmup_max_lr": 1e-5}} } EOTexport PL_DEEPSPEED_CONFIG_PATH=$config_json TRAINER_ARGS="--max_epochs 1 \--num_nodes 2 \--gpus 8 \--strategy deepspeed_stage_${ZERO_STAGE}_offload \--default_root_dir $ROOT_DIR \--dirpath $ROOT_DIR/ckpt \--save_top_k 3 \--monitor train_loss \--mode min \--save_last \ "DATA_DIR=/home/liuzhaofeng/nlg_pipeline/gpt2/dialog/datasets DATA_ARGS="--data_dir $DATA_DIR \--max_seq_length 64 \--train_batchsize $MICRO_BATCH_SIZE \--valid_batchsize $MICRO_BATCH_SIZE \--train_data test_train.txt \--valid_data test.txt \--test_data test.txt "PRETRAINED_MODEL_PATH="IDEA-CCNL/Wenzhong2.0-GPT2-3.5B-chinese" MODEL_ARGS="--pretrained_model_path ${PRETRAINED_MODEL_PATH} \--output_save_path $ROOT_DIR/predict.json \--learning_rate 1e-4 \--weight_decay 0.1 \--warmup 0.01 \ "SCRIPTS_PATH=${ROOT_DIR}/finetune_gpt2.pyexport CMD=" \$SCRIPTS_PATH \$TRAINER_ARGS \$MODEL_ARGS \$DATA_ARGS \ "export NCCL_IB_DISABLE=1 export NCCL_DEBUG_SUBSYS=INIT,P2P,SHM,NET,GRAPH,ENV export NCCL_SOCKET_IFNAME=enp129s0f1torchrun ${CMD}壓力測(cè)試
| 1 | 64 | 153G | 5669M |
| 128 | 153G | 6041M | |
| 256 | 154G | 7139M | |
| 512 | 154G | 9383M | |
| 2 | 64 | 153G | 6367M |
| 128 | 153G | 7439M | |
| 256 | 153G | 9651M | |
| 512 | OOM | ||
| 3 | 64 | 167G | 6431M |
| 128 | 166G | 8273M | |
| 256 | 163G | 11103M |
總結(jié)
以上是生活随笔為你收集整理的震惊！我竟然在1080Ti上加载了一个35亿参数的模型（ZeRO, Zero Redundancy Optimizer）的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。
- 上一篇： 【Zotero文献管理】使用zotero
- 下一篇： OTG--miniUSB的工作原理