Persist async checkpoint plans across resume - #353
Conversation
Greptile SummaryThis PR fixes a metadata persistence bug in the async checkpoint planner: local save plans were previously stored only as a dynamic
Confidence Score: 3/5Not safe to merge: the regression test added by the PR is expected to fail with the chosen planner, and the underlying persistence fix is only wired into the decentralized code path, leaving the default reduce_scatter path unchanged. The core idea is sound, but the implementation has two gaps that need resolving before this can land. First, _set_all_local_plans is called only inside the can_run_decentralized_global_plan branch; the reduce_scatter path's global_step function still creates global_metadata without persisting the plans into planner_data, so the fix does nothing for any job that uses a standard planner. Second, the new regression test uses DefaultSavePlanner, which lacks can_run_decentralized_global_plan, so execution always takes the reduce_scatter path, planner_data remains None after save, and assert loaded_all_plans is not None will fail on every run of the test. Both changed files need attention: state_dict_saver.py needs _set_all_local_plans called (or an equivalent) inside global_step for the reduce_scatter path, and test_async_writer.py needs a planner that actually exercises the decentralized path (or the assertion scope needs to be aligned with what DefaultSavePlanner can produce). Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[save_state_dict_async_plan] --> B{validated_cache_reuse?}
B -- yes --> C[Skip planning, use cached_central_plan]
B -- no --> D{can_run_decentralized_global_plan?}
D -- yes --> E[local_step]
E --> F[verify_global_md_reuse]
F --> G{loaded_all_plans valid?}
G -- no --> H[gather_object all ranks]
H --> I[planner.create_global_plan]
I --> J["_set_all_local_plans(global_metadata, all_local_plans)"]
J --> K["metadata.all_local_plans = plans (legacy attr)"]
J --> L["metadata.planner_data[__nvrx_all_local_plans] = plans (persistent)"]
G -- yes --> M[global_metadata = None, reuse cached]
D -- no --> N["reduce_scatter(local_step, global_step)"]
N --> O["global_step: planner.create_global_plan (NO _set_all_local_plans)"]
style O fill:#ffaaaa
subgraph Reload
P["FileSystemReader.read_metadata()"] --> Q["_get_all_local_plans(metadata)"]
Q --> R{metadata.all_local_plans attr?}
R -- present --> S[Return in-memory attr]
R -- absent --> T{"planner_data[key]?"}
T -- present --> U[Return persisted plans]
T -- absent --> V[Return None]
end
L --> P
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
A[save_state_dict_async_plan] --> B{validated_cache_reuse?}
B -- yes --> C[Skip planning, use cached_central_plan]
B -- no --> D{can_run_decentralized_global_plan?}
D -- yes --> E[local_step]
E --> F[verify_global_md_reuse]
F --> G{loaded_all_plans valid?}
G -- no --> H[gather_object all ranks]
H --> I[planner.create_global_plan]
I --> J["_set_all_local_plans(global_metadata, all_local_plans)"]
J --> K["metadata.all_local_plans = plans (legacy attr)"]
J --> L["metadata.planner_data[__nvrx_all_local_plans] = plans (persistent)"]
G -- yes --> M[global_metadata = None, reuse cached]
D -- no --> N["reduce_scatter(local_step, global_step)"]
N --> O["global_step: planner.create_global_plan (NO _set_all_local_plans)"]
style O fill:#ffaaaa
subgraph Reload
P["FileSystemReader.read_metadata()"] --> Q["_get_all_local_plans(metadata)"]
Q --> R{metadata.all_local_plans attr?}
R -- present --> S[Return in-memory attr]
R -- absent --> T{"planner_data[key]?"}
T -- present --> U[Return persisted plans]
T -- absent --> V[Return None]
end
L --> P
Reviews (2): Last reviewed commit: "Persist async checkpoint plans across re..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| :caption: Contents: | ||
|
|
||
| usage_guide | ||
| metadata_caching |
There was a problem hiding this comment.
metadata_caching is added to the toctree but no corresponding metadata_caching.rst file exists in docs/source/checkpointing/async/. Sphinx will raise an error like WARNING: toctree contains reference to nonexistent document 'checkpointing/async/metadata_caching' and fail strict builds. The referenced file needs to be created before this can land.
| def test_cached_metadata_plans_survive_reload(self, tmp_path_dist_ckpt, async_queue): | ||
| """Cached local plans are restored from metadata after a checkpoint reload.""" | ||
| Utils.initialize_distributed() | ||
| model = FSDP(Model((1024, 1024), 8)) | ||
| state_dict = model.state_dict() | ||
| planner = DefaultSavePlanner() | ||
| metadata_cache = CheckpointMetadataCache() | ||
|
|
||
| with TempNamedDir(tmp_path_dist_ckpt / 'cached_metadata_reload', sync=True) as ckpt_dir: | ||
| self.async_save_checkpoint( | ||
| ckpt_dir, | ||
| state_dict, | ||
| planner, | ||
| async_queue, | ||
| caching=True, | ||
| metadata_cache=metadata_cache, | ||
| ) | ||
| async_queue.maybe_finalize_async_calls(blocking=True, no_dist=False) | ||
|
|
||
| loaded_metadata = FileSystemReader(ckpt_dir).read_metadata() | ||
| resumed_metadata_cache = CheckpointMetadataCache() | ||
| resumed_metadata_cache.set_cached_global_metadata(loaded_metadata) | ||
| _, _, _, loaded_all_plans = resumed_metadata_cache.get_cache_metadata() | ||
|
|
||
| assert loaded_all_plans is not None | ||
| assert len(loaded_all_plans) == torch.distributed.get_world_size() | ||
|
|
||
| def test_cached_data_structure(self, tmp_path_dist_ckpt): | ||
| """ |
There was a problem hiding this comment.
DefaultSavePlanner
_set_all_local_plans is only called inside the elif branch in save_state_dict_async_plan (lines 361–366 of state_dict_saver.py), which requires both the planner and writer to expose can_run_decentralized_global_plan = True. FileSystemWriterAsync satisfies that condition, but the standard torch.distributed.checkpoint.DefaultSavePlanner does not define this attribute, so getattr(planner, 'can_run_decentralized_global_plan', False) returns False. Execution falls through to the else/reduce_scatter path, _set_all_local_plans is never invoked, planner_data in the saved .metadata file stays None, and the assertion assert loaded_all_plans is not None will always fail. The test should either use a planner that supports the decentralized path or the fix should also be applied inside global_step.
| def _set_all_local_plans(metadata: Metadata, all_local_plans: List[SavePlan]) -> None: | ||
| """Store local plans in serialized metadata and preserve the legacy in-memory attr.""" | ||
| metadata.all_local_plans = all_local_plans | ||
| if metadata.planner_data is None: | ||
| metadata.planner_data = {} | ||
| metadata.planner_data[_ALL_LOCAL_PLANS_KEY] = all_local_plans |
There was a problem hiding this comment.
planner_data will raise TypeError
_set_all_local_plans initialises planner_data to {} only when it is None. If an external planner has already populated planner_data with a non-dict type (a dataclass, a list, etc.) the subscript assignment metadata.planner_data[_ALL_LOCAL_PLANS_KEY] = all_local_plans will raise TypeError. Similarly, _get_all_local_plans relies on planner_data.get(...) which would raise AttributeError for any non-dict falsy value that is not None. Both helpers should guard with isinstance(metadata.planner_data, dict) rather than an is None check.
Store cached local save plans in serialized DCP planner metadata so the first save after restart can validate and reuse loaded global metadata. Previously, NVRx attached the gathered local plans only as a dynamic `Metadata.all_local_plans` attribute. Newer PyTorch DCP writers may rebuild the `Metadata` dataclass before pickling `.metadata`, which drops dynamic attributes and prevents resumed jobs from finding `loaded_all_plans`. Persist the plans under a reserved `Metadata.planner_data` key while preserving the legacy in-memory attribute for compatibility. Add a regression test that loads metadata back from disk and verifies a fresh metadata cache can recover the saved local plans.
7ec2144 to
4f8bbe2
Compare
hexinw-nvidia
left a comment
hexinw-nvidia
left a comment
There was a problem hiding this comment.
Have an offline discussion with @sbak5 to understand the change.
The change looks good overall.
I will leave it to @shurkat-nvidia for the final approval. Thanks.
|
Added "ci-approved" label to trigger the GitLab CI. |
3 participants
Store cached local save plans in serialized DCP planner metadata so the first save after restart can validate and reuse loaded global metadata.
Previously, NVRx attached the gathered local plans only as a dynamic
Metadata.all_local_plansattribute. Newer PyTorch DCP writers may rebuild theMetadatadataclass before pickling.metadata, which drops dynamic attributes and prevents resumed jobs from findingloaded_all_plans.Persist the plans under a reserved
Metadata.planner_datakey while preserving the legacy in-memory attribute for compatibility. Add a regression test that loads metadata back from disk and verifies a fresh metadata cache can recover the saved local plans.