[Cherry-Pick][Performance] Add num_workers to paddle.load for parallel payload reading (#79786) - #79787
Conversation
…l payload reading (PaddlePaddle#79786) Cherry-pick of PaddlePaddle#79786. paddle.load moves tensor bytes with a single thread, which caps the load bandwidth well below what the storage can deliver: an 8GB checkpoint still takes 6.7s from tmpfs, where there is no disk involved at all. Intercept the unpickler's readinto() for out-of-frame payloads, record and skip them so the first pass only parses metadata, then read the payloads in parallel with os.preadv() straight into the final buffers. The file format is unchanged and the fast path is opt-in via paddle.load(..., num_workers=n), falling back to the serial implementation in every unsupported case. dcp.load_state_dict gains a matching num_workers argument (default 1). This branch still reads checkpoint files directly in local_load_state_dict and load_state_dict_impl, so the argument is threaded through those call sites instead of the _load_checkpoint_data_file helper introduced on develop.
Paddle-Bot Review Board （review完成)
Powered by Nyanpasu claude with Opus 4.8 默认推理级别, please check the suggestions carefully. |
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
| # ``b`` views the buffer the numpy array will use; keeping it here | ||
| # also keeps it alive until the fill pass. | ||
| self.holes.append((self._f.tell(), b)) |
There was a problem hiding this comment.
readinto(b) 收到的 b 是 CPython 用 PyMemoryView_FromMemory 构造的裸指针视图，它并不持有被写入 payload 的那个对象，所以注释里 “keeping it here also keeps it alive until the fill pass” 不成立。如果某个 >= 1MB 的 payload 没有被反序列化结果引用（pickle 流中该 bytes 对象后面被丢弃），扫描阶段结束、unpickler 释放后这块内存就会被回收，_fill_holes 变成写入已释放内存。
用本模块的真实代码可以复现：构造一个含 2MB BINBYTES8 payload、但反序列化结果不引用它的文件，num_workers=4 时同尺寸的新分配会复用该内存并被覆盖，随后进程以 malloc(): invalid next size (unsorted) 中止；num_workers=1 的串行路径正常返回。
建议在填充前确认每个 hole 的内存仍被结果图中的存活 buffer 覆盖（否则回退串行），或至少修正注释并明确 “被 punch 的 payload 必须被结果引用” 这一前提。
| 'keep_name_table', | ||
| 'return_numpy', | ||
| 'safetensors', | ||
| 'num_workers', |
There was a problem hiding this comment.
num_workers 已加入 supported_configs 并在 _parse_load_config 中透传，但没有同步加到 python/paddle/framework/io.py:74 的 _LoadOptions；而 load() 的签名是 **configs: Unpack[_LoadOptions]，类型检查会把 paddle.load(path, num_workers=8) 报成非法参数。建议补上 num_workers: NotRequired[int]（如需保留 None 语义可写成 int | None）。
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## release/3.4 #79787 +/- ##
==============================================
Coverage ? 99.13%
==============================================
Files ? 7
Lines ? 346
Branches ? 0
==============================================
Hits ? 343
Misses ? 3
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Cherry-pick of the review fixes on PaddlePaddle#79786. The memoryview handed to readinto does not own the payload buffer, so filling it after the parse is only safe when the parsed object still references that memory. It does not hold for payloads the unpickler copies or decodes while parsing, e.g. a bytearray at protocol 4, which is built from a temporary bytes object that is freed right away. Writing into that freed memory corrupted data in 3 of 6 runs. Read each payload in parallel from inside readinto instead, so the buffer is complete before the unpickler continues and no lifetime assumption is left. This also drops the second full parse for files with no large payload. Validate num_workers in _parse_load_config so invalid values are rejected on macOS too, where paddle.load takes the _pickle_loads_mac path. Skip the tests that need os.preadv on Windows and the ones that need the fast path on macOS. Add regression tests for bytearray, bytes and str payloads.
|
感谢 review，两条都已修复（commit 25ee96c，release/3.4 对应 PR #79787 同步修复）。 P1 确认成立，已改为在 根因与你的判断一致： 修复方式不是加白名单判断 payload 类型（无法在 代价是失去了跨 payload 的重叠，单文件收益从 ~5x 降到 ~3.3x（2.10GB 新增回归用例 P3 随之自动解决。 现在没有任何字节被跳过，第一趟解析本身就是完整的， 另外这一版顺带修了 CI：
|
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
| chunks = [ | ||
| (offset + start, mv[start : start + _READ_CHUNK_SIZE]) | ||
| for start in range(0, n, _READ_CHUNK_SIZE) | ||
| ] | ||
| # Iterating the map result waits for every chunk and re-raises the first | ||
| # exception, so the buffer is complete once readinto returns. | ||
| for _ in self._pool.map( | ||
| lambda chunk: _pread_exact(self._fd, *chunk, self._path), chunks | ||
| ): | ||
| pass |
There was a problem hiding this comment.
readinto 现在是同步调用：unpickler 必须等当前 payload 读完才能继续，因此并行度只来自单个 payload 内部的 32MB 分块，跨 payload 的并行被取消了。当 payload ≤ 32MB 时 chunks 只有 1 个（2MB/4MB/16MB 这类常见张量都属此类），num_workers 实际不起作用，整个文件退化为逐个 payload 串行 preadv。
用带计数（每次调用注入 4ms 延迟以便观测重叠）的 os.preadv 包装实测：60 个互不相同的 2MB payload、num_workers=8 时，本实现最大并发为 1，上一版 ad5cf79e 的实现为 8。
PR 描述与 _READ_CHUNK_SIZE 的注释仍按被替换掉的 _HolePunchFile/_fill_holes 跨 payload 方案表述，2.8x~5.8x 的加速数据也来自旧实现，建议一并更新。若希望常见尺寸的张量也能用满线程池，可以按 payload 自适应分块（例如 chunk ≈ max(1MB, ceil(n / num_workers))）后重新测量。
`load_state_dict` decides whether every rank can restore straight from its own `.distcp` before any resharding happens. Answering that question costs a full `paddle.load` of the `*.metadata` file plus a `MetadataManager` build, because `check_resumable_locally` needs `storage_metadata`. On a fully replicated checkpoint `storage_metadata` holds `num_keys * world_size` entries -- hundreds of MiB and tens of seconds -- all to produce one boolean. This adds a fast path that answers the same question from two much cheaper sources: - `metadata_reader.load_state_dict_metadata` parses only the `state_dict_metadata` field of the metadata file. It relies on that field being pickled immediately before `storage_metadata`: planting a `STOP` over the opcode that would push the `"storage_metadata"` key ends the stream with exactly the wanted value on the stack, so nothing past the cut is read. Any stream that does not match that expectation falls back to a full `paddle.load`, so the function is always safe to call. It is also exported as `paddle.distributed.load_state_dict_metadata`. - `distcp_reader.scan_tensor_shapes` reads the keys and shapes a `.distcp` really holds by walking its pickle header and seeking over the weight payloads, which costs a few hundred bytes per tensor instead of the file. `fast_resumable.check_resumable_locally_fast` combines the two. When every tensor of the checkpoint is stored whole, a tensor key identifies its `LocalTensorIndex` uniquely and `storage_metadata` carries no information the key alone does not. When any tensor is sharded, shard identity depends on `global_offset` / `flattened_range`, which live nowhere but `storage_metadata`, so the fast check declines (returns None) and the stock check runs unchanged. Two properties keep it a drop-in replacement: the applicability verdict is a pure function of the shared metadata file, so all ranks decide identically before communicating, and exactly one `all_gather_object` is issued when the check applies and none when it declines -- the same collective count as the stock check on either branch. The check is never more permissive than the stock one: it validates the file's real contents, so a truncated `.distcp` (which the stock check accepts, since it only calls `os.path.isfile`) makes this one fall back to resharding. The behaviour is controlled by the new `load_state_dict(..., fast_resumable_check=True)` argument, so the original code path stays reachable. test/flex_checkpoint/test_fast_resumable_check.py covers the three new modules, mostly differentially against `utils.check_resumable_locally`: pickle protocols 2/4/5, all supported dtypes, scalar and empty shapes, the 64 KiB frame boundary, dedup'ed checkpoints, `ShardedWeight` inputs, uninitialized tensors, truncated and corrupt files, the byte-level framing helpers, and the collective branch.
64fc94d
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
本轮针对新提交 64fc94d（跳过 local-resume 检查中的全量 metadata 加载）进行复核：metadata_reader 部分解析、distcp_reader 头部扫描与 fast_resumable 三态检查的集合通信次数、shape 比较与回退语义均与原 check_resumable_locally 一致，且更严格（截断/损坏文件回退 reshard），未发现新的问题。
以下两条既有建议仍未处理，保持关注：
paddle.load的num_workers尚未加入_LoadOptions，Unpack[_LoadOptions]下类型检查仍会报参数非法（P3）。readinto改为同步读取后，跨 payload 并行被取消，≤32MB 的常见张量下num_workers实际不生效，相关加速数据与注释建议同步更新（P2）。
详情见对应评论线程。
|
@risemeup1111 review一下 |
|
@DanielSun11 已对当前 head 新提交「跳过 local-resume 检查中的全量 metadata 加载」已复核： 以下两条既有建议仍待处理，详见对应评论线程：
具体条目与状态见 Review Board。 |
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
本轮针对新提交 426b4d1（修复 metadata API 示例的类型检查）进行复核：该提交仅移除 metadata_reader.py 中 load_state_dict_metadata docstring 内一处会导致 doctest 类型检查失败的示例输出（state_dict_metadata["w1"][0].local_shape），不涉及任何运行时逻辑，未引入新的问题。
以下两条既有建议仍未处理，保持关注：
paddle.load的num_workers尚未加入_LoadOptions，Unpack[_LoadOptions]下类型检查仍会将paddle.load(path, num_workers=8)报为非法参数（P3）。readinto改为同步读取后，跨 payload 并行被取消，≤32MB 的常见张量下num_workers实际不生效，相关加速数据与注释建议同步更新（P2）。
详情见对应评论线程。
This branch has not been deployed
6 participants
PR Category
Performance Optimization
PR Types
Performance
Description
devPR: #79786
本 PR 含两处改动，针对
dist.load_state_dict加载路径上两个不同的瓶颈：一、
paddle.load(num_workers=...)：并行读取张量数据（cherry-pick of #79786）paddle.load全程单线程搬运张量数据。把 8GB checkpoint 放到/dev/shm后仍需 6.7s，说明瓶颈不是 IO 而是单线程拷贝：pickle 对张量数据不做任何变换，元信息只占极小比例，耗时几乎全在“把 GB 级字节搬一遍”。protocol >= 4下 >= 64KB 的 payload 写在 pickle frame 之外，可以先只读元信息建好完整对象图（此时 numpy 数组已是空缓冲区上的零拷贝视图），再用线程池的os.preadv把这些 payload 直接读进最终内存，preadv期间释放 GIL 所以能真正并行。pickle 格式与解析逻辑均不变，安全反序列化照常生效。默认
num_workers=1即原串行行为。BytesIO、macOS、无os.preadv、protocol == 2、文件中没有大 payload，以及并行过程中的任何异常，都自动回退串行。二、跳过 local-resume 检查中的全量 metadata 加载
dist.load_state_dict在 reshard 前先判断“每张卡能否直接从自己的.distcp恢复”。这个布尔值目前的代价是对*.metadata做完整paddle.load再构建MetadataManager，因为check_resumable_locally需要storage_metadata——而它是Metadata里最大的字段，完全副本下有num_keys * world_size条记录，真实 154GB checkpoint 上是 227MiB 文件、十几秒 unpickle，全部只为换一个 bool。改为由两个只读极少字节的读取器回答同一个问题：
metadata_reader.load_state_dict_metadata：state_dict_metadata恰好紧邻storage_metadata之前被 pickle，在后者的 key opcode 处植入STOP，只读文件前缀即可拿到需要的字段，其后的字节一个都不从磁盘读。不符合该假设时退回完整paddle.load，因此调用永远安全。该函数同时导出为paddle.distributed.load_state_dict_metadata。distcp_reader.scan_tensor_shapes：只走 pickle 头部得到本卡文件真实的{key: shape}；payload 长度显式写在 opcode 里，故可直接seek跳过。fast_resumable.check_resumable_locally_fast组合两者，是三态的：只有当 checkpoint 里每个张量都整块存放时，张量名才唯一确定其LocalTensorIndex，storage_metadata不携带额外信息，结论可只由本卡文件头得出；一旦有张量被切分，shard 身份依赖仅存在于storage_metadata的global_offset/flattened_range，此时返回None，调用方原样走原检查。两点使它成为 drop-in 替换：是否适用只取决于共享的 metadata 文件，所有卡在任何通信之前就得到相同判断；适用时恰好发一次
all_gather_object、不适用时零次，与原实现的集合通信次数一致。它也不比原检查宽松——校验文件的真实内容而非 metadata 的声明，因此没写完或被截断的.distcp（原检查只做os.path.isfile）会退回 reshard。由
load_state_dict(..., fast_resumable_check: bool = True)控制，置False强制走原路径。是否引起精度变化
否。改动一的并行路径读的是同一文件的同一偏移，与串行结果逐位一致；改动二只替换
check_resumable_locally的布尔判断，张量数据通路完全未动。