目錄FSDP 的設(shè)計(jì)動機(jī)ZeRO-3 分片原理FSDP 的通信模式FSDP 的工程實(shí)現(xiàn)FSDP 與 DDP 的對比FSDP 的邊界與失效模式摘要FSDPFully Sharded Data Parallel將模型參數(shù)、梯度和優(yōu)化器狀態(tài)分片到多個 GPU 上每個 GPU 只存儲部分參數(shù)通過通信獲取完整參數(shù)進(jìn)行前向和反向傳播。本文從 FSDP 的設(shè)計(jì)動機(jī)出發(fā)分析 ZeRO-3 的分片原理、通信模式以及與 DDP 的對比。1. FSDP 的設(shè)計(jì)動機(jī)DDP 在訓(xùn)練時每個 GPU 持有完整的模型副本。當(dāng)模型參數(shù)量超過單 GPU 顯存時DDP 無法使用。FSDP 通過將模型參數(shù)分片到多個 GPU 上使超大模型可以在有限的顯存上訓(xùn)練。1.1 為什么需要 FSDP模型規(guī)模參數(shù)量顯存需求FP16DDP 可行性FSDP 可行性7B7B14GB可80GB 卡可13B13B26GB可80GB 卡可30B30B60GB可80GB 卡可70B70B140GB不可需 2 卡可8 卡175B175B350GB不可可32 卡1.2 FSDP 的核心思想FSDP 的核心思想是將模型參數(shù)分片到多個 GPU 上每個 GPU 只存儲部分參數(shù)需要時通過通信獲取完整參數(shù)。模型參數(shù)GPU 0: 參數(shù)分片 0GPU 1: 參數(shù)分片 1GPU 2: 參數(shù)分片 2GPU 3: 參數(shù)分片 3前向: All-Gather 獲取完整參數(shù)計(jì)算后丟棄非本分片參數(shù)反向: 再次 All-Gather計(jì)算梯度后 Reduce-Scatter更新本分片參數(shù)1.3 FSDP 的歷史演進(jìn)數(shù)據(jù)并行DDP→ ZeRO-1優(yōu)化器分片, 2019→ ZeRO-2梯度分片, 2020→ ZeRO-3全參數(shù)分片, 2020→ FSDPPyTorch 實(shí)現(xiàn), 2022。1.4 FSDP 的產(chǎn)業(yè)應(yīng)用模型參數(shù)規(guī)模GPU 配置策略LLaMA 7B7B8 A100FSDPLLaMA 13B13B8 A100FSDPLLaMA 70B70B64 A100FSDP 張量并行GPT-3 175B175B256 A100FSDP 3D 并行1.5 FSDP 的局限性FSDP 的局限性包括通信量增加前向和反向都需要 All-Gather、實(shí)現(xiàn)復(fù)雜度高需要手動管理分片以及小模型效率低小模型下 FSDP 不如 DDP 效率高。2. ZeRO-3 分片原理2.1 ZeRO-3 的分片策略ZeRO-3 將模型參數(shù)、梯度和優(yōu)化器狀態(tài)全部在 GPU 間分片分片內(nèi)容分片前分片后顯存節(jié)省模型參數(shù)14GB7B14GB/N1/N梯度14GB7B14GB/N1/N優(yōu)化器狀態(tài)28GBAdam28GB/N1/N總計(jì)56GB56GB/N1/N2.2 分片比例N GPU 數(shù)量 N \text{GPU 數(shù)量}NGPU數(shù)量GPU 數(shù)量7B 模型顯存70B 模型顯存156GB560GB87GB70GB321.75GB17.5GB640.875GB8.75GB2.3 分片與恢復(fù)FSDP 在計(jì)算時需要將分片的參數(shù)恢復(fù)為完整參數(shù)classFSDPUnit:FSDP 分片單元def__init__(self,param,world_size,rank):self.paramparam self.world_sizeworld_size self.rankrank self.sharded_paramself._shard(param)def_shard(self,param):分片參數(shù)chunk_sizeparam.numel()//self.world_size startself.rank*chunk_size end(self.rank1)*chunk_sizereturnparam.detach().flatten()[start:end].clone()defgather(self):收集完整參數(shù)# All-Gather 收集所有分片full_params[torch.zeros_like(self.sharded_param)for_inrange(self.world_size)]dist.all_gather(full_params,self.sharded_param)returntorch.cat(full_params).reshape(self.param.shape)defscatter(self,grad):分發(fā)梯度# Reduce-Scatter 分發(fā)梯度chunk_sizegrad.numel()//self.world_size sharded_gradtorch.zeros(chunk_size,devicegrad.device)dist.reduce_scatter(sharded_grad,grad.view(self.world_size,chunk_size))returnsharded_grad3. FSDP 的通信模式3.1 前向通信前向傳播時FSDP 需要從其他 GPU 收集完整參數(shù)All-Gather → Forward → Discard Non-local Params \text{All-Gather} \rightarrow \text{Forward} \rightarrow \text{Discard Non-local Params}All-Gather→Forward→Discard Non-local Params3.2 反向通信反向傳播時FSDP 再次收集完整參數(shù)計(jì)算梯度然后分發(fā)梯度All-Gather → Backward → Reduce-Scatter \text{All-Gather} \rightarrow \text{Backward} \rightarrow \text{Reduce-Scatter}All-Gather→Backward→Reduce-Scatter3.3 通信量對比策略前向通信反向通信總通信量DDP02 × Model2 × ModelFSDP1 × Model2 × Model3 × Model3.4 通信與計(jì)算重疊FSDP 通過預(yù)取Prefetching實(shí)現(xiàn)通信與計(jì)算重疊classFSDPWithPrefetch:帶預(yù)取的 FSDPdef__init__(self,modules,world_size,rank):self.modulesmodules self.world_sizeworld_size self.rankrankdefforward(self,x):fori,moduleinenumerate(self.modules):# 預(yù)取下一個模塊的參數(shù)ifi1len(self.modules):self.modules[i1].prefetch_params()# 收集當(dāng)前模塊參數(shù)module.gather_params()# 前向傳播xmodule(x)# 丟棄非本分片參數(shù)module.discard_non_local()returnx4. FSDP 的工程實(shí)現(xiàn)4.1 FSDP 的基本使用fromtorch.distributed.fsdpimportFullyShardedDataParallelasFSDPfromtorch.distributed.fsdp.wrapimporttransformer_auto_wrap_policy# 創(chuàng)建 FSDP 模型modelFSDP(model,sharding_strategyShardingStrategy.FULL_SHARD,# ZeRO-3auto_wrap_policytransformer_auto_wrap_policy,# 自動包裝device_idrank,mixed_precisionTrue# 混合精度)# 訓(xùn)練循環(huán)與 DDP 類似forbatchindataloader:lossmodel(batch)loss.backward()optimizer.step()optimizer.zero_grad()4.2 FSDP 的訓(xùn)練配置參數(shù)值說明sharding_strategyFULL_SHARDZeRO-3 全分片auto_wrap_policytransformer_auto_wrap_policy自動包裝策略mixed_precisionTrue混合精度訓(xùn)練device_idrank當(dāng)前 GPU 設(shè)備 ID4.3 FSDP 的包裝策略FSDP 提供多種包裝策略策略描述適用場景默認(rèn)包裝整個模型包裝為一個 FSDP 單元小模型按層包裝每層一個 FSDP 單元通用自定義包裝指定哪些模塊包裝為 FSDP 單元大模型4.4 FSDP 與混合精度訓(xùn)練fromtorch.distributed.fsdpimportMixedPrecision# 混合精度配置mp_configMixedPrecision(param_dtypetorch.bfloat16,# 參數(shù)精度reduce_dtypetorch.bfloat16,# 通信精度buffer_dtypetorch.bfloat16# 緩沖精度)modelFSDP(model,mixed_precisionmp_config,sharding_strategyShardingStrategy.FULL_SHARD)5. FSDP 與 DDP 的對比5.1 性能對比對比維度DDPFSDP顯存占用高完整模型低分片模型通信量2 × Model3 × Model訓(xùn)練速度快慢通信多適用模型 單 GPU 顯存 單 GPU 顯存5.2 擴(kuò)展性對比GPU 數(shù)量DDP 擴(kuò)展比FSDP 擴(kuò)展比推薦87.6x7.0xDDP3228.0x25.0xDDP6450.0x45.0x根據(jù)模型大小256150.0x140.0xFSDP5.3 選擇指南場景推薦策略原因模型 單 GPU 顯存DDP通信少速度快模型 單 GPU 顯存FSDP分片節(jié)省顯存模型 多個 GPU 顯存FSDP 張量并行分片 模型并行6. FSDP 的邊界與失效模式6.1 通信瓶頸問題表現(xiàn)解決方案通信量大訓(xùn)練速度慢增加 GPU 數(shù)量通信延遲高同步等待時間長使用更高速網(wǎng)絡(luò)通信不平衡某些 GPU 負(fù)載高優(yōu)化通信拓?fù)?.2 顯存管理問題表現(xiàn)解決方案顯存不足訓(xùn)練失敗減小模型或增加 GPU顯存碎片訓(xùn)練不穩(wěn)定使用內(nèi)存池峰值顯存前向時顯存峰值高啟用梯度檢查點(diǎn)6.3 FSDP 的優(yōu)缺點(diǎn)總結(jié)優(yōu)點(diǎn)缺點(diǎn)節(jié)省顯存通信量增加支持超大模型實(shí)現(xiàn)復(fù)雜度高與 DDP 兼容小模型效率低支持混合精度調(diào)試?yán)щy7. FSDP 的最佳實(shí)踐7.1 分片策略選擇策略顯存節(jié)省通信量適用場景NO_SHARD0%2 × Model等同于 DDPSHARD_GRAD50%2 × ModelZeRO-2FULL_SHARD66%3 × ModelZeRO-3推薦HYBRID_SHARD66%3 × Model跨節(jié)點(diǎn)分片7.2 性能優(yōu)化優(yōu)化策略描述效果預(yù)取預(yù)取下一個模塊的參數(shù)減少通信等待梯度檢查點(diǎn)減少前向激活顯存節(jié)省 30% 顯存混合精度BF16 訓(xùn)練減少 50% 顯存參數(shù)卸載卸載到 CPU 內(nèi)存節(jié)省更多顯存7.3 監(jiān)控與調(diào)試指標(biāo)描述告警閾值通信時間通信占總時間比例30%顯存使用各 GPU 顯存使用率90%吞吐量每秒處理的樣本數(shù)低于預(yù)期 50%8. FSDP 的高級配置8.1 分片策略選擇FSDP 支持多種分片策略適用于不同的訓(xùn)練場景策略參數(shù)分片梯度分片優(yōu)化器分片顯存節(jié)省通信量NO_SHARD否否否0%2 × ModelSHARD_GRAD否是是50%2 × ModelFULL_SHARD是是是66%3 × ModelHYBRID_SHARD節(jié)點(diǎn)內(nèi)分片節(jié)點(diǎn)內(nèi)分片節(jié)點(diǎn)內(nèi)分片66%3 × Model8.2 自動包裝策略fromtorch.distributed.fsdp.wrapimport(transformer_auto_wrap_policy,size_based_auto_wrap_policy,lambda_auto_wrap_policy)# 基于 Transformer 層自動包裝wrap_policypartial(transformer_auto_wrap_policy,transformer_layer_cls{TransformerBlock})# 基于參數(shù)大小自動包裝wrap_policypartial(size_based_auto_wrap_policy,min_num_params1e7# 10M 參數(shù)以上包裝)# 自定義包裝策略defcustom_wrap_policy(module,recurse,nonwrapped_numel):ifisinstance(module,(AttentionLayer,FFNLayer)):returnTruereturnFalsemodelFSDP(model,auto_wrap_policycustom_wrap_policy,sharding_strategyShardingStrategy.FULL_SHARD)8.3 參數(shù)卸載FSDP 支持將參數(shù)卸載到 CPU 內(nèi)存進(jìn)一步節(jié)省 GPU 顯存fromtorch.distributed.fsdpimportCPUOffload modelFSDP(model,cpu_offloadCPUOffload(offload_paramsTrue),# 參數(shù)卸載到 CPUsharding_strategyShardingStrategy.FULL_SHARD)9. FSDP 在訓(xùn)練中的性能分析9.1 通信時間分析模型規(guī)模GPU 數(shù)量通信時間占比計(jì)算時間占比吞吐量7B825%75%100%13B1630%70%85%70B6440%60%70%175B25650%50%50%9.2 顯存占用對比defmemory_comparison():FSDP vs DDP 顯存對比model_size7e9# 7B 參數(shù)fp16_bytes2# 每個參數(shù) 2 字節(jié)gpu_count8ddp_memory{parameters:model_size*fp16_bytes,gradients:model_size*fp16_bytes,optimizer_states:model_size*fp16_bytes*2,# Adam 動量 方差total:model_size*fp16_bytes*4}fsdp_memory{parameters:model_size*fp16_bytes/gpu_count,gradients:model_size*fp16_bytes/gpu_count,optimizer_states:model_size*fp16_bytes*2/gpu_count,total:model_size*fp16_bytes*4/gpu_count}returnddp_memory,fsdp_memory9.3 優(yōu)化建議問題優(yōu)化建議效果通信時間占比高增大 batch size減少通信次數(shù)顯存不足啟用梯度檢查點(diǎn)節(jié)省 30% 顯存訓(xùn)練速度慢使用混合精度加速 2x擴(kuò)展效率低優(yōu)化通信拓?fù)涮岣邤U(kuò)展比10. FSDP 的調(diào)試與故障排除10.1 常見問題問題表現(xiàn)解決方案顯存不足OOM 錯誤減小 batch size 或啟用參數(shù)卸載通信超時訓(xùn)練卡住NCCL_DEBUGINFO 查看通信狀態(tài)梯度爆炸loss 變成 NaN梯度裁剪或降低學(xué)習(xí)率數(shù)值不穩(wěn)定混合精度下 loss 不穩(wěn)定使用 BF16 替代 FP1610.2 性能調(diào)試defprofile_fsdp_training(model,dataloader,profiler):FSDP 訓(xùn)練性能分析forbatchindataloader:profiler.step()withprofiler.record_function(forward):lossmodel(batch)withprofiler.record_function(backward):loss.backward()withprofiler.record_function(optimizer):optimizer.step()optimizer.zero_grad()10.3 日志記錄defsetup_fsdp_logging(rank,log_dirlogs):設(shè)置 FSDP 日志loggerlogging.getLogger(ffsdp_rank_{rank})logger.setLevel(logging.INFO)# 文件日志fhlogging.FileHandler(f{log_dir}/fsdp_rank_{rank}.log)fh.setLevel(logging.INFO)# 控制臺日志chlogging.StreamHandler()ch.setLevel(logging.INFO)formatterlogging.Formatter(%(asctime)s - %(name)s - %(levelname)s - %(message)s)fh.setFormatter(formatter)ch.setFormatter(formatter)logger.addHandler(fh)logger.addHandler(ch)returnlogger11. FSDP 的未來方向11.1 異構(gòu) FSDPFSDP 擴(kuò)展到異構(gòu)設(shè)備不同型號 GPU、不同代 GPU異構(gòu)場景挑戰(zhàn)解決方案不同型號 GPU計(jì)算速度不同動態(tài)負(fù)載均衡不同顯存容量不同自適應(yīng)分片不同帶寬通信速度不同自適應(yīng)通信11.2 FSDP 張量并行FSDP 與張量并行結(jié)合支持更大模型的訓(xùn)練# FSDP 張量并行modelTensorParallel(model,tp_size2)modelFSDP(model,sharding_strategyShardingStrategy.FULL_SHARD)11.3 FSDP 的自動調(diào)優(yōu)FSDP 的自動調(diào)優(yōu)根據(jù)模型大小和硬件配置自動選擇最優(yōu)的分片策略、包裝策略和通信方案?？偨Y(jié)FSDP 是 PyTorch 實(shí)現(xiàn) ZeRO-3 全參數(shù)分片訓(xùn)練的核心庫。通過將模型參數(shù)、梯度和優(yōu)化器狀態(tài)分片到多個 GPU 上FSDP 使超大模型可以在有限的顯存上訓(xùn)練。FSDP 的通信量比 DDP 多但顯存節(jié)省顯著。FSDP 適用于模型超過單 GPU 顯存容量的場景在小模型場景下 DDP 更高效。外部引用ZeRO 原始論文https://arxiv.org/abs/1910.02054PyTorch FSDP 文檔https://pytorch.org/docs/stable/fsdp.htmlFSDP 分片策略https://pytorch.org/docs/stable/fsdp.htmlFSDP 混合精度https://pytorch.org/docs/stable/fsdp.htmlFSDP 與 DDP 對比https://pytorch.org/docs/stable/fsdp.htmlZeRO-3 全分片原理https://arxiv.org/abs/1910.02054FSDP 預(yù)取優(yōu)化https://pytorch.org/docs/stable/fsdp.htmlFSDP 梯度檢查點(diǎn)https://pytorch.org/docs/stable/checkpoint.html分布式訓(xùn)練擴(kuò)展性分析https://arxiv.org/abs/2303.04226FSDP 最佳實(shí)踐https://pytorch.org/docs/stable/fsdp.html