[Performance] Add num_workers to paddle.load for parallel payload reading - #79786
DanielSun11 wants to merge 2 commits into
Conversation
…ding 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).
Paddle-Bot Review Board （review完成)
|
risemeup1111
left a comment
risemeup1111
left a comment
There was a problem hiding this comment.
| def readinto(self, b) -> int: | ||
| n = len(b) | ||
| if n >= self._thresh: | ||
| # ``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)) | ||
| self._f.seek(n, 1) | ||
| return n | ||
| return self._f.readinto(b) |
There was a problem hiding this comment.
这里假设 readinto 收到的 b 就是 payload 的最终存储、解析结束后交给 _fill_holes 写入即可，但 unpickler 传入的是 PyMemoryView_FromMemory 创建的 memoryview（obj is None），它不持有底层缓冲区，只有“解析结果仍然引用这块内存”时缓冲区才有效。对在 unpickle 过程中被拷贝的 payload 这个前提不成立：例如 ≥1MB 的 bytearray（pickle 中以 bytearray(BINBYTES) 保存），解析时 bytearray(...) 已经从仍为空白的缓冲区拷贝过一次，承载它的临时 bytes 随解析结束被释放，_fill_holes 随后便把文件字节写入已释放/已被复用的内存。
实测（pickle.dump({"a": bytearray(b"a"*3_000_000), "b": bytearray(b"b"*3_000_000)}, f, protocol=4) 后调用 parallel_safe_load_pickle(path, f, 8)）连续 6 次运行：
- 5 次：
b的 3,000,000 字节中混入 112 字节来自a的 payload、10 字节为 0，其余才正确； - 1 次：
a只有前 8,144 字节正确，其余为 0； - 6 次都没有 warning、也没有触发串行回退，即静默返回错误数据。
同一输入结果随运行变化，说明这里依赖的是分配器巧合（写入已释放内存）。建议扫描完成后先校验每个 hole 缓冲区仍是返回值的存储（例如收集结果中所有 np.ndarray 的 data 区间、要求 hole 落在其中，否则回退串行），并补一个 ≥1MB bytearray payload 的用例覆盖该分支。
There was a problem hiding this comment.
已在新 head 25ee96c5 验证修复：改为在 readinto 内同步完成并行读取后，不再对缓冲区生命周期做任何假设。用 bytearray/bytes/str/大端 numpy 数组/嵌套结构等 12 种 payload 场景各重复多轮与串行逐字节比对，结果全部一致且无 warning，原场景不再出现坏字节。
| if not holes: | ||
| # protocol < 3, or tiny tensors only: nothing to parallelize. | ||
| f.seek(0) | ||
| return safe_load_pickle(f, encoding=encoding) |
There was a problem hiding this comment.
holes 为空时第一趟已经完整读过整个流（没有跳过任何字节），result 本身可用；这里却丢弃它并用 f 重新解析，等于对同一文件做两遍完整解析。实测 21MB、protocol 2 的文件在 num_workers=8 下耗时 95ms，而串行只需 50ms（约 1.9x），这类收益为 0 的文件要付双倍解析成本。若无其他考虑可直接 return result；若确实需要重新解析，建议在 docstring 中说明这份额外开销。
There was a problem hiding this comment.
同一提交改为单趟解析、去掉了无 payload 时的重复解析分支，该开销已消失：同样 21MB / protocol 2 的文件在本机实测串行 53.3ms、num_workers=8 53.5ms（比值 1.00，原先约 1.9x）。此条已解决。
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## develop #79786 +/- ##
==========================================
Coverage ? 97.75%
==========================================
Files ? 3
Lines ? 89
Branches ? 0
==========================================
Hits ? 87
Misses ? 2
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Review feedback (P1): 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 (P3): nothing is skipped, so the first parse is always usable. 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.
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.
PR 描述「接口」一节把 5 种情况统一写成"自动回退到原串行实现，只 warn 不报错"，但实现中只有并行读取抛异常这一条路径会告警（python/paddle/framework/parallel_pickle_load.py:156-161），其余都是静默走串行：
BytesIO/ 无os.preadv/ macOS：_fast_path_available直接返回 False（python/paddle/framework/parallel_pickle_load.py:121-130），不打印任何信息；macOS 在python/paddle/framework/io.py:1300就已分流到_pickle_loads_mac。- 文件中没有 >= 1MB 的 payload：仍会经过一次
_ParallelPayloadFile解析，只是所有readinto都转发给底层文件对象（python/paddle/framework/parallel_pickle_load.py:100-103），同样没有告警。
paddle.load 的 docstring（python/paddle/framework/io.py:1132-1135）写的是 "falls back to the serial implementation otherwise"，与实现一致，因此只需把描述改成"静默回退，仅并行读取出错时才 warn"，避免使用者以为传了 num_workers 一定会收到提示。
3 participants
PR Category
Performance Optimization
PR Types
Performance
Description
背景
paddle.load加载 checkpoint 时全程单线程搬运张量数据，带宽远低于存储与主机的能力。把 8GB checkpoint 放到/dev/shm（完全排除磁盘）后paddle.load仍需 6.7s（1.2 GB/s），说明瓶颈不在 IO 而在单线程拷贝：pickle 对张量数据不做任何变换，protocol >= 3存下来的字节与内存逐位相同，元信息只占极小比例（实测一个 2.1GB 的.distcp文件中 opcode + 名字 + dtype/shape 合计仅 1,106 字节，占 0.00005%），因此耗时几乎全在“把 GB 级字节搬一遍”。方案
保持 pickle 格式完全不变，只改变“谁来搬”：
_ParallelPayloadFile包装 checkpoint 文件。CPython 的 C 版 unpickler 对写在 pickle frame 之外的 payload 会调用文件对象的readinto()（protocol >= 4的 framer 对 >= 64KB 的数据就写在 frame 之外）；readinto(b)收到 >= 1MB 的请求时，把memoryview(b)按 32MB 切块交给线程池的os.preadv并行读，等所有块写完才返回，随后把底层文件指针前进len(b)；小于阈值的请求原样转发给串行readinto；os.preadv是系统调用，执行期间释放 GIL，所以线程能真正并行；切块互不重叠，每个字节只被一个线程写一次，无需加锁。切块的原因是张量数量少时（2.1GB 的.distcp只有 3 个 payload）“一个 payload 一个线程”用不满线程池。关键点是数据在
readinto返回前已经就位：unpickler 观察到的行为与普通readinto完全一致，代码对这块内存的生命周期不做任何假设，因此bytearray/bytes/str这类在解析期就会拷贝或解码的 payload 同样安全。（初版曾用“先跳过、解析完再回填”的写法，对 protocol 4 的bytearray会写进已释放的临时缓冲区，属于越界写；已按 review 意见重写，并补了对应单测。）pickle 的解析工作一步没少，安全反序列化（
RestrictedUnpickler）照常生效，下游_pack_loaded_dict/StructuredToParameterName@@/_parse_load_result等后处理无需改动。接口
默认关闭，行为与改动前完全一致：
以下情况自动回退到原串行实现，只 warn 不报错：传入
BytesIO、macOS（走_pickle_loads_mac）、无os.preadv（Windows）、文件中没有 >= 1MB 的 payload（含protocol == 2，payload 以 latin1 文本存储），以及并行读取过程中的任何异常（整体丢弃结果、seek(0)后重新串行加载，不会返回半填充的张量）。num_workers的校验放在_parse_load_config，因此非法参数在所有平台上都会报错。实测数据
本地 NVMe（ext4），加载真实的 flex_checkpoint
.distcp文件：dist.load_state_dict（256MB checkpoint）：num_workers=10.473s ->num_workers=80.214s（2.2x）收益上限取决于存储带宽：在冷读带宽约 1.8 GB/s 的网络盘上只有约 1.4x（IO 本身已经是下界），本地盘或 page cache 命中时可达 3~6x。
单测
新增
test/legacy_test/test_paddle_load_num_workers.py，29 个 case：Layer.state_dict()（含参数 name 恢复）、tensor place 与串行一致、不同num_workers、重复加载稳定性。全部使用逐字节比对而非allclose；bytearray/bytes/str在 protocol 4/5 下重复加载校验，覆盖上面提到的缓冲区生命周期问题；TypeError/ValueError，None/0/1走串行；BytesIO、缺少os.preadv、sys.platform == 'darwin'、注入preadv短读；_BUFFER_SIZE = 8192的取值）、小 payload 不进快路径、短读抛EOFError。Linux 29 个全部通过；Windows（无
os.preadv）与 macOS（paddle.load走_pickle_loads_mac）下依赖快路径的用例用skipUnless跳过，其余用例仍校验回退结果与串行逐字节一致。是否引起精度变化
否。
并行路径与串行路径读取同一文件的同一偏移，结果逐位一致：2.10GB 文件 10/10 个张量、7.83GB 文件 161/161 个张量的
dtype/shape相同且tobytes()完全相等；单测中对 NaN/±Inf/-0.0 也做了逐字节校验。并行读取过程中的任何异常都会整体丢弃结果、退回串行加载，不会返回半填充的张量。