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