1. 項(xiàng)目概述大模型時(shí)代的效率革命如果你最近在搞大模型相關(guān)的項(xiàng)目無論是微調(diào)、推理還是應(yīng)用開發(fā)大概率會(huì)碰到一個(gè)讓人頭疼的問題顯存。動(dòng)輒幾十億、上百億參數(shù)的模型哪怕只是加載到GPU里看一眼都可能直接把你的顯存撐爆。更別提同時(shí)訓(xùn)練多個(gè)模型或者進(jìn)行復(fù)雜的多任務(wù)學(xué)習(xí)了。傳統(tǒng)的分布式訓(xùn)練方法比如數(shù)據(jù)并行Data Parallelism雖然能把數(shù)據(jù)分到多張卡上但每張卡上依然要保存一份完整的模型副本。模型一大這招就不好使了。這就是“ZERO”系列技術(shù)誕生的背景。它不是什么具體的算法模型而是一套由微軟DeepSpeed團(tuán)隊(duì)提出的、用于極致優(yōu)化大規(guī)模模型訓(xùn)練內(nèi)存和效率的分布式訓(xùn)練策略。你可以把它理解為一套“組合拳”專門對(duì)付大模型訓(xùn)練中的顯存“怪獸”。掌握ZERO意味著你能用有限的硬件資源比如幾塊消費(fèi)級(jí)顯卡去挑戰(zhàn)以前需要昂貴計(jì)算集群才能完成的任務(wù)。這不僅僅是技術(shù)上的優(yōu)化更是成本控制和研發(fā)效率的革命。無論是學(xué)生、研究者還是工程師只要你的工作與大模型沾邊ZERO就是你繞不開的必修課。2. ZERO核心思想與三級(jí)策略深度解析ZERO的核心思想非常直觀既然完整的模型參數(shù)、梯度和優(yōu)化器狀態(tài)太占地方那我們就把它們“拆開”分散到不同的GPU上去。通過精密的通信和協(xié)調(diào)在需要的時(shí)候再把它們組合起來。這種“分而治之”的思路直接擊中了分布式訓(xùn)練的痛點(diǎn)。ZERO主要分為三個(gè)等級(jí)ZERO-1, ZERO-2, ZERO-3它們像俄羅斯套娃一樣層層遞進(jìn)優(yōu)化得越來越徹底當(dāng)然實(shí)現(xiàn)的復(fù)雜度和通信開銷也會(huì)相應(yīng)增加。2.1 ZERO-1優(yōu)化器狀態(tài)分區(qū)這是入門級(jí)也是性價(jià)比最高的一級(jí)。它的目標(biāo)很明確干掉最占地方的“元兇”——優(yōu)化器狀態(tài)。為什么是優(yōu)化器狀態(tài)以最常用的Adam優(yōu)化器為例。對(duì)于每一個(gè)模型參數(shù)Adam需要維護(hù)兩個(gè)狀態(tài)一階動(dòng)量m和二階動(dòng)量v。假設(shè)模型有Ψ個(gè)參數(shù)使用FP16混合精度訓(xùn)練那么參數(shù)本身占用2Ψ字節(jié)FP16。梯度占用2Ψ字節(jié)FP16。Adam狀態(tài)包括FP32格式的參數(shù)副本、一階動(dòng)量、二階動(dòng)量共3 * 4Ψ 12Ψ字節(jié)。計(jì)算一下比例12Ψ / (2Ψ 2Ψ 12Ψ) 12Ψ / 16Ψ 75%?？吹搅藛醿?yōu)化器狀態(tài)吃掉了高達(dá)75%的顯存ZERO-1就專門對(duì)付它。具體如何操作假設(shè)我們有Nd張GPU數(shù)據(jù)并行維度。在傳統(tǒng)數(shù)據(jù)并行中每張卡都有一份完整的優(yōu)化器狀態(tài)。ZERO-1的做法是分區(qū)將完整的優(yōu)化器狀態(tài)均勻地分割成Nd份。分發(fā)每個(gè)GPU只保存其中一份。例如GPU0保存第1到第Ψ/Nd個(gè)參數(shù)的優(yōu)化器狀態(tài)GPU1保存下一份以此類推。通信與更新前向和反向傳播時(shí)每張卡都有完整的參數(shù)和梯度與傳統(tǒng)DP一樣。在優(yōu)化器更新參數(shù)時(shí)情況變了。每個(gè)GPU只負(fù)責(zé)更新自己“管轄”的那部分參數(shù)。因?yàn)樗粨碛心遣糠謪?shù)的優(yōu)化器狀態(tài)。更新完成后它需要通過“全體收集”All-Gather操作將自己更新好的那部分參數(shù)廣播給所有其他GPU確保所有GPU上的參數(shù)保持一致。效果與權(quán)衡顯存節(jié)省優(yōu)化器狀態(tài)的顯存占用直接降為原來的1/Nd。這是巨大的提升。通信開銷引入了額外的All-Gather通信來同步參數(shù)。但通常優(yōu)化器步驟的計(jì)算量不大通信開銷相對(duì)可控性價(jià)比極高。實(shí)操心得ZERO-1幾乎是所有大模型訓(xùn)練的起點(diǎn)。在DeepSpeed中你只需要在配置文件中將stage設(shè)置為1并指定optimizer和scheduler就能輕松啟用。它帶來的顯存收益是立竿見影的而增加的通信成本在大多數(shù)網(wǎng)絡(luò)環(huán)境下都是可以接受的。對(duì)于很多場(chǎng)景僅用ZERO-1就能讓你把模型規(guī)模擴(kuò)大近一倍。2.2 ZERO-2梯度分區(qū)解決了優(yōu)化器狀態(tài)下一個(gè)目標(biāo)就是梯度。在反向傳播結(jié)束后每張卡上都會(huì)計(jì)算出完整的梯度張量。ZERO-2在ZERO-1的基礎(chǔ)上進(jìn)一步對(duì)梯度進(jìn)行分區(qū)。工作原理梯度計(jì)算與分區(qū)反向傳播過程中每張卡計(jì)算出完整的梯度后立即對(duì)其進(jìn)行分區(qū)Nd份然后通過“規(guī)約散射”Reduce-Scatter操作。Reduce-Scatter操作這個(gè)操作可以理解為兩步的合并。首先所有GPU將各自梯度對(duì)應(yīng)的分區(qū)進(jìn)行“規(guī)約”Reduce通常是求和得到全局梯度的分區(qū)。然后這個(gè)結(jié)果被“散射”Scatter到對(duì)應(yīng)的GPU上。最終每個(gè)GPU只持有全局梯度的一部分。優(yōu)化器更新由于ZERO-1已經(jīng)將優(yōu)化器狀態(tài)分區(qū)每個(gè)GPU正好用自己持有的那一份梯度去更新自己持有的那一份優(yōu)化器狀態(tài)和參數(shù)。參數(shù)同步更新后同樣通過All-Gather操作同步所有參數(shù)。效果與權(quán)衡顯存節(jié)省梯度顯存占用也降為原來的1/Nd。結(jié)合ZERO-1現(xiàn)在參數(shù)和梯度是完整的各占2Ψ字節(jié)但優(yōu)化器狀態(tài)和梯度都只有1/Nd。顯存節(jié)省進(jìn)一步擴(kuò)大。通信開銷引入了Reduce-Scatter操作。與ZERO-1的All-Gather相比Reduce-Scatter的通信量是相同的都是傳輸Ψ個(gè)元素但算法不同。通?，F(xiàn)代GPU集群使用NCCL后端對(duì)這兩種集體操作都有高度優(yōu)化實(shí)際開銷增加并不顯著。注意事項(xiàng)ZERO-2的通信發(fā)生在反向傳播結(jié)束后這可能會(huì)略微增加反向傳播階段的時(shí)間。但在整體訓(xùn)練流程中由于顯存大幅降低你可能可以使用更大的批量大小batch size來彌補(bǔ)甚至獲得更高的吞吐量。需要根據(jù)實(shí)際硬件和模型進(jìn)行 profiling 來權(quán)衡。2.3 ZERO-3參數(shù)分區(qū)這是最激進(jìn)、也是最徹底的優(yōu)化級(jí)別。ZERO-3的思想是既然優(yōu)化器狀態(tài)和梯度都分區(qū)了為什么不把模型參數(shù)本身也分區(qū)呢讓每張GPU在平時(shí)只保存整個(gè)模型的1/Nd。工作原理前向傳播當(dāng)某層需要計(jì)算時(shí)通過All-Gather操作從所有GPU上收集該層所需的完整參數(shù)。計(jì)算完成后立即釋放掉從其他GPU收集來的參數(shù)只保留自己負(fù)責(zé)的那一部分。這被稱為“參數(shù)卸載”。反向傳播與正向類似需要時(shí)收集參數(shù)計(jì)算梯度。梯度計(jì)算完成后同樣通過Reduce-Scatter操作將梯度分區(qū)并分發(fā)。優(yōu)化器更新每個(gè)GPU用自己分區(qū)的梯度更新自己分區(qū)的優(yōu)化器狀態(tài)和參數(shù)。效果與權(quán)衡顯存節(jié)省達(dá)到極致。每張GPU上持久化保存的只有1/Nd的參數(shù) 1/Nd的梯度 1/Nd的優(yōu)化器狀態(tài)。模型顯存占用幾乎與GPU數(shù)量成線性反比。通信開銷急劇增加。因?yàn)樵谇跋蚝头聪虻拿恳粚佣伎赡苄枰M(jìn)行All-Gather和Reduce-Scatter操作。通信頻率和總量遠(yuǎn)高于前兩個(gè)階段。計(jì)算粒度ZERO-3通常與模型并行Model Parallelism或流水線并行Pipeline Parallelism結(jié)合使用以更細(xì)的粒度如層內(nèi)進(jìn)行參數(shù)分區(qū)從而減少每次通信的數(shù)據(jù)量這也就是所謂的“ZERO-Infinity”或“ZERO”所做的進(jìn)一步優(yōu)化。踩坑實(shí)錄ZERO-3雖然省顯存但絕不是“無腦開”。通信開銷可能成為嚴(yán)重的性能瓶頸尤其是在節(jié)點(diǎn)間網(wǎng)絡(luò)帶寬不足的情況下。開啟ZERO-3后訓(xùn)練速度可能會(huì)顯著下降。我們的經(jīng)驗(yàn)是只有當(dāng)模型大到連ZERO-2都無法加載時(shí)才考慮ZERO-3并且一定要搭配高速互聯(lián)如NVLink, InfiniBand和細(xì)致的性能剖析。3. 實(shí)戰(zhàn)使用DeepSpeed配置與調(diào)優(yōu)ZERO理論懂了關(guān)鍵還得落地。微軟的DeepSpeed庫是目前實(shí)現(xiàn)和使用ZERO最方便、最強(qiáng)大的工具。下面我們以一個(gè)具體的例子看看如何配置和調(diào)優(yōu)。3.1 基礎(chǔ)配置與啟動(dòng)假設(shè)我們有一個(gè)基于Hugging Face Transformers的模型訓(xùn)練腳本train.py。使用DeepSpeed的第一步是創(chuàng)建一個(gè)配置文件ds_config.json。一個(gè)啟用ZERO-1的基礎(chǔ)配置{ train_batch_size: 32, gradient_accumulation_steps: 1, fp16: { enabled: true, loss_scale: 0, loss_scale_window: 1000, initial_scale_power: 16 }, zero_optimization: { stage: 1, // 啟用ZERO-1 allgather_partitions: true, allgather_bucket_size: 5e8, overlap_comm: true, reduce_scatter: true, reduce_bucket_size: 5e8 }, optimizer: { type: AdamW, params: { lr: 5e-5, betas: [0.9, 0.999], eps: 1e-8, weight_decay: 0.01 } }, scheduler: { type: WarmupLR, params: { warmup_min_lr: 0, warmup_max_lr: 5e-5, warmup_num_steps: 1000 } } }關(guān)鍵字段解析stage: ZERO的階段12或3。allgather_bucket_size和reduce_bucket_size: 通信桶大小。將大量小張量的通信聚合成少量大張量的通信能極大提升效率。一般設(shè)置為5e8500MB左右是個(gè)不錯(cuò)的起點(diǎn)。overlap_comm: 是否重疊通信和計(jì)算。開啟后在通信進(jìn)行的同時(shí)GPU可以繼續(xù)做其他計(jì)算能有效隱藏通信延遲。強(qiáng)烈建議開啟。啟動(dòng)訓(xùn)練的命令也很簡(jiǎn)單deepspeed --num_gpus4 train.py --deepspeed ds_config.json3.2 進(jìn)階調(diào)優(yōu)與參數(shù)解讀當(dāng)你需要啟用ZERO-2或ZERO-3時(shí)配置需要更精細(xì)的調(diào)整。ZERO-2配置示例zero_optimization: { stage: 2, contiguous_gradients: true, overlap_comm: true, reduce_scatter: true, reduce_bucket_size: 5e8, allgather_bucket_size: 5e8, cpu_offload: false // 謹(jǐn)慎開啟見下文 }contiguous_gradients: 在反向傳播前將梯度緩沖區(qū)置為連續(xù)內(nèi)存。這能提升Reduce-Scatter操作的效率建議開啟。ZERO-3配置示例zero_optimization: { stage: 3, contiguous_gradients: true, stage3_max_live_parameters: 1e9, stage3_max_reuse_distance: 1e9, stage3_prefetch_bucket_size: 5e8, stage3_param_persistence_threshold: 1e6, overlap_comm: true, reduce_bucket_size: 5e8, allgather_bucket_size: 2e8, // ZERO-3下可以調(diào)小一些 cpu_offload: false }ZERO-3的配置參數(shù)更為復(fù)雜stage3_max_live_parameters: 控制同一時(shí)刻駐留在GPU上的參數(shù)數(shù)量上限。調(diào)小可以省顯存但可能增加通信次數(shù)。stage3_prefetch_bucket_size: 參數(shù)預(yù)取桶大小。DeepSpeed會(huì)嘗試在需要參數(shù)之前提前發(fā)起All-Gather通信與計(jì)算重疊。這個(gè)參數(shù)控制預(yù)取的量。stage3_param_persistence_threshold: 參數(shù)持久化閾值。小于此大小的參數(shù)張量不會(huì)被分區(qū)/卸載而是常駐在所有GPU上。這對(duì)于偏置bias、層歸一化LayerNorm參數(shù)等小張量很有效能避免為它們付出昂貴的通信代價(jià)。關(guān)于CPU Offload配置中有一個(gè)cpu_offload選項(xiàng)。當(dāng)設(shè)置為true時(shí)可以將優(yōu)化器狀態(tài)、梯度甚至參數(shù)卸載到CPU內(nèi)存。這能進(jìn)一步釋放GPU顯存讓你跑起更大的模型。血淚教訓(xùn)CPU Offload是一把雙刃劍。雖然顯存省了但數(shù)據(jù)在CPU和GPU之間的傳輸PCIe帶寬會(huì)成為巨大瓶頸訓(xùn)練速度可能下降一個(gè)數(shù)量級(jí)。除非你的模型真的巨大無比且訓(xùn)練時(shí)間不是首要考慮因素例如只為獲取一個(gè)預(yù)訓(xùn)練權(quán)重否則不要輕易開啟。我們的原則是優(yōu)先用盡GPU顯存和通信優(yōu)化最后才考慮CPU Offload。3.3 性能監(jiān)控與瓶頸分析開啟DeepSpeed后如何知道性能瓶頸在哪DeepSpeed提供了豐富的日志和性能分析工具。查看日志在訓(xùn)練命令中加入--deepspeed_config ds_config.json --deepspeed 21 | tee train.log日志中會(huì)包含各個(gè)階段的耗時(shí)。使用Timeline在配置文件中啟用wall_clock_breakdown: true和flops_profiler: {enabled: true}可以生成更詳細(xì)的時(shí)間線分析看清是計(jì)算耗時(shí)多還是通信耗時(shí)多。調(diào)整桶大小reduce_bucket_size和allgather_bucket_size是最關(guān)鍵的調(diào)優(yōu)參數(shù)。如果通信耗時(shí)占比高可以嘗試增大它們但不要超過單個(gè)張量的最大限制。如果GPU顯存利用率低可以嘗試減小它們。這是一個(gè)需要反復(fù)試驗(yàn)的過程。結(jié)合NVProf/Nsight Systems對(duì)于更深度的性能分析可以結(jié)合NVIDIA的性能分析工具查看CUDA Kernel執(zhí)行和通信操作的具體耗時(shí)。4. 常見問題排查與實(shí)戰(zhàn)技巧在實(shí)際部署中你會(huì)遇到各種各樣的問題。這里記錄了一些典型場(chǎng)景和解決方法。4.1 內(nèi)存溢出OOM問題即使開了ZERO也可能OOM?，F(xiàn)象訓(xùn)練剛開始或中途報(bào)CUDA out of memory。排查思路檢查配置階段確認(rèn)stage設(shè)置正確。ZERO-3比ZERO-2省顯存。檢查激活值內(nèi)存ZERO主要優(yōu)化參數(shù)、梯度、優(yōu)化器狀態(tài)。但前向傳播中產(chǎn)生的激活值A(chǔ)ctivations也可能占用大量顯存?？梢詥⒂眉せ钪禉z查點(diǎn)Activation Checkpointing或梯度檢查點(diǎn)用計(jì)算換內(nèi)存。在DeepSpeed配置中可以通過activation_checkpointing部分配置。檢查批量大小train_batch_size是全局批量大小。DeepSpeed會(huì)自動(dòng)根據(jù)GPU數(shù)量計(jì)算每張卡的本地批量大小。如果你手動(dòng)設(shè)置了gradient_accumulation_steps確保train_batch_size per_gpu_batch_size * num_gpus * gradient_accumulation_steps。檢查模型大小用torch.cuda.max_memory_allocated()在關(guān)鍵位置打印顯存使用定位內(nèi)存峰值。4.2 訓(xùn)練速度慢或不穩(wěn)定現(xiàn)象開啟ZERO后迭代時(shí)間變長(zhǎng)或Loss曲線震蕩劇烈。排查思路通信瓶頸這是ZERO-2/3最常見的問題。使用wall_clock_breakdown分析時(shí)間。如果通信占比過高如30%嘗試增大reduce_bucket_size和allgather_bucket_size。確保使用了overlap_comm: true。檢查硬件單機(jī)多卡確保使用NVLink多機(jī)確保使用InfiniBand等高速網(wǎng)絡(luò)。精度問題FP16混合精度訓(xùn)練可能不穩(wěn)定特別是當(dāng)模型中有非常小或非常大的梯度值時(shí)?？梢試L試使用DeepSpeed的fp16: {loss_scale: 0}動(dòng)態(tài)損失縮放?；蛘呖紤]使用BFloat16如果硬件支持其數(shù)值范圍比FP16更穩(wěn)健。優(yōu)化器狀態(tài)同步在ZERO-1/2下確保優(yōu)化器步驟后參數(shù)同步正確?？梢远ㄆ跈z查不同GPU上同一參數(shù)的數(shù)值是否一致。4.3 保存與加載檢查點(diǎn)這是一個(gè)容易踩坑的地方。在ZERO-3下模型參數(shù)是分區(qū)的不能直接用torch.save(model.state_dict(), ...)。正確做法使用DeepSpeed提供的engine.save_checkpoint()和engine.load_checkpoint()API。保存engine.save_checkpoint(save_dir, tag, client_state{...})會(huì)以分布式的方式保存所有分區(qū)。加載load_checkpoint(save_dir, tag, load_optimizer_statesTrue, load_lr_scheduler_statesTrue)。關(guān)鍵點(diǎn)保存的檢查點(diǎn)目錄結(jié)構(gòu)是特定的不要手動(dòng)修改。加載時(shí)需要先初始化DeepSpeed引擎再用該引擎加載。4.4 與其它并行策略的混合使用在實(shí)際的超大模型訓(xùn)練中ZERO常與模型并行MP、流水線并行PP結(jié)合。與模型并行結(jié)合DeepSpeed支持自動(dòng)化的張量并行Tensor Parallelism在配置中通過tensor_parallel: {tp_size: 2}來設(shè)置。ZERO負(fù)責(zé)數(shù)據(jù)并行維度的內(nèi)存優(yōu)化MP負(fù)責(zé)模型層內(nèi)的切分。此時(shí)ZERO的Nd指的是數(shù)據(jù)并行組的規(guī)模它會(huì)自動(dòng)調(diào)整。與流水線并行結(jié)合通過pipeline_parallel: {pp_size: 2}啟用。流水線并行將模型按層切分到不同GPUZERO則在每個(gè)流水線階段內(nèi)進(jìn)行數(shù)據(jù)并行優(yōu)化。配置相對(duì)復(fù)雜需要仔細(xì)規(guī)劃micro-batch和梯度累積步驟。我個(gè)人在多次項(xiàng)目中的體會(huì)是ZERO不是一個(gè)“設(shè)置完就忘”的黑盒。它是一套需要你根據(jù)具體模型、硬件和數(shù)據(jù)流進(jìn)行精細(xì)調(diào)優(yōu)的工具集。從ZERO-1開始逐步推進(jìn)密切監(jiān)控日志和性能指標(biāo)小步快跑地調(diào)整參數(shù)是掌握它的不二法門。當(dāng)你看到原本需要8張A100才能訓(xùn)練的模型在4張3090上穩(wěn)定跑起來時(shí)那種感覺就是工程師的快樂。