Skip to content

[Optimization] Incremental checkpoint save for dcp on torch 2.7.x (ARM CPU optimization)#1525

Merged
CyCle1024 merged 1 commit into
InternLM:mainfrom
tina-wen:dcp_save
Mar 26, 2026
Merged

[Optimization] Incremental checkpoint save for dcp on torch 2.7.x (ARM CPU optimization)#1525
CyCle1024 merged 1 commit into
InternLM:mainfrom
tina-wen:dcp_save

Conversation

@tina-wen

@tina-wen tina-wen commented Mar 3, 2026

Copy link
Copy Markdown
Contributor

Description

This PR optimizes dcp.save performance on ARM CPUs by implementing incremental metadata saving for torch 2.7.1.

Implementation

  • Incremental save: Only save metadata changes after first checkpoint
  • xtuner framework patch: Added patch_for_dcp_finish config flag
  • API update: Switch to storage_writer/planner for dcp.save

Performance

Checkpoint saving performance improved by up to 85%

Compatibility

✅ Works with existing ckpt_save
✅ No precision issues on recovery
✅ No PyTorch/PTA source changes

Comment thread xtuner/v1/patch/__init__.py
Comment thread xtuner/v1/engine/xtuner_storage.py Outdated
Comment thread xtuner/v1/engine/xtuner_storage.py Outdated

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR targets faster distributed checkpoint (DCP) saves on ARM CPUs (torch 2.7.1) by introducing incremental/cached planning and write-result handling, plus an optional monkeypatch to reduce finish-time overhead.

Changes:

  • Add a patch_for_dcp_finish config flag to optionally monkeypatch torch DCP internals.
  • Switch TrainEngine.save_dcp() to use storage_writer + planner on torch 2.7.x via new XtunnerWriter and XtunerCacheSavePlanner.
  • Introduce new engine utilities for caching save plans/metadata and write results.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 8 comments.

Show a summary per file
File Description
xtuner/v1/train/trainer.py Adds patch_for_dcp_finish config/plumbing to enable a DCP finish monkeypatch.
xtuner/v1/patch/torch_dcp_planner.py Adds a patched _save_state_dict implementation and a function to apply the monkeypatch.
xtuner/v1/patch/init.py Exposes the new patch function from the patch package.
xtuner/v1/engine/xtuner_storage.py New FileSystemWriter subclass that can cache write results to reduce repeated overhead.
xtuner/v1/engine/xtuner_cache_planner.py New DefaultSavePlanner subclass that caches global plan/metadata to support incremental saves.
xtuner/v1/engine/train_engine.py Uses the new writer/planner for torch 2.7.x DCP saves.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread xtuner/v1/train/trainer.py
Comment thread xtuner/v1/patch/torch_dcp_planner.py
Comment thread xtuner/v1/patch/torch_dcp_planner.py Outdated
Comment thread xtuner/v1/engine/xtuner_storage.py Outdated
Comment thread xtuner/v1/engine/xtuner_storage.py Outdated
Comment thread xtuner/v1/engine/train_engine.py Outdated
Comment thread xtuner/v1/engine/train_engine.py Outdated
Comment thread xtuner/v1/train/trainer.py
@tina-wen tina-wen force-pushed the dcp_save branch 2 times, most recently from 42b4328 to 26e7aa7 Compare March 10, 2026 07:59
@HAOCHENYE HAOCHENYE added the npu label Mar 11, 2026
@CyCle1024 CyCle1024 force-pushed the dcp_save branch 2 times, most recently from b31b2bc to bef8ad2 Compare March 19, 2026 08:36
@tina-wen tina-wen force-pushed the dcp_save branch 2 times, most recently from 367ff35 to b2101fa Compare March 25, 2026 13:27
@CyCle1024 CyCle1024 merged commit 9c3d549 into InternLM:main Mar 26, 2026
5 checks passed
@tina-wen tina-wen deleted the dcp_save branch May 14, 2026 07:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants