已开启
[RFC]: DCP 加载路径——广播下沉到读循环,读盘 rank 按字节均衡 #361
zhangbuxue创建于  18 天前
zhangbuxue成员
18 天前 创建

动机(Motivation)

DCP 加载是大模型作业启动与断点续训的主要耗时之一。master 上已经有"副本张量只读一次 + 广播"的能力(broadcast_from_minimum_rank),但当前形态存在四个结构性问题:

  1. 读与通信完全串行。 should_load_shard 在计划阶段把要广播的张量打上 BROADCAST_INFOexecute_read 把所有文件读完之后,在末尾调用一次 broadcast_loaded_tensors 统一发送。读的时候没有通信在走,发的时候没有读在走,两段时间无法互相掩盖。

  2. 读盘 rank 固定为 min(group_ranks) 同一组内的最小 rank 往往就是全局靠前的那几张卡,读盘压力反复压在少数卡上,其余卡空等。

  3. 广播粒度过细。 一个张量一次广播。实测 4 卡 Ascend 上,64 KiB 与 1 MiB 的广播代价几乎相同(约 145 µs),要到几 MB 才跑满带宽。30B MoE 这类模型每个文件约 9000 个张量,大量小分片使时间花在发起通信而非搬运数据上。

  4. 读盘本身有可观的常数开销。 逐分片重开 safetensors 文件(0.70 ms vs 复用的 0.045 ms);逐分片重建 available_keys 集合(约 7 ms × 1.8 万个分片)。

目标设计

把广播下沉到读循环内部,与读重叠;并把"谁读、谁收"的决策载体从 state_dict 的张量属性搬到加载计划的 ReadItem 上。

当前 master 本 RFC
广播时机 全部文件读完后统一发一轮 分片一就位即发,与后续读取重叠
决策载体 张量属性 BROADCAST_INFO ReadItem.source 上的 BroadcastSource
粒度 每个张量 每个分片(同一张量的两个分片可分给不同 rank)
读盘 rank min(group_ranks) 组内已分配字节最少者
并发 逐个同步等待 最多 8 个在飞
小分片 一片一次广播 小于 6 MiB 合批
文件句柄 每分片重开 LRU 保留(torch 8 个 / ms 1 个)
key 集合 每分片重建 每文件一次

决策载体的搬迁是前提:信息落在计划上而非张量对象上,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_confignpu: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 内)。原因有二:

  • checkpoint 文件本就按 key 字典序首尾相接排布,与全局序 (fqn, offset, index) 完全吻合(实测 9218/9218 匹配、零空洞),单文件内的读已经是顺序访问,没有可优化空间;
  • 拆成两个序之后,读线程按文件序碰页、主线程按全局序 copy_,两个访问序在同一文件上交错,打乱内核的顺序检测,预读失效。

方案 B:保留末尾统一广播,仅做负载均衡与聚批。
改动更小,但读与通信仍然串行,拿不到重叠收益,也无法解决大模型下通信段独占时间的问题。

方案 C(早期版本中曾实现,现已移除):读预取线程领先主线程若干批次。
该方案在两种存储上均实测无收益,故从实现中删除,读改为惰性生成器(取用顺序与时机不变)。数据见"性能影响分析"。结论要点:mmap 上的 disk→host 不是一段可以搬到旁路线程的等待——内核 readahead 已将其异步化;预取线程不是把这段时间从主线程移走,而是在总量上新增一遍全量遍历的 CPU 开销。

对现有代码的影响范围

文件 变更性质
core/distributed_checkpoint/planner.py 新增 BroadcastSourceLoadPlan.identity()ReadItem.broadcastable
core/distributed_checkpoint/standard_planner.py build_global_plan 重写,返回本 rank 的 LoadPlan(原为 list[LoadPlan]);移除 should_load_shardbuild_local_plan 标注目的缓冲是否在 host 上
core/distributed_checkpoint/filesystem_storage.py execute_read 重写;_load_* 族拆为 _fetch_*_apply_fetched;新增 _OpenFiles
core/distributed_checkpoint/util.py 移除 BROADCAST_INFO / broadcast_loaded_tensors 一套;新增 BroadcastBatcher 与建组、销毁、发送函数
core/distributed_checkpoint/metadata.py 新增 dtype_element_size
core/distributed_checkpoint/api.py load() 参数改名;计划选择逻辑随返回类型简化
platform/{platform,torch/platform}.py 新增三个接口;mindspore 后端未动,见下文

向后兼容性

存在对外接口不兼容变更(一处改名 + 两个默认值翻转),详见下节。checkpoint 文件格式、目录结构、metadata 均不变,已保存的 checkpoint 无需转换:没有共享 .metadata 的旧 checkpoint 会自动回退到 {rank}.metadata 并关掉 plan 交换。显式传 use_collectives=False / broadcast_replicated_tensors=False 时行为与当前一致。

性能影响分析

各优化项的微基准(均在 8× Ascend 910B3 上测得,已记录于对应代码注释):

数据
广播下界 64 KiB 与 1 MiB 均约 145 µs(4 卡);聚批阈值交点实测约 4.6 MiB,取 6 MiB
文件句柄复用 0.70 ms → 0.045 ms(400 张量的文件)
key 集合缓存 约 7 ms × 1.8 万个分片

端到端(8 卡 Ascend 910B3,本地 NVMe),10.45 GiB / 3147 张量的 Qwen3-30B-A3B 形状 checkpoint(由 master 侧写入):

broadcast=ON broadcast=OFF
master 29.5 s 3.45 s
本设计 3.0 s 3.37 s

放大到 48 层(56.89 GiB / 18867 张量),broadcast=ON:master 166.85 s → 本设计 16.86 s。按 HP_LOG_CONFIG=DCP:INFO 拆解,masterbroadcast_loaded_tensorsexecute_read 的 92%,其中 all_gather_object 8.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 22.26 s
本设计 无预读 22.65 s
预读=2(仅取 mmap 视图) 22.74 s
预读=2 + 预取线程内强制缺页 27.27 s
预读=2 + 预取线程内整块物化 25.02 s

共享存储(同机 NFS,裸盘 8 文件并行冷读 0.11 GiB/s,10.45 GiB / 3147 张量):

配置 冷加载
8 文件并行裸读 T_io(时间下界) 95.33 s
master 100.10 s
本设计 无预读 99.99 s
预读=2(仅取 mmap 视图) 100.14 s
预读=2 + 预取线程内强制缺页 100.05 s
预读=2 + 预取线程内整块物化 100.15 s

共享存储上加载时间贴在裸读下界上方 5%(100.0 vs 95.3 s),说明管道已被内核 readahead 喂满;五种配置极差 0.16 s。这同时给出一个对本设计有利的旁证:共享存储上瓶颈是存储带宽本身,唯一有效的杠杆是少读字节,正是"副本分片只读一次 + 广播"所做的事。

尚未覆盖:共享存储上 broadcast=ON/OFF 的直接对比(见下方后续项)。

已结项:万卡下的读序与文件句柄

masterexecute_readfile-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 读的正是它自己写的那个文件:

[rank=0] 共享读的文件: ['_rank0_.safetensors']
...
[rank=7] 共享读的文件: ['_rank7_.safetensors']

因此下述两个候选在主场景下都不成立

  • 方案 B(分窗 file-major)——否决。 前提是"一个目的分片要从多个文件取数",而这在实际负载中概率极低(同构续训恒为 1 个文件,等比例缩并行度也只有 2~4 个)。纯 Python 模拟对比三种访问模式后,LRU=256 w=1 在每一种下都优于 LRU=8 w=256,而分窗对"一片一文件、邻片跳文件"这一真实形态毫无作用。已实现后回退。
  • 方案 A(LRU 容量按 plan 定)——实现并实测后不采纳。 在构造的"轮转 + 回头访问"负载(64 文件 × 3000 张量、6400 分片)上确有 25.7× 收益(固定 8:30.44 s / 6400 次开文件;按 plan 定容量:1.17 s / 64 次)。但主场景下每个 rank 只开 1 个文件,容量下限本就是 8,该分支恒不触发;8 卡端到端 A/B 中位数 7.12/7.21 s(改前)对 7.17/7.06 s(改后),无可测差异。按"resharding 只保正确性、性能不作要求"的取舍,不为其增加复杂度。

实际采纳的是第三条,起因于一个此前未识别的问题:非张量(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 个叶子):

每 rank open() 次数 涉及文件数
逐分片读 80 3
按文件分组读 3 3

改动execute_read 把"无人共享的分片"按其所在文件归组,一个文件一次读。这类分片没有任何 rank 在等,因此不受"全局一致顺序"约束,可以任意重排;共享分片仍严格按各 rank 一致的顺序逐片读发,广播时序不变。

同时删掉读路径上两处冗余检查,它们说的都只是下一行自己会说的话:

被删的检查 原开销 谁来报
set(f.keys()) 成员预检查 每次开文件 6.5 ms(9402 张量的文件) safetensors:File does not contain tensor <name>
os.path.exists 存在性预检查 每文件一次 stat,共享存储上是一次网络往返 safetensors 与内建 open 均抛 FileNotFoundError 并带路径

后者连"路径是目录"都挡不住(exists() 对目录返回 True),新增用例钉住"缺文件由底层报错并指明文件名"。

仍未覆盖:换切分维度(Shard(0)Shard(1))的重切分下句柄抖动的实测。该场景按上述取舍只保正确性,resharding_api ST 已覆盖其正确性(含 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 原先默认 Falsesave() 默认 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=Trueuse_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

对外接口变更(不兼容)

load(..., broadcast_from_minimum_rank=False)   →   load(..., broadcast_replicated_tensors=False)
load(..., use_collectives=False)               →   load(..., use_collectives=True)
load(..., broadcast_replicated_tensors=False)  →   load(..., broadcast_replicated_tensors=True)

第一行是改名:旧名称在新设计下已不准确,读盘 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) 返回 work handle 而不等待,支撑多个广播在飞
copy_each(dests, srcs) 批量 pair-wise 拷贝;torch 侧走 _foreach_copy_
get_world_group() 取全世界通信组,避免为全副本参数重复建组

mindspore 后端不实现这三个方法new_group 也保持它原有的 NotImplementedError——该后端后续将被移除,不值得为这条路径适配。影响是 MS 上执行会真正广播的 load() 会在 broadcast_asyncNotImplementedError_shard_buffer 里的 obj.detach() 同样走不通,MS 的 Tensor/Parameter 没有 detach,见 mindspore/platform.pyPlatform.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 未实现之后,该分支实际已不可达。

关联

  • PR:!1322
  • 无依赖的其他 RFC。

完整的反馈期限

两周。若涉及接口命名的过渡策略需要更长讨论,可延长。

其他补充说明

本 RFC 对应的实现已完成并提交 PR !1322,UT/ST 验证结果见 PR 描述。提出 RFC 的目的是就设计选择对外接口变更获取评审意见,特别是:

  1. 负载均衡的确定性约束是否有更稳妥的表达方式;
  2. broadcast_from_minimum_rank 改名是否需要过渡期;
  3. 聚批阈值 6 MiB 是否应做成可配置项(当前可通过 broadcast_batch_bytes 覆盖,默认值硬编码)。附带一提:broadcast_replicated_tensors / broadcast_groups / broadcast_batch_bytes 三个强耦合参数曾尝试收敛为单个 broadcast: Union[bool, BroadcastOptions],实现后评估认为多态参数损害可发现性、且与本 PR 已有的改名叠加成两次破坏性变更,已回退,load() 维持三个独立参数;
  4. 非张量条目的去冗:它们在 save 侧参与负载均衡、在 load 侧不广播,于是被摊到多个文件、又被每个 rank 各读一遍。是否应当索性不对其做均衡(全部归到同一个 rank,让所有读者只开一个文件);
  5. broadcast_replicated_tensorsuse_collectives 的组合是否应显式报错 —— 已自行否决,理由见"已处理"一节:该组合的降级结果正确,且现有调用方依赖它。
likedislike
Zzhangbuxue成员
18 天前 添加了label:RFC
Zzhangbuxue成员
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
Zzhangbuxue成员
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
Zzhangbuxue成员
18 天前 关联了pull request:perf(dcp): overlap the replicated-shard broadcast with the read that feeds it
Zzhangbuxue成员
18 天前 修改标题为 “[RFC]: DCP 加载路径优化——读播流水化与逐分片开销削减”,原标题为“[RFC]: DCP 加载路径——广播下沉到读循环,读盘 rank 按字节均衡”
Zzhangbuxue成员
18 天前 修改标题为 “[RFC]: DCP 加载路径——广播下沉到读循环,读盘 rank 按字节均衡”,原标题为“[RFC]: DCP 加载路径优化——读播流水化与逐分片开销削减”
Zzhangbuxue成员
13 天前 修改了issue 的描述
Zzhangbuxue成员
12 天前 修改了issue 的描述
Zzhangbuxue成员
10 天前 修改了issue 的描述
Zzhangbuxue成员
9 天前 修改了issue 的描述
Zzhangbuxue成员
9 天前 修改了issue 的描述
Zzhangbuxue成员
9 天前 修改了issue 的描述