分片實(shí)戰(zhàn)：按需聚集機(jī)制拆解與 stage3 配置調(diào)優(yōu)指南)
DeepSpeed ZeRO-3 參數(shù)分片實(shí)戰(zhàn)按需聚集機(jī)制拆解與 stage3 配置調(diào)優(yōu)指南【免費(fèi)下載鏈接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.項(xiàng)目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed[rank 1] CUDA out of memory. Tried to allocate 112.00 GiB. Process 0 has 27.97 GiB in use. Of which 26.41 GiB is allocated by PyTorch.微調(diào) 70B 模型時(shí),普通數(shù)據(jù)并行要求每張 GPU 持有完整參數(shù)副本,單卡直接 OOM。DeepSpeed ZeRO-3 的參數(shù)分區(qū)(把參數(shù)切成 N 份,每張卡只留自己那 1/N)配合按需聚集,把單卡顯存占用從全量副本降到局部分片臨時(shí)聚集。本文從ZeroParamStatus狀態(tài)機(jī)講起,覆蓋預(yù)取、釋放全生命周期,以及 7 個(gè) stage3 關(guān)鍵參數(shù)的調(diào)優(yōu)邊界。訓(xùn)練 1T 參數(shù)模型時(shí),每張卡實(shí)際駐留的參數(shù)只有總量的 1/512,這是 ZeRO-3 與張量/流水線并行的核心差異。官方發(fā)布的內(nèi)存開銷對(duì)比圖如下:為什么 OOM:ZeRO-3 用狀態(tài)機(jī)替代常駐副本數(shù)據(jù)并行下,參數(shù)是常駐副本:每張卡從初始化到結(jié)束都持有全部權(quán)重。ZeRO-3 反其道而行,把參數(shù)狀態(tài)顯式建模為一個(gè)三態(tài)枚舉,每個(gè)參數(shù)在任意時(shí)刻只能處于其中一種狀態(tài):# 來源deepspeed/runtime/zero/partition_parameters.py class ZeroParamStatus(Enum): # parameters are fully present and ready for use on all processes AVAILABLE 1 # parameters are either partitioned or remote in some or all process NOT_AVAILABLE 2 # parameters are being gathered. INFLIGHT 3設(shè)計(jì)決策鏈在這里很清楚:為什么不做全量加載延遲釋放:常駐副本就是數(shù)據(jù)并行,顯存天花板等于模型大小,分區(qū)毫無意義。為什么不做隱式懶加載:PyTorch 的Parameter默認(rèn)隨時(shí)可讀,但 ZeRO-3 下未聚集的參數(shù)底層存儲(chǔ)只有 0 個(gè)元素,誤讀會(huì)得到 shape 為(0,)的張量,前向直接算出 NaN。所以必須用顯式狀態(tài)位(ds_status)在讀寫路徑上攔截。為什么引入 INFLIGHT 中間態(tài):聚集走異步allgather,提交后數(shù)據(jù)尚未落地。沒有第三態(tài)就無法防止聚集還沒完成就被釋放或同一參數(shù)被重復(fù)提交兩類競(jìng)態(tài)——代碼里InflightParamRegistry會(huì)對(duì)重復(fù)注冊(cè)直接拋RuntimeError(partitioned_param_coordinator.py)。預(yù)取、使用、釋放:三層鉤子怎么掛進(jìn)前向/反向狀態(tài)機(jī)只是賬本,真正驅(qū)動(dòng)狀態(tài)遷移的是掛在每個(gè)子模塊上的四組鉤子,入口在 parameter_offload.py 的pre/post_sub_module_forward/backward_function。以反向?yàn)槔?# 來源deepspeed/runtime/zero/parameter_offload.py def pre_sub_module_backward_function(self, sub_module): param_coordinator self.get_param_coordinator() param_coordinator.trace_prologue(sub_module) if param_coordinator.is_record_trace(): param_coordinator.record_module(sub_module) param_coordinator.fetch_sub_module(sub_module, forwardFalse)fetch_sub_module內(nèi)部做兩件事:把本模塊NOT_AVAILABLE的參數(shù)批量 allgather 提交(forward_fetch_submit與forward_prefetch_submit是兩個(gè)獨(dú)立計(jì)時(shí)項(xiàng),見 coordinator 類常量),然后從預(yù)取隊(duì)列里彈出下一個(gè)即將使用的模塊提前發(fā)起通信。這就是預(yù)取的落點(diǎn)——不是猜,而是第一步前向時(shí)錄制的模塊訪問順序(trace)決定的,ZeRoTraceMode會(huì)在錄制/復(fù)用/失配三種模式間切換。兩個(gè)值得注意的工程細(xì)節(jié):通信 dtype 自動(dòng)降級(jí):get_allgather_dtype檢測(cè)參數(shù)上是否打了comm_dtype標(biāo)記(autocast 場(chǎng)景),有則用 bf16/fp16 通信,沒有則按原 dtype,通信量和顯存峰值都受它影響。釋放帶流同步保護(hù):free_param在丟棄存儲(chǔ)前先record_stream,防止 CUDA 異步執(zhí)行時(shí)內(nèi)核還在讀這塊顯存:# 來源deepspeed/runtime/zero/partition_parameters.py if get_accelerator().on_accelerator(param.data): if not get_accelerator().is_synchronized_device(): param.data.record_stream(get_accelerator().current_stream()) # param.data doesnt store anything meaningful in partitioned state param.data torch.empty(0, dtypeparam.dtype, deviceparam.device) param.ds_status ZeroParamStatus.NOT_AVAILABLE數(shù)據(jù)佐證:512 卡 1T 參數(shù),49 TFLOPS/卡ZeRO-3 比數(shù)據(jù)并行多出約 50% 的通信量(每層參數(shù)前后各一次 allgather),官方實(shí)現(xiàn)靠計(jì)算-通信重疊把它幾乎全部藏進(jìn)計(jì)算時(shí)間里。公開數(shù)據(jù)(引自 2021-03-08-zero3-offload.md):對(duì)比維度ZeRO-2 OffloadZeRO-3 Offload單 V100(32GB)1.5TB 內(nèi)存可訓(xùn)模型13B40B(約 3 倍)訓(xùn)練 1T 參數(shù)所需 GPU 數(shù)不可行512 張 V100(3D 并行需約 1600 張)1T 參數(shù)持續(xù)算力(512 V100)—25 PFLOPS,約 49 TFLOPS/卡相對(duì) ZeRO-2 Offload 單卡吞吐基線13B 模型上約 2 倍7 個(gè) stage3 配置項(xiàng):默認(rèn)值、作用與調(diào)錯(cuò)后果所有項(xiàng)定義在 zero/config.py,按配置項(xiàng)逐行給出:配置項(xiàng)默認(rèn)值作用調(diào)錯(cuò)會(huì)怎樣stage13 才啟用參數(shù)分區(qū)停留在 2 只分區(qū)梯度優(yōu)化器狀態(tài),70B 依然放不下stage3_prefetch_bucket_size5e7異步 allgather 預(yù)取的桶大小(元素?cái)?shù))太小→通信碎、帶寬利用率低;太大→臨時(shí)聚集顯存峰值升高stage3_max_live_parameters1e9允許同時(shí)駐留 GPU 的參數(shù)元素上限調(diào)大省通信、費(fèi)顯存;OOM 時(shí)優(yōu)先調(diào)它stage3_max_reuse_distance1e9兩次使用距離超過該值就釋放參數(shù)調(diào)小→重算場(chǎng)景反復(fù) gather;調(diào)大→顯存堆積stage3_param_persistence_threshold1e5小于該規(guī)模的參數(shù)永久駐留不分區(qū)默認(rèn)值對(duì)大量小 bias 已經(jīng)有效,一般不用動(dòng)offload_param.devicenone分片卸載到cpu/nvme設(shè)為 cpu 后吞吐下降約 20%-30%,換取數(shù)倍顯存空間stage3_gather_16bit_weights_on_model_savefalse存 checkpoint 時(shí)自動(dòng)聚齊全量權(quán)重不開則拿不到可直接加載的完整權(quán)重文件動(dòng)手鏈路:從單行配置到生產(chǎn)級(jí)第 1 步,最簡(jiǎn)配置——只加一行stage: 3:{ zero_optimization: { stage: 3 } }做錯(cuò)了會(huì)怎樣:模型超過單卡顯存時(shí),deepspeed.initialize之后第一次前向就會(huì) OOM——因?yàn)?8B 以上模型不能直接在 GPU 上完整構(gòu)造,需要配合第 3 步。第 2 步,模型構(gòu)造期就分區(qū)——用zero.Init上下文,參數(shù)一分配就切片到各卡:# 來源docs/_tutorials/zero.md with deepspeed.zero.Init(data_parallel_groupmpu.get_data_parallel_group(), remote_deviceget_args().remote_device, enabledget_args().zero_stage 3): model GPT2Model(num_tokentypes0, parallel_outputTrue)做錯(cuò)了會(huì)怎樣:不用zero.Init直接model.to(device),完整權(quán)重先在 GPU 上存在一次再分區(qū),70B 在這一步就 OOM。第 3 步,生產(chǎn)級(jí)疊加卸載與調(diào)參:{ zero_optimization: { stage: 3, // 參數(shù)梯度優(yōu)化器狀態(tài)全分區(qū) contiguous_gradients: true, // 梯度連續(xù)化,減少碎片 stage3_prefetch_bucket_size: 1e7, // 預(yù)取桶,通信帶寬不足時(shí)調(diào)大 stage3_max_live_parameters: 1e9, // 顯存緊張時(shí)調(diào)小 stage3_param_persistence_threshold: 1e5, // 小參數(shù)常駐 stage3_gather_16bit_weights_on_model_save: true, offload_optimizer: { device: cpu }, // 優(yōu)化器狀態(tài)卸載 offload_param: { device: cpu } // 參數(shù)分片卸載,最后手段 } }調(diào)參順序建議:先stage3_max_live_parameters(顯存旋鈕),再stage3_prefetch_bucket_size(通信旋鈕),offload_param放最后——CPU 帶寬是 PCIe 的零頭,每多卸一層吞吐?lián)p失越明顯。三個(gè)踩坑實(shí)錄??坑一:模塊外讀權(quán)重,前向輸出全 NaN?，F(xiàn)象:某個(gè)模塊在forward里用到了兄弟模塊的weight,結(jié)果算出 NaN,單步調(diào)試發(fā)現(xiàn)該參數(shù)ds_status NOT_AVAILABLE,讀到的數(shù)據(jù)是 0 元素張量。 根因:ZeRO-3 只協(xié)調(diào)屬主模塊內(nèi)的參數(shù)訪問,跨模塊引用不在它的依賴圖里,沒人負(fù)責(zé)為它發(fā)起 gather。 修復(fù):在需要該參數(shù)的模塊上注冊(cè)外部依賴,register_external_parameter會(huì)把它并入該模塊的前向/反向 gather 范圍:# 來源deepspeed/runtime/zero/partition_parameters.py class ModuleZ3(torch.nn.Module): def __init__(self, *args): super().__init__(*args) self.layer1 SomeLayer() self.layer2 OtherLayer() deepspeed.zero.register_external_parameter(self, self.layer1.weight)順帶一提:如果參數(shù)是掛在OrderedDict里、且訪問發(fā)生在前向內(nèi),ZeROOrderedDict.__getitem__會(huì)自動(dòng)替你注冊(cè)并 gather(parameter_offload.py 第 82 行),這是為什么有些代碼沒注冊(cè)也能跑?？佣?Cannot free a ZeRO-3 parameter while it is still active in submodules?，F(xiàn)象:訓(xùn)練中途free_param拋RuntimeError。 根因:在GatheredParameters上下文里修改了參數(shù),但modifier_rankNone,各 rank 的副本狀態(tài)不一致,屬主模塊還沒釋放,coordinator 就嘗試回收。 修復(fù):給GatheredParameters(..., modifier_rankrank)指定一個(gè) rank 作為修改者,由它廣播更新,保證所有副本一致后再釋放?？尤?log_trace_cache_warnings持續(xù)告警,預(yù)取失效?，F(xiàn)象:日志反復(fù)出現(xiàn) trace 與當(dāng)前 forward/backward 不匹配,吞吐量明顯下降。 根因:訓(xùn)練循環(huán)里用了動(dòng)態(tài)控制流(按 batch 長(zhǎng)度走不同分支),第一步錄制的模塊訪問順序(trace)不再覆蓋實(shí)際路徑,預(yù)取退化為同步 gather。 修復(fù):固定控制流路徑;或打開log_trace_cache_warnings: true觀察失配頻率,必要時(shí)改用stage3_max_reuse_distance放寬釋放策略兜底。行動(dòng)出口 可直接執(zhí)行的命令序列:git clone https://gitcode.com/GitHub_Trending/de/DeepSpeed cd DeepSpeed pip install -e . deepspeed --num_gpus8 train.py \ --deepspeed zero3_config.json \ --deepspeed_autotuning_profile # 先用默認(rèn) 3 跑通,再按上表逐項(xiàng)調(diào)進(jìn)一步閱讀路徑(均為倉庫內(nèi)相對(duì)路徑):官方教程與 JSON 示例:docs/_tutorials/zero.md全部 zero 配置項(xiàng)說明:docs/_pages/config-json.md狀態(tài)機(jī)與聚集實(shí)現(xiàn):deepspeed/runtime/zero/partition_parameters.py鉤子掛載與外部參數(shù)自動(dòng)注冊(cè):deepspeed/runtime/zero/parameter_offload.py發(fā)布數(shù)據(jù)與 3D 并行對(duì)比:docs/_posts/2021-03-08-zero3-offload.md【免費(fèi)下載鏈接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.項(xiàng)目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考 DeepSpeed ZeRO-3 參數(shù)分片實(shí)戰(zhàn)：按需聚集機(jī)制拆解與 stage3 配置調(diào)優(yōu)指南
DeepSpeed ZeRO-3 參數(shù)分片實(shí)戰(zhàn)按需聚集機(jī)制拆解與 stage3 配置調(diào)優(yōu)指南【免費(fèi)下載鏈接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.項(xiàng)目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed[rank 1] CUDA out of memory. Tried to allocate 112.00 GiB. Process 0 has 27.97 GiB in use. Of which 26.41 GiB is allocated by PyTorch.微調(diào) 70B 模型時(shí),普通數(shù)據(jù)并行要求每張 GPU 持有完整參數(shù)副本,單卡直接 OOM。DeepSpeed ZeRO-3 的參數(shù)分區(qū)(把參數(shù)切成 N 份,每張卡只留自己那 1/N)配合按需聚集,把單卡顯存占用從全量副本降到局部分片臨時(shí)聚集。本文從ZeroParamStatus狀態(tài)機(jī)講起,覆蓋預(yù)取、釋放全生命周期,以及 7 個(gè) stage3 關(guān)鍵參數(shù)的調(diào)優(yōu)邊界。訓(xùn)練 1T 參數(shù)模型時(shí),每張卡實(shí)際駐留的參數(shù)只有總量的 1/512,這是 ZeRO-3 與張量/流水線并行的核心差異。官方發(fā)布的內(nèi)存開銷對(duì)比圖如下:為什么 OOM:ZeRO-3 用狀態(tài)機(jī)替代常駐副本數(shù)據(jù)并行下,參數(shù)是常駐副本:每張卡從初始化到結(jié)束都持有全部權(quán)重。ZeRO-3 反其道而行,把參數(shù)狀態(tài)顯式建模為一個(gè)三態(tài)枚舉,每個(gè)參數(shù)在任意時(shí)刻只能處于其中一種狀態(tài):# 來源deepspeed/runtime/zero/partition_parameters.py class ZeroParamStatus(Enum): # parameters are fully present and ready for use on all processes AVAILABLE 1 # parameters are either partitioned or remote in some or all process NOT_AVAILABLE 2 # parameters are being gathered. INFLIGHT 3設(shè)計(jì)決策鏈在這里很清楚:為什么不做全量加載延遲釋放:常駐副本就是數(shù)據(jù)并行,顯存天花板等于模型大小,分區(qū)毫無意義。為什么不做隱式懶加載:PyTorch 的Parameter默認(rèn)隨時(shí)可讀,但 ZeRO-3 下未聚集的參數(shù)底層存儲(chǔ)只有 0 個(gè)元素,誤讀會(huì)得到 shape 為(0,)的張量,前向直接算出 NaN。所以必須用顯式狀態(tài)位(ds_status)在讀寫路徑上攔截。為什么引入 INFLIGHT 中間態(tài):聚集走異步allgather,提交后數(shù)據(jù)尚未落地。沒有第三態(tài)就無法防止聚集還沒完成就被釋放或同一參數(shù)被重復(fù)提交兩類競(jìng)態(tài)——代碼里InflightParamRegistry會(huì)對(duì)重復(fù)注冊(cè)直接拋RuntimeError(partitioned_param_coordinator.py)。預(yù)取、使用、釋放:三層鉤子怎么掛進(jìn)前向/反向狀態(tài)機(jī)只是賬本,真正驅(qū)動(dòng)狀態(tài)遷移的是掛在每個(gè)子模塊上的四組鉤子,入口在 parameter_offload.py 的pre/post_sub_module_forward/backward_function。以反向?yàn)槔?# 來源deepspeed/runtime/zero/parameter_offload.py def pre_sub_module_backward_function(self, sub_module): param_coordinator self.get_param_coordinator() param_coordinator.trace_prologue(sub_module) if param_coordinator.is_record_trace(): param_coordinator.record_module(sub_module) param_coordinator.fetch_sub_module(sub_module, forwardFalse)fetch_sub_module內(nèi)部做兩件事:把本模塊NOT_AVAILABLE的參數(shù)批量 allgather 提交(forward_fetch_submit與forward_prefetch_submit是兩個(gè)獨(dú)立計(jì)時(shí)項(xiàng),見 coordinator 類常量),然后從預(yù)取隊(duì)列里彈出下一個(gè)即將使用的模塊提前發(fā)起通信。這就是預(yù)取的落點(diǎn)——不是猜,而是第一步前向時(shí)錄制的模塊訪問順序(trace)決定的,ZeRoTraceMode會(huì)在錄制/復(fù)用/失配三種模式間切換。兩個(gè)值得注意的工程細(xì)節(jié):通信 dtype 自動(dòng)降級(jí):get_allgather_dtype檢測(cè)參數(shù)上是否打了comm_dtype標(biāo)記(autocast 場(chǎng)景),有則用 bf16/fp16 通信,沒有則按原 dtype,通信量和顯存峰值都受它影響。釋放帶流同步保護(hù):free_param在丟棄存儲(chǔ)前先record_stream,防止 CUDA 異步執(zhí)行時(shí)內(nèi)核還在讀這塊顯存:# 來源deepspeed/runtime/zero/partition_parameters.py if get_accelerator().on_accelerator(param.data): if not get_accelerator().is_synchronized_device(): param.data.record_stream(get_accelerator().current_stream()) # param.data doesnt store anything meaningful in partitioned state param.data torch.empty(0, dtypeparam.dtype, deviceparam.device) param.ds_status ZeroParamStatus.NOT_AVAILABLE數(shù)據(jù)佐證:512 卡 1T 參數(shù),49 TFLOPS/卡ZeRO-3 比數(shù)據(jù)并行多出約 50% 的通信量(每層參數(shù)前后各一次 allgather),官方實(shí)現(xiàn)靠計(jì)算-通信重疊把它幾乎全部藏進(jìn)計(jì)算時(shí)間里。公開數(shù)據(jù)(引自 2021-03-08-zero3-offload.md):對(duì)比維度ZeRO-2 OffloadZeRO-3 Offload單 V100(32GB)1.5TB 內(nèi)存可訓(xùn)模型13B40B(約 3 倍)訓(xùn)練 1T 參數(shù)所需 GPU 數(shù)不可行512 張 V100(3D 并行需約 1600 張)1T 參數(shù)持續(xù)算力(512 V100)—25 PFLOPS,約 49 TFLOPS/卡相對(duì) ZeRO-2 Offload 單卡吞吐基線13B 模型上約 2 倍7 個(gè) stage3 配置項(xiàng):默認(rèn)值、作用與調(diào)錯(cuò)后果所有項(xiàng)定義在 zero/config.py,按配置項(xiàng)逐行給出:配置項(xiàng)默認(rèn)值作用調(diào)錯(cuò)會(huì)怎樣stage13 才啟用參數(shù)分區(qū)停留在 2 只分區(qū)梯度優(yōu)化器狀態(tài),70B 依然放不下stage3_prefetch_bucket_size5e7異步 allgather 預(yù)取的桶大小(元素?cái)?shù))太小→通信碎、帶寬利用率低;太大→臨時(shí)聚集顯存峰值升高stage3_max_live_parameters1e9允許同時(shí)駐留 GPU 的參數(shù)元素上限調(diào)大省通信、費(fèi)顯存;OOM 時(shí)優(yōu)先調(diào)它stage3_max_reuse_distance1e9兩次使用距離超過該值就釋放參數(shù)調(diào)小→重算場(chǎng)景反復(fù) gather;調(diào)大→顯存堆積stage3_param_persistence_threshold1e5小于該規(guī)模的參數(shù)永久駐留不分區(qū)默認(rèn)值對(duì)大量小 bias 已經(jīng)有效,一般不用動(dòng)offload_param.devicenone分片卸載到cpu/nvme設(shè)為 cpu 后吞吐下降約 20%-30%,換取數(shù)倍顯存空間stage3_gather_16bit_weights_on_model_savefalse存 checkpoint 時(shí)自動(dòng)聚齊全量權(quán)重不開則拿不到可直接加載的完整權(quán)重文件動(dòng)手鏈路:從單行配置到生產(chǎn)級(jí)第 1 步,最簡(jiǎn)配置——只加一行stage: 3:{ zero_optimization: { stage: 3 } }做錯(cuò)了會(huì)怎樣:模型超過單卡顯存時(shí),deepspeed.initialize之后第一次前向就會(huì) OOM——因?yàn)?8B 以上模型不能直接在 GPU 上完整構(gòu)造,需要配合第 3 步。第 2 步,模型構(gòu)造期就分區(qū)——用zero.Init上下文,參數(shù)一分配就切片到各卡:# 來源docs/_tutorials/zero.md with deepspeed.zero.Init(data_parallel_groupmpu.get_data_parallel_group(), remote_deviceget_args().remote_device, enabledget_args().zero_stage 3): model GPT2Model(num_tokentypes0, parallel_outputTrue)做錯(cuò)了會(huì)怎樣:不用zero.Init直接model.to(device),完整權(quán)重先在 GPU 上存在一次再分區(qū),70B 在這一步就 OOM。第 3 步,生產(chǎn)級(jí)疊加卸載與調(diào)參:{ zero_optimization: { stage: 3, // 參數(shù)梯度優(yōu)化器狀態(tài)全分區(qū) contiguous_gradients: true, // 梯度連續(xù)化,減少碎片 stage3_prefetch_bucket_size: 1e7, // 預(yù)取桶,通信帶寬不足時(shí)調(diào)大 stage3_max_live_parameters: 1e9, // 顯存緊張時(shí)調(diào)小 stage3_param_persistence_threshold: 1e5, // 小參數(shù)常駐 stage3_gather_16bit_weights_on_model_save: true, offload_optimizer: { device: cpu }, // 優(yōu)化器狀態(tài)卸載 offload_param: { device: cpu } // 參數(shù)分片卸載,最后手段 } }調(diào)參順序建議:先stage3_max_live_parameters(顯存旋鈕),再stage3_prefetch_bucket_size(通信旋鈕),offload_param放最后——CPU 帶寬是 PCIe 的零頭,每多卸一層吞吐?lián)p失越明顯。三個(gè)踩坑實(shí)錄??坑一:模塊外讀權(quán)重,前向輸出全 NaN?，F(xiàn)象:某個(gè)模塊在forward里用到了兄弟模塊的weight,結(jié)果算出 NaN,單步調(diào)試發(fā)現(xiàn)該參數(shù)ds_status NOT_AVAILABLE,讀到的數(shù)據(jù)是 0 元素張量。 根因:ZeRO-3 只協(xié)調(diào)屬主模塊內(nèi)的參數(shù)訪問,跨模塊引用不在它的依賴圖里,沒人負(fù)責(zé)為它發(fā)起 gather。 修復(fù):在需要該參數(shù)的模塊上注冊(cè)外部依賴,register_external_parameter會(huì)把它并入該模塊的前向/反向 gather 范圍:# 來源deepspeed/runtime/zero/partition_parameters.py class ModuleZ3(torch.nn.Module): def __init__(self, *args): super().__init__(*args) self.layer1 SomeLayer() self.layer2 OtherLayer() deepspeed.zero.register_external_parameter(self, self.layer1.weight)順帶一提:如果參數(shù)是掛在OrderedDict里、且訪問發(fā)生在前向內(nèi),ZeROOrderedDict.__getitem__會(huì)自動(dòng)替你注冊(cè)并 gather(parameter_offload.py 第 82 行),這是為什么有些代碼沒注冊(cè)也能跑?？佣?Cannot free a ZeRO-3 parameter while it is still active in submodules?，F(xiàn)象:訓(xùn)練中途free_param拋RuntimeError。 根因:在GatheredParameters上下文里修改了參數(shù),但modifier_rankNone,各 rank 的副本狀態(tài)不一致,屬主模塊還沒釋放,coordinator 就嘗試回收。 修復(fù):給GatheredParameters(..., modifier_rankrank)指定一個(gè) rank 作為修改者,由它廣播更新,保證所有副本一致后再釋放?？尤?log_trace_cache_warnings持續(xù)告警,預(yù)取失效?，F(xiàn)象:日志反復(fù)出現(xiàn) trace 與當(dāng)前 forward/backward 不匹配,吞吐量明顯下降。 根因:訓(xùn)練循環(huán)里用了動(dòng)態(tài)控制流(按 batch 長(zhǎng)度走不同分支),第一步錄制的模塊訪問順序(trace)不再覆蓋實(shí)際路徑,預(yù)取退化為同步 gather。 修復(fù):固定控制流路徑;或打開log_trace_cache_warnings: true觀察失配頻率,必要時(shí)改用stage3_max_reuse_distance放寬釋放策略兜底。行動(dòng)出口 可直接執(zhí)行的命令序列:git clone https://gitcode.com/GitHub_Trending/de/DeepSpeed cd DeepSpeed pip install -e . deepspeed --num_gpus8 train.py \ --deepspeed zero3_config.json \ --deepspeed_autotuning_profile # 先用默認(rèn) 3 跑通,再按上表逐項(xiàng)調(diào)進(jìn)一步閱讀路徑(均為倉庫內(nèi)相對(duì)路徑):官方教程與 JSON 示例:docs/_tutorials/zero.md全部 zero 配置項(xiàng)說明:docs/_pages/config-json.md狀態(tài)機(jī)與聚集實(shí)現(xiàn):deepspeed/runtime/zero/partition_parameters.py鉤子掛載與外部參數(shù)自動(dòng)注冊(cè):deepspeed/runtime/zero/parameter_offload.py發(fā)布數(shù)據(jù)與 3D 并行對(duì)比:docs/_posts/2021-03-08-zero3-offload.md【免費(fèi)下載鏈接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.項(xiàng)目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考