深度解析 parameter_server 密鑰分配機(jī)制大規(guī)模參數(shù)分片如何優(yōu)雅落地【免費(fèi)下載鏈接】parameter_servermoved to https://github.com/dmlc/ps-lite項(xiàng)目地址: https://gitcode.com/gh_mirrors/pa/parameter_serverparameter_server參數(shù)服務(wù)器框架的 KeyRange 密鑰分配機(jī)制是分布式訓(xùn)練中誰負(fù)責(zé)存哪些參數(shù)這個(gè)問題的核心答案。它用一段簡單的閉區(qū)間 [begin, end)把整個(gè) 64 位密鑰空間均勻切給各個(gè) Server 節(jié)點(diǎn)再自動(dòng)按密鑰歸屬把訓(xùn)練消息拆包發(fā)送——新手無需手寫任何路由邏輯就能讓億級參數(shù)輕松落到多臺機(jī)器上。下面帶你一步步看穿這套機(jī)制。一、為什么需要參數(shù)分片在分布式機(jī)器學(xué)習(xí)里模型參數(shù)可能有上億個(gè)比如大規(guī)模推薦系統(tǒng)的稀疏特征單臺機(jī)器存不下、也算不動(dòng)。parameter_server 采用經(jīng)典的Worker-Server 架構(gòu)Scheduler調(diào)度器負(fù)責(zé)節(jié)點(diǎn)注冊、分配密鑰范圍、廣播拓?fù)渥兓疭erver 組每臺 Server 只保存自己名分內(nèi)的那段密鑰對應(yīng)的參數(shù)Worker 組各自讀數(shù)據(jù)、算梯度把結(jié)果推給負(fù)責(zé)的 Server再拉回更新后的參數(shù)。關(guān)鍵問題隨之而來一臺 Worker 算出的梯度里可能同時(shí)包含屬于 3 臺不同 Server 的參數(shù)怎么保證每個(gè) key 都發(fā)給正確的 Server且不重不漏答案就是 KeyRange 機(jī)制。二、Range一個(gè)閉區(qū)間撐起全局路由 整個(gè)機(jī)制的地基是 range.h 中的模板類RangeT它表示半開區(qū)間[begin, end)只提供幾個(gè)恰到好處的操作操作作用在分片中的角色EvenDivide(n, i)把區(qū)間均分為 n 份取第 i 份調(diào)度器給 Server 分地SetIntersection(other)求兩個(gè)區(qū)間的交集判斷這包消息和你有沒有關(guān)系Project(v)把值投影到區(qū)間邊界切包時(shí)定位切割點(diǎn)contains(k)判斷 key 是否落在區(qū)間內(nèi)本地索引換算注意區(qū)間用長雙精度long double做均分計(jì)算即使密鑰空間是2^64 - 1這么大切給任意數(shù)量的 Server 也基本能做到誤差為 0。這正是大規(guī)模分片能優(yōu)雅落地的第一個(gè)細(xì)節(jié)。三、調(diào)度器如何自動(dòng)分配密鑰范圍全局密鑰范圍由兩個(gè)啟動(dòng)參數(shù)定義見 manager.cckey_start默認(rèn) 0和key_end默認(rèn) uint64 最大值。也就是說默認(rèn)整個(gè) 64 位空間都是可用密鑰覆蓋任何規(guī)模的稀疏特征。當(dāng)一臺 Server 啟動(dòng)后它向調(diào)度器發(fā)送REGISTER_NODE請求調(diào)度器隨即執(zhí)行 assigner.h 中NodeAssigner::assign()的兩步動(dòng)作定排名為 Server 分配自增的rank0、1、2……Worker 同理分地盤調(diào)用key_range_.EvenDivide(num_servers, rank)把全局范圍的第 rank 份寫入該節(jié)點(diǎn)信息。舉一個(gè)最直觀的例子3 臺 Server、密鑰范圍 [0, 120)Server 0 → [0, 40) Server 1 → [40, 80) Server 2 → [80, 120)之后調(diào)度器通過ADD_NODE把節(jié)點(diǎn) ID 角色 密鑰范圍廣播給全網(wǎng)。于是每個(gè)節(jié)點(diǎn)都持有一份完整的分片地圖后續(xù)路由不需要任何中心化查詢——這就是分布式系統(tǒng)的精髓控制面只在啟動(dòng)時(shí)集中決策數(shù)據(jù)面全程去中心化。四、消息自動(dòng)拆包Slice 的一刀多段 ??Worker 算完一輪梯度后調(diào)用參數(shù)對象的Push把 (key, value) 推給整個(gè) Server 組。系統(tǒng)并不會傻乎乎地廣播給所有 Server而是觸發(fā) customer.h 中Customer::Slice的分包流程核心實(shí)現(xiàn)在 message.h 的SliceKOFVMessage前提key 有序。Worker 側(cè)要求消息里的 key 已排序參數(shù)類如 kv_map.h、kv_vector.h 都遵循這一約定這為二分查找鋪平了道路定位切割點(diǎn)對每臺 Server 的密鑰區(qū)間端點(diǎn)做Project投影再用lower_bound在有序 key 數(shù)組中二分出切分位置按段切包第 i 段消息只保留落在 Server i 區(qū)間內(nèi)的 key 和對應(yīng) value 切片空交集剪枝若某 Server 的區(qū)間與本次消息的密鑰范圍沒有交集SetIntersection為空該分包直接標(biāo)記valid false一條網(wǎng)絡(luò)流量都不會浪費(fèi)。一句話總結(jié)一條包含十萬個(gè) key 的梯度消息會在發(fā)出瞬間被精確裁剪成 N 份每份只含有對方負(fù)責(zé)的參數(shù)。五、Server 端只收自己的賬本 Server 側(cè)的參數(shù)類實(shí)現(xiàn)得非常薄。以KVMap為例kv_map.hGetValue(msg)按 key 從本地哈希表取值填回消息SetValue(msg)把收到的 value 寫進(jìn)本地表Slice默認(rèn)轉(zhuǎn)調(diào)SliceKOFVMessage保證 pull 響應(yīng)同樣按區(qū)間切分。應(yīng)用層如 darlin.h 里的稀疏模型還會把收到的全局密鑰區(qū)間換算成本地行下標(biāo)先SetIntersection判斷請求和自己有無交集無交集直接跳過省 CPU再用FindRange把全局區(qū)間映射成SizeR本地范圍連多線程分塊也復(fù)用同一個(gè)EvenDivide——區(qū)間抽象在網(wǎng)絡(luò)分片 → 本地計(jì)算兩級里被統(tǒng)一復(fù)用非常優(yōu)雅。六、實(shí)踐指南改密鑰范圍與節(jié)點(diǎn)重分配 ?如何自定義密鑰空間啟動(dòng)時(shí)傳--key_start0 --key_end10000000即可分片會自動(dòng)按新范圍均分。Server 數(shù)量變了怎么辦密鑰范圍與 Server 數(shù)解耦EvenDivide是純函數(shù)改num_servers后重新注冊即按新份數(shù)切分。項(xiàng)目內(nèi)置了演示程序 reassign_server_key_range_ps.cc展示如何對在線節(jié)點(diǎn)組執(zhí)行partitionServerKeyRange重新切分并下發(fā)UPDATE_NODE指令——即密鑰范圍支持運(yùn)行時(shí)重分配為集群擴(kuò)縮容預(yù)留了通道。新手落地三步走啟動(dòng) 1 個(gè) Scheduler--rolescheduler、N 個(gè) Server、M 個(gè) Worker確認(rèn)日志里每個(gè) Server 打印的[begin,end)區(qū)間互不重疊且首尾相接觀察 Worker 推送時(shí)每個(gè) Server 收到的 key 數(shù)量大致相等——均分即負(fù)載均衡。七、小結(jié)簡單即強(qiáng)大 ?parameter_server 的 KeyRange 機(jī)制沒有復(fù)雜的哈希環(huán)或一致性哈希僅靠區(qū)間均分 有序切包 交集剪枝三板斧就實(shí)現(xiàn)了大規(guī)模參數(shù)分片的完整閉環(huán)對調(diào)度器注冊即分片一次EvenDivide定終身對Worker無感路由消息自動(dòng)拆到正確的 Server對Server只存只算自己那一段水平擴(kuò)展毫無壓力。這種用最少的抽象解決最難的路由問題的設(shè)計(jì)正是參數(shù)服務(wù)器類框架值得新手研讀的地方。想要?jiǎng)邮謱?shí)驗(yàn)可以從 hello_ps.cc 這樣的最小示例跑起來再逐層深入 kv_layer_ps.cc 感受分片全流程?！久赓M(fèi)下載鏈接】parameter_servermoved to https://github.com/dmlc/ps-lite項(xiàng)目地址: https://gitcode.com/gh_mirrors/pa/parameter_server創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考