[FlexCheckpoint] Avoid the full metadata load in the local-resume check - #79807
Open
DanielSun11 wants to merge 2 commits into
Open
DanielSun11 wants to merge 2 commits into
DanielSun11 wants to merge 2 commits into
Conversation
`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.
risemeup1111
approved these changes
Sep 22, 2026
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## develop #79807 +/- ##
==========================================
Coverage ? 99.63%
==========================================
Files ? 5
Lines ? 274
Branches ? 0
==========================================
Hits ? 273
Misses ? 1
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
LiYuRio
previously approved these changes
Sep 23, 2026
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
risemeup1111
approved these changes
Sep 24, 2026
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR Category
Distributed Strategy
PR Types
Performance
Description
dist.load_state_dict在 reshard 之前先判断“每张卡能否直接从自己的.distcp恢复”，只有答案为否才需要 reshard。这个布尔值目前的代价是对*.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强制走原路径，无需环境变量即可按调用点回退。release/3.4 对应 PR 为 #79787。
是否引起精度变化
否。快速路径只替换
check_resumable_locally的布尔判断，张量数据通路完全未动。