已开启
[RFC]: DCP 加载路径——广播下沉到读循环,读盘 rank 按字节均衡 #361
zhangbuxue创建于 18 天前
18 天前 添加了label:RFC
Zzhangbuxue
18 天前 关联了pull request:perf(dcp): overlap the replicated-shard broadcast with the read that feeds it
18 天前 关联了pull request:perf(dcp): overlap the replicated-shard broadcast with the read that feeds it
Zzhangbuxue
18 天前 删除了关联的pull request:perf(dcp): overlap the replicated-shard broadcast with the read that feeds it
18 天前 删除了关联的pull request:perf(dcp): overlap the replicated-shard broadcast with the read that feeds it
Zzhangbuxue
18 天前 关联了pull request:perf(dcp): overlap the replicated-shard broadcast with the read that feeds it
18 天前 关联了pull request:perf(dcp): overlap the replicated-shard broadcast with the read that feeds it
Zzhangbuxue
18 天前 删除了关联的pull request:perf(dcp): overlap the replicated-shard broadcast with the read that feeds it
18 天前 删除了关联的pull request:perf(dcp): overlap the replicated-shard broadcast with the read that feeds it
Zzhangbuxue
18 天前 关联了pull request:perf(dcp): overlap the replicated-shard broadcast with the read that feeds it
18 天前 关联了pull request:perf(dcp): overlap the replicated-shard broadcast with the read that feeds it
Zzhangbuxue
18 天前 修改标题为 “[RFC]: DCP 加载路径优化——读播流水化与逐分片开销削减”,原标题为“[RFC]: DCP 加载路径——广播下沉到读循环,读盘 rank 按字节均衡”
18 天前 修改标题为 “[RFC]: DCP 加载路径优化——读播流水化与逐分片开销削减”,原标题为“[RFC]: DCP 加载路径——广播下沉到读循环,读盘 rank 按字节均衡”
Zzhangbuxue
18 天前 修改标题为 “[RFC]: DCP 加载路径——广播下沉到读循环,读盘 rank 按字节均衡”,原标题为“[RFC]: DCP 加载路径优化——读播流水化与逐分片开销削减”
18 天前 修改标题为 “[RFC]: DCP 加载路径——广播下沉到读循环,读盘 rank 按字节均衡”,原标题为“[RFC]: DCP 加载路径优化——读播流水化与逐分片开销削减”
13 天前 修改了issue 的描述
12 天前 修改了issue 的描述
10 天前 修改了issue 的描述
9 天前 修改了issue 的描述
9 天前 修改了issue 的描述
9 天前 修改了issue 的描述
动机(Motivation)
DCP 加载是大模型作业启动与断点续训的主要耗时之一。
master上已经有"副本张量只读一次 + 广播"的能力(broadcast_from_minimum_rank),但当前形态存在四个结构性问题:读与通信完全串行。
should_load_shard在计划阶段把要广播的张量打上BROADCAST_INFO,execute_read把所有文件读完之后,在末尾调用一次broadcast_loaded_tensors统一发送。读的时候没有通信在走,发的时候没有读在走,两段时间无法互相掩盖。读盘 rank 固定为
min(group_ranks)。 同一组内的最小 rank 往往就是全局靠前的那几张卡,读盘压力反复压在少数卡上,其余卡空等。广播粒度过细。 一个张量一次广播。实测 4 卡 Ascend 上,64 KiB 与 1 MiB 的广播代价几乎相同(约 145 µs),要到几 MB 才跑满带宽。30B MoE 这类模型每个文件约 9000 个张量,大量小分片使时间花在发起通信而非搬运数据上。
读盘本身有可观的常数开销。 逐分片重开 safetensors 文件(0.70 ms vs 复用的 0.045 ms);逐分片重建
available_keys集合(约 7 ms × 1.8 万个分片)。目标设计
把广播下沉到读循环内部,与读重叠;并把"谁读、谁收"的决策载体从 state_dict 的张量属性搬到加载计划的
ReadItem上。BROADCAST_INFOReadItem.source上的BroadcastSourcemin(group_ranks)决策载体的搬迁是前提:信息落在计划上而非张量对象上,
execute_read才能在读到某个分片的当下就知道该发给谁,不必等整个 state_dict 填完。广播的单位是目的分片的整块本地 buffer,而非
ReadItem。重切分时一个目的分片可能由多个ReadItem从多个文件拼出,因此该分片的全部读必须先聚齐、整批落位后才发送,且计重时须按其全部ReadItem累加。参考实现见 PR !1322。
需要设计评审DFX建议
关键设计点
1. 全局分片序是正确性基础,不是性能选择。
所有 rank 以
sorted(key=(fqn, offset, index))独立推导分片顺序,结果天然一致。集合通信要求同一通信域的各 rank 以相同顺序入队,边读边发能成立正是依赖这一点——收发双方不需要额外一轮通信就能就顺序达成一致。2. 负载均衡的确定性要求。
分片按字节数从大到小排序,逐个交给组内已分配字节最少的 rank。字节数相等时必须保持
replicas的填充顺序(跟随 gather 序,各 rank 一致)。若各 rank 在此产生分歧,会选出不同的 reader 并互相等待对方的广播而挂死。这是 DFX 上最需要评审关注的一点。3. 建组决策必须全卡一致。
DeviceMesh.from_group只记录本 rank 所在的组,进程级缓存天然不对称。若一个 rank 因命中缓存而跳过建组、另一个去建,new_group(世界级集合操作)会挂死。设计上采用:all-gather 各 rank 缺哪些组 → 任一 rank 缺则全体新建,都不缺才复用。采用
new_group而非create_group:后者把 rank 列表当作模板展开成全世界的划分并写入全局缓存,且会拒绝 PP 下跨非相邻 stage 的绑定参数。4. 通信组销毁前必须排空设备流。
集合通信被 wait 之后仅保证与设备流有序,传输仍可能在队列中。此时拆除 communicator 会连带取消传输,接收方静默保留旧数据。销毁前调用
platform.synchronize(),且销毁本身为 best-effort(失败告警而非失败加载)。5. 执行序:先共享、后私有。
共享分片按全局序处理并 flush 之后,才处理无人共享的私有分片。先读没人要的分片会让等待广播的同伴一直空转。
6. 留在 host 内存的条目不进广播。
目的缓冲在 host 上的分片不打广播标签,每个 rank 各读各的,与 BYTE_IO 同一待遇。两条理由,按先后:一是不值得发——host 上留的东西又少又小,torch 的 AdamW 每个参数张量留一个 fp32 的
step标量(capturable=False时用torch.tensor(0.0)建,不带 device),九千个参数张量也才 36 KiB/卡,且它们同处一个文件、句柄已开、正好被第 5 点的"私有分片按文件成批读"覆盖;换成广播要付一次集合通信启动加 staging 两次拷贝,大集群下若组里真有 gloo 后端则是主机网络上的 TCP 树,单次时延就盖过读它。二是常常发不了——加速器库起的通信组没有 host 内存的后端,实测get_backend_config为npu:hccl,广播 host 张量直接RuntimeError: No backend type associated with device type cpu。判断落在
build_local_plan(整条 load 路径上唯一同时握着目的对象与其 ReadItem 的地方),结果记在新增的ReadItem.broadcastable上随identity()过 all_gather。必须过 gather:第 2 点的load_bytes是全局累加的贪心,一个 rank 少算一个分片就会让后续所有分片的 reader 选择偏掉;而"某 fqn 在不在 host 上"本 rank 只能对自己持有的 fqn 回答。收敛规则为任一 rank 报告不可发则整个分片不发——广播要全员参与,一个收不了就定了全体。替代方案及否决理由
方案 A:读线程按文件序走,主线程仍按全局序消费(中间用完成表代替 FIFO)。
动机是消除跨文件跳的开销。30B 实测反而慢 38%(15.3 s → 21.1 s,三个变体交错跑两轮,轮间偏差 0.2 s 内)。原因有二:
(fqn, offset, index)完全吻合(实测 9218/9218 匹配、零空洞),单文件内的读已经是顺序访问,没有可优化空间;copy_,两个访问序在同一文件上交错,打乱内核的顺序检测,预读失效。方案 B:保留末尾统一广播,仅做负载均衡与聚批。
改动更小,但读与通信仍然串行,拿不到重叠收益,也无法解决大模型下通信段独占时间的问题。
方案 C(早期版本中曾实现,现已移除):读预取线程领先主线程若干批次。
该方案在两种存储上均实测无收益,故从实现中删除,读改为惰性生成器(取用顺序与时机不变)。数据见"性能影响分析"。结论要点:mmap 上的 disk→host 不是一段可以搬到旁路线程的等待——内核 readahead 已将其异步化;预取线程不是把这段时间从主线程移走,而是在总量上新增一遍全量遍历的 CPU 开销。
对现有代码的影响范围
core/distributed_checkpoint/planner.pyBroadcastSource、LoadPlan.identity()、ReadItem.broadcastablecore/distributed_checkpoint/standard_planner.pybuild_global_plan重写,返回本 rank 的LoadPlan(原为list[LoadPlan]);移除should_load_shard;build_local_plan标注目的缓冲是否在 host 上core/distributed_checkpoint/filesystem_storage.pyexecute_read重写;_load_*族拆为_fetch_*与_apply_fetched;新增_OpenFilescore/distributed_checkpoint/util.pyBROADCAST_INFO/broadcast_loaded_tensors一套;新增BroadcastBatcher与建组、销毁、发送函数core/distributed_checkpoint/metadata.pydtype_element_sizecore/distributed_checkpoint/api.pyload()参数改名;计划选择逻辑随返回类型简化platform/{platform,torch/platform}.py向后兼容性
存在对外接口不兼容变更(一处改名 + 两个默认值翻转),详见下节。checkpoint 文件格式、目录结构、metadata 均不变,已保存的 checkpoint 无需转换:没有共享
.metadata的旧 checkpoint 会自动回退到{rank}.metadata并关掉 plan 交换。显式传use_collectives=False/broadcast_replicated_tensors=False时行为与当前一致。性能影响分析
各优化项的微基准(均在 8× Ascend 910B3 上测得,已记录于对应代码注释):
端到端(8 卡 Ascend 910B3,本地 NVMe),10.45 GiB / 3147 张量的 Qwen3-30B-A3B 形状 checkpoint(由
master侧写入):master放大到 48 层(56.89 GiB / 18867 张量),broadcast=ON:
master166.85 s → 本设计 16.86 s。按HP_LOG_CONFIG=DCP:INFO拆解,master侧broadcast_loaded_tensors占execute_read的 92%,其中all_gather_object8.96 s、建组与 3147 次广播约 16 s、销毁 0.66 s;本设计将 3147 个分片合并为 410 次广播,且全程无新建通信组。值得评审注意的是第一行:
master上开启广播反而比不开慢约 8.5 倍,即现有的broadcast_from_minimum_rank在这个负载形状下是净亏损,实际不具备可用性。本设计的目标应理解为让一个当前开启即亏损的特性变得可用。读侧优化的边界(据此否决方案 C)。 在两种差异极大的存储上分别复测,每次读前
posix_fadvise(DONTNEED)清缓存并以mincore复核驻留率归零:本地 NVMe(裸盘 8 文件并行冷读 6.65 GiB/s,48 层 56.89 GiB):
master共享存储(同机 NFS,裸盘 8 文件并行冷读 0.11 GiB/s,10.45 GiB / 3147 张量):
T_io(时间下界)master共享存储上加载时间贴在裸读下界上方 5%(100.0 vs 95.3 s),说明管道已被内核 readahead 喂满;五种配置极差 0.16 s。这同时给出一个对本设计有利的旁证:共享存储上瓶颈是存储带宽本身,唯一有效的杠杆是少读字节,正是"副本分片只读一次 + 广播"所做的事。
尚未覆盖:共享存储上
broadcast=ON/OFF的直接对比(见下方后续项)。已结项:万卡下的读序与文件句柄
master的execute_read是 file-major 的(_group_items_by_file(plan)对整个 plan 分组,每个文件只开一次);本设计为了"分片就位即发"改成了 shard-major,文件局部性由_OpenFiles(LRU,torch 8 个)兜底。RFC 初稿曾就此列出两个候选方案征求评审,现已全部实测,结论如下。先厘清一个前提:在不做 resharding 的主场景里,读侧和写侧本来就是对齐的,因为两侧跑的是同一套贪心(大在前、给当前字节最少的那个、并列取小 rank),喂的是同一批分片。8 卡实测(3147 张量、TP+DP、broadcast=ON)8/8 全中,每个 rank 读的正是它自己写的那个文件:
因此下述两个候选在主场景下都不成立:
LRU=256 w=1在每一种下都优于LRU=8 w=256,而分窗对"一片一文件、邻片跳文件"这一真实形态毫无作用。已实现后回退。实际采纳的是第三条,起因于一个此前未识别的问题:非张量(BYTE_IO)条目的读放大。
非张量条目在 save 侧照常参与去冗(
plan_ownership_masks不区分类型,tensor_storage_size()返回None时按权重 1 计),被分给"当前字节最少"的 rank;而 load 侧build_global_plan只标记LoadItemType.TENSOR,非张量条目永不广播。两条合起来:每个 rank 都要把自己需要的全部非张量条目读一遍,而它们散在多个 rank 的.bytes文件里,且_fetch_bytes_file不走句柄缓存,逐分片读时每条都真开一次文件。4 卡实测(40 个嵌套 dict,flatten_state_dict摊成 80 个叶子):open()次数改动:
execute_read把"无人共享的分片"按其所在文件归组,一个文件一次读。这类分片没有任何 rank 在等,因此不受"全局一致顺序"约束,可以任意重排;共享分片仍严格按各 rank 一致的顺序逐片读发,广播时序不变。同时删掉读路径上两处冗余检查,它们说的都只是下一行自己会说的话:
set(f.keys())成员预检查File does not contain tensor <name>os.path.exists存在性预检查open均抛FileNotFoundError并带路径后者连"路径是目录"都挡不住(
exists()对目录返回 True),新增用例钉住"缺文件由底层报错并指明文件名"。仍未覆盖:换切分维度(
Shard(0)→Shard(1))的重切分下句柄抖动的实测。该场景按上述取舍只保正确性,resharding_apiST 已覆盖其正确性(含 broadcast 开关两种组合)。DFX 与可测试性
新增用例覆盖:负载均衡的确定性(与 plan 内项序无关)、无分片的 rank 仍参与 all-gather、仅部分 rank 持有的组全体重建、销毁前排空设备流、读失败必须导致加载失败、批内数据回到各自分片;重切分场景下同一
dest_index的多个ReadItem须全部标记同一BroadcastSource、且分片按其全部ReadItem累加计重。新增用例均经变异验证(人为引入对应缺陷后,确认目标用例失败)。ST 覆盖 PP 场景(含跨非相邻 stage 的绑定参数)与重切分 × 广播组合,且已在真卡 HCCL 上复跑——这一步不可省略:销毁通信组前的
platform.synchronize()排空的是设备流,而 gloo 是主机侧的、没有设备流,该路径在 CPU 后端下从未被真正执行。已处理:
load()与save()的默认值不对称。load()的use_collectives原先默认False(save()默认True),而广播只在前者为真时生效,于是load(broadcast_replicated_tensors=True)单独传是静默失效的——新增 ST 曾因此假通过。更要紧的是训练器dcp_checkpointer.py:630两个都没传,本 RFC 做的这条路径真实训练从未走到过。两个默认值现已都改为True。改默认值带来一条约束(
save()一直有):默认通信组里每个 rank 必须一起调load(),plan 交换跑在该组上,缺席者会让其余 rank 卡在凑不齐的 gather 上。真卡实测:test_dcp_save_and_load_save_8card_load_4card显式让一半 rank 跳过load(),翻默认后 503 s 超时死锁,改为全员调用、不参与者传空 state_dict 后 120 s 通过。生产上 save 与 load 是两个任务、各自 world 内全员参与,不存在这种用法。而
broadcast_replicated_tensors=True配use_collectives=False不应报错,这是之前建议里判断错的一点。这个组合的结果是:plan 不经 all_gather →build_global_plan首行返回 → 没有 item 拿到source→ 不建组、不发任何集合通信 → 每个 rank 自己读自己要的。数据完整、数值正确、不会挂,只是退化成master的行为。而integration/llamafactory/utils.py:1122正是显式传use_collectives=False的调用方,翻默认后会隐式带上broadcast_replicated_tensors=True——加了显式校验它当场就崩。相关的RFCs和API
对外接口变更(不兼容)
第一行是改名:旧名称在新设计下已不准确,读盘 rank 按字节负载挑选,不再固定为组内最小 rank;语义不变。
后两行是默认值改变,理由见上一节。
core/distributed_checkpoint/之外的调用方有两个:components/checkpoint/dcp_checkpointer.py(不传参,直接受益)与integration/llamafactory/utils.py(显式use_collectives=False,行为不变)。docs/guide/distributed_checkpoint.md已同步。请评审两点:(1) 是否需要保留旧参数名一个版本作为过渡(deprecation warning);(2) 默认值改变是否需要在 release note 中显著标注"所有 rank 必须一起调
load()"。新增对外接口
Platform基类新增三个方法,只在 torch 后端实现:broadcast_async(data, src, group, group_src)copy_each(dests, srcs)_foreach_copy_get_world_group()mindspore 后端不实现这三个方法,
new_group也保持它原有的NotImplementedError——该后端后续将被移除,不值得为这条路径适配。影响是 MS 上执行会真正广播的load()会在broadcast_async抛NotImplementedError(_shard_buffer里的obj.detach()同样走不通,MS 的 Tensor/Parameter 没有detach,见mindspore/platform.py中Platform.detach的注释)。tests/mindspore/st/distributed_checkpoint/目前只覆盖 reshard handler 与 offline convert,没有调用load()的用例,无 ST 受影响。顺带说明
filesystem_storage._broadcast_batch_bytes在 MS 上返回 0(即关闭小分片聚批)的原因:聚批把若干分片装进一块staging,两端都靠staging[off:off+n].reshape(shape)与staging共享存储——发送端写 view 必须落进 staging,接收端广播写 staging 后必须能从 view 读到。torch 对"一维连续张量的连续切片 + reshape"有明确的 view 语义,MS 没有等价契约。这是正确性护栏,不是性能取舍;在上述broadcast_async未实现之后,该分支实际已不可达。关联
完整的反馈期限
两周。若涉及接口命名的过渡策略需要更长讨论,可延长。
其他补充说明
本 RFC 对应的实现已完成并提交 PR !1322,UT/ST 验证结果见 PR 描述。提出 RFC 的目的是就设计选择与对外接口变更获取评审意见,特别是:
broadcast_from_minimum_rank改名是否需要过渡期;broadcast_batch_bytes覆盖,默认值硬编码)。附带一提:broadcast_replicated_tensors/broadcast_groups/broadcast_batch_bytes三个强耦合参数曾尝试收敛为单个broadcast: Union[bool, BroadcastOptions],实现后评估认为多态参数损害可发现性、且与本 PR 已有的改名叠加成两次破坏性变更,已回退,load()维持三个独立参数;—— 已自行否决,理由见"已处理"一节:该组合的降级结果正确,且现有调用方依赖它。broadcast_replicated_tensors与use_collectives的组合是否应显式报错