Skip to content

Commit 9c3d549

Browse files
tina-wenwentiange
tina-wen
and
wentiange
authored
[Optimization] Incremental checkpoint save for dcp on torch 2.7.x (ARM CPU optimization) (#1525)
[Feature] incremental save for dcp.save Co-authored-by: wentiange <tiangewen@qq.com>
1 parent 401ab4f commit 9c3d549

5 files changed

Lines changed: 377 additions & 3 deletions

File tree

‎xtuner/v1/patch/__init__.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,10 @@
11
from . import torch_shape_env_simplify_pt28
2-
from .torch_dcp_planner import patch_default_save_plan
2+
from .torch_dcp_planner import patch_dcp_save_state_dict, patch_dcp_save_with_cache_storage, patch_default_save_plan
33

44

5-
__all__ = ["patch_default_save_plan", "torch_shape_env_simplify_pt28"]
5+
__all__ = [
6+
"patch_default_save_plan",
7+
"torch_shape_env_simplify_pt28",
8+
"patch_dcp_save_state_dict",
9+
"patch_dcp_save_with_cache_storage",
10+
]

‎xtuner/v1/patch/torch_dcp_planner.py‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,19 @@
1+
import inspect
2+
import warnings
3+
from pathlib import Path
4+
from typing import Optional
5+
6+
import torch
7+
import torch.distributed as dist
8+
import torch.distributed.checkpoint as dcp
19
import torch.distributed.checkpoint.default_planner as torch_default_runner
10+
from torch.distributed.checkpoint import state_dict_saver
11+
from torch.distributed.checkpoint.default_planner import DefaultSavePlanner
12+
from torch.distributed.checkpoint.logger import _dcp_method_logger
13+
from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE, Metadata
14+
from torch.distributed.checkpoint.planner import SavePlan, SavePlanner
15+
from torch.distributed.checkpoint.storage import StorageWriter
16+
from torch.distributed.checkpoint.utils import _DistWrapper
217

318

419
def fake_validate_global_plan(*args, **kwargs):
@@ -7,3 +22,111 @@ def fake_validate_global_plan(*args, **kwargs):
722

823
def patch_default_save_plan():
924
torch_default_runner._validate_global_plan = fake_validate_global_plan
25+
26+
27+
def _xtuner_save_state_dict(
28+
state_dict: STATE_DICT_TYPE,
29+
storage_writer: StorageWriter,
30+
process_group: Optional[dist.ProcessGroup] = None,
31+
coordinator_rank: int = 0,
32+
no_dist: bool = False,
33+
planner: Optional[SavePlanner] = None,
34+
) -> Metadata:
35+
torch._C._log_api_usage_once("torch.distributed.checkpoint.save_state_dict")
36+
37+
distW = _DistWrapper(process_group, not no_dist, coordinator_rank)
38+
if planner is None:
39+
planner = DefaultSavePlanner()
40+
assert planner is not None
41+
42+
global_metadata = None
43+
44+
ckpt_kwargs = {}
45+
if (ckpt_id := getattr(storage_writer, "checkpoint_id", None)) is not None:
46+
ckpt_kwargs["checkpoint_id"] = ckpt_id
47+
ckpt_kwargs["process_group"] = distW.group
48+
49+
@_dcp_method_logger(**ckpt_kwargs)
50+
def local_step():
51+
assert planner is not None
52+
storage_meta = storage_writer.storage_meta()
53+
if "storage_meta" not in inspect.signature(planner.set_up_planner).parameters:
54+
warnings.warn(
55+
"The function definition for SavePlanner.set_up_planner has been updated"
56+
" to include the storage_meta argument. Please update your implementation"
57+
" to include this parameter."
58+
)
59+
planner.set_up_planner(state_dict, distW.is_coordinator) # type: ignore[call-arg, arg-type]
60+
else:
61+
planner.set_up_planner(
62+
state_dict=state_dict,
63+
storage_meta=storage_meta,
64+
is_coordinator=distW.is_coordinator,
65+
)
66+
storage_writer.set_up_storage_writer(distW.is_coordinator)
67+
68+
local_plan = planner.create_local_plan()
69+
local_plan = storage_writer.prepare_local_plan(local_plan)
70+
return local_plan
71+
72+
@_dcp_method_logger(**ckpt_kwargs)
73+
def global_step(all_local_plans):
74+
nonlocal global_metadata
75+
76+
assert planner is not None
77+
all_local_plans, global_metadata = planner.create_global_plan(all_local_plans)
78+
all_local_plans = storage_writer.prepare_global_plan(all_local_plans)
79+
return all_local_plans
80+
81+
central_plan: SavePlan = distW.reduce_scatter("plan", local_step, global_step)
82+
83+
@_dcp_method_logger(**ckpt_kwargs)
84+
def write_data():
85+
assert planner is not None
86+
final_local_plan = planner.finish_plan(central_plan)
87+
all_writes = storage_writer.write_data(final_local_plan, planner)
88+
89+
all_writes.wait()
90+
return all_writes.value()
91+
92+
@_dcp_method_logger(**ckpt_kwargs)
93+
def finish_checkpoint(all_results):
94+
assert global_metadata is not None
95+
storage_writer.finish(metadata=global_metadata, results=all_results)
96+
# return global_metadata
97+
return Metadata(state_dict_metadata={}) # This is a patch to avoid broadcast overhead.
98+
99+
return distW.all_reduce("write", write_data, finish_checkpoint)
100+
101+
102+
def patch_dcp_save_state_dict():
103+
if hasattr(state_dict_saver, "_save_state_dict"):
104+
original = getattr(state_dict_saver, "_save_state_dict")
105+
if callable(original):
106+
state_dict_saver._save_state_dict = _xtuner_save_state_dict
107+
108+
109+
def patch_dcp_save_with_cache_storage():
110+
original_dcp_save = dcp.save
111+
112+
def dcp_save_with_cache_storage(state_dict, **kwargs):
113+
checkpoint_id = kwargs.get("checkpoint_id", None)
114+
assert checkpoint_id is not None, "checkpoint_id is required for caching mechanism."
115+
checkpoint_id = checkpoint_id if isinstance(checkpoint_id, Path) else Path(checkpoint_id)
116+
planner = kwargs.get("planner", None)
117+
storage_writer = kwargs.get("storage_writer", None)
118+
119+
if storage_writer is None and planner is None:
120+
from xtuner.v1.patch.xtuner_cache_planner import XtunerCacheSavePlanner
121+
from xtuner.v1.patch.xtuner_storage import XtunerCacheWriter
122+
123+
planner = XtunerCacheSavePlanner(enable_plan_caching=True, cache_key_prefix=checkpoint_id.stem)
124+
storage_writer = XtunerCacheWriter(
125+
checkpoint_id, enable_write_result_caching=True, cache_key_prefix=checkpoint_id.stem
126+
)
127+
kwargs["planner"] = planner
128+
kwargs["storage_writer"] = storage_writer
129+
130+
return original_dcp_save(state_dict, **kwargs)
131+
132+
dcp.save = dcp_save_with_cache_storage

‎xtuner/v1/patch/xtuner_cache_planner.py‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
from typing import Optional
2+
3+
from torch.distributed.checkpoint import DefaultSavePlanner, Metadata, SavePlan, SavePlanner
4+
from torch.distributed.checkpoint.planner_helpers import ( # type: ignore[attr-defined]
5+
_compare_save_plans,
6+
_merge_delta_local_plans,
7+
)
8+
9+
10+
# copy from torch 2.8.0 planner_helpers.py
11+
def _contains_usable_plan(delta_plans: list[SavePlan]) -> bool:
12+
"""Check if any delta plan is usable, indicating the plan has changed.
13+
14+
Args:
15+
delta_plans (List[SavePlan]): A list of delta plans to check.
16+
Returns:
17+
True if any delta plan is usable, False otherwise.
18+
"""
19+
return any(delta_plan and delta_plan.usable for delta_plan in delta_plans) # type: ignore[attr-defined]
20+
21+
22+
class XtunerCacheSavePlanner(DefaultSavePlanner):
23+
# Metadata for the global checkpoint plan as computed by `create_global_plan` API.
24+
# Cached on the coordinator rank.
25+
_cached_metadata: dict[str, Metadata] = {}
26+
27+
def __init__(
28+
self,
29+
flatten_state_dict: bool = True,
30+
flatten_sharded_tensors: bool = True,
31+
dedup_replicated_tensors: Optional[bool] = None,
32+
dedup_save_to_lowest_rank: bool = False,
33+
enable_plan_caching: bool = False,
34+
cache_key_prefix: str = "",
35+
) -> None:
36+
super().__init__(
37+
flatten_state_dict,
38+
flatten_sharded_tensors,
39+
dedup_replicated_tensors,
40+
dedup_save_to_lowest_rank,
41+
enable_plan_caching, # type: ignore[call-arg]
42+
)
43+
self._cached_plans_key: str = cache_key_prefix + self.__class__.__name__
44+
45+
def _create_global_plan_with_caching(
46+
self, all_plans: list[SavePlan]
47+
) -> tuple[list[SavePlan], list[SavePlan], Metadata]:
48+
if hasattr(SavePlanner, "_cached_metadata"):
49+
# adaptor for torch >= 2.8.0
50+
return super()._create_global_plan_with_caching(all_plans) # type: ignore[misc]
51+
52+
# ONLY cache ``_cached_metadata`` in XtunerCacheSavePlanner
53+
global_plan_delta: list[SavePlan] = []
54+
55+
if self._cached_plans_key not in SavePlanner._cached_all_plans: # type: ignore[attr-defined]
56+
# Case 1: If the plans are not cached, the cache will be hydrated with the
57+
# all_plans, global_plans (Deduped), and metadata.
58+
59+
# Cache the original all_plans
60+
SavePlanner._cached_all_plans[self._cached_plans_key] = all_plans # type: ignore[attr-defined]
61+
global_plan, metadata = self._create_global_plan(all_plans) # type: ignore[attr-defined]
62+
# Cache the deduped and validated global_plan
63+
SavePlanner._cached_global_plan[self._cached_plans_key] = global_plan # type: ignore[attr-defined]
64+
# Cache the metadata
65+
XtunerCacheSavePlanner._cached_metadata[self._cached_plans_key] = metadata
66+
# If plans are not cached, global_plan delta will be the same as global plan.
67+
return global_plan, global_plan, metadata
68+
69+
# Case 2: Plans are cached
70+
if not _contains_usable_plan(all_plans):
71+
# Case 2.1: Plans are cached and the local plans have NOT changed (No usable plans).
72+
# Global plan delta will be empty plans to avoid the collective overhead.
73+
# We can reuse the deduped global plan and metadata from the cache directly.
74+
global_plan_delta = [SavePlan([], usable=False)] * len(all_plans) # type: ignore[call-arg]
75+
global_plan = SavePlanner._cached_global_plan[self._cached_plans_key] # type: ignore[attr-defined]
76+
metadata = XtunerCacheSavePlanner._cached_metadata[self._cached_plans_key]
77+
else:
78+
# Case 2.2: Plans are cached but the local plans have changed.
79+
# We will merge the changed local plans with the cached local plans.
80+
# Updated plans will overwrite the cached plans. New global plan and metadata will be created and cached.
81+
# Global plan delta will be created by comparing the new global plan with the cached global plan.
82+
# Only the global plan delta (updated ones) will be sent to the coordinator to avoid the collective overhead.
83+
merged_plans = _merge_delta_local_plans(SavePlanner._cached_all_plans[self._cached_plans_key], all_plans) # type: ignore[attr-defined]
84+
# Cache the updated local plans
85+
SavePlanner._cached_all_plans[self._cached_plans_key] = merged_plans # type: ignore[attr-defined]
86+
global_plan, metadata = self._create_global_plan(merged_plans) # type: ignore[attr-defined]
87+
88+
if self._cached_plans_key in self._cached_global_plan: # type: ignore[attr-defined]
89+
for cached_plan, new_plan in zip(SavePlanner._cached_global_plan[self._cached_plans_key], global_plan): # type: ignore[attr-defined]
90+
if _compare_save_plans(cached_plan, new_plan):
91+
global_plan_delta.append(SavePlan([], usable=False)) # type: ignore[call-arg]
92+
else:
93+
global_plan_delta.append(new_plan)
94+
95+
# Cache the new global plan and the metadata
96+
SavePlanner._cached_global_plan[self._cached_plans_key] = global_plan # type: ignore[attr-defined]
97+
XtunerCacheSavePlanner._cached_metadata[self._cached_plans_key] = metadata
98+
99+
return global_plan_delta, global_plan, metadata

‎xtuner/v1/patch/xtuner_storage.py‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
import os
2+
from collections.abc import Sequence
3+
from typing import Any, Optional, Union
4+
5+
import torch
6+
from packaging import version
7+
from torch.distributed.checkpoint import FileSystemWriter, Metadata, SavePlan, SavePlanner
8+
from torch.distributed.checkpoint._extension import (
9+
StreamTransformExtension,
10+
)
11+
from torch.distributed.checkpoint.storage import (
12+
WriteResult,
13+
)
14+
from torch.futures import Future
15+
16+
17+
# PyTorch 2.7+ introduced _extensions parameter for FileSystemWriter
18+
_TORCH_DCP_FSWRITER_HAS_EXTENSIONS = version.parse(torch.__version__) >= version.parse("2.7.0")
19+
20+
21+
def _compare_write_results(write_results: list[WriteResult], other_write_results: list[WriteResult]) -> bool:
22+
"""Compare two lists of WriteResults for equality.
23+
24+
Args:
25+
write_results: First list of WriteResults to compare.
26+
other_write_results: Second list of WriteResults to compare.
27+
28+
Returns:
29+
True if both lists have the same length and all elements are equal,
30+
False otherwise.
31+
"""
32+
33+
# Both the plans should have the same number of items
34+
if len(write_results) != len(other_write_results):
35+
return False
36+
37+
# Both the plans should have the same write items.
38+
for write_item, other_write_item in zip(write_results, other_write_results):
39+
# Write item type should be same
40+
if write_item != other_write_item:
41+
return False
42+
43+
return True
44+
45+
46+
def _contains_new_write_results(results: list[list[WriteResult]]) -> bool:
47+
return any(delta_result for delta_result in results)
48+
49+
50+
class XtunerCacheWriter(FileSystemWriter):
51+
# Save write results for the current rank as computed by `write_data` API
52+
# Cached on the local rank.
53+
_cache_write_results: dict[str, list[WriteResult]] = {}
54+
55+
# Collection of all the write results from all the ranks.
56+
# This is the ``results`` input to the `finish` API.
57+
# Cached on the coordinator rank.
58+
_cached_all_write_results: dict[str, list[list[WriteResult]]] = {}
59+
60+
def __init__(
61+
self,
62+
path: Union[str, os.PathLike],
63+
single_file_per_rank: bool = True,
64+
sync_files: bool = True,
65+
thread_count: int = 1,
66+
per_thread_copy_ahead: int = 10_000_000,
67+
cache_staged_state_dict: bool = False,
68+
overwrite: bool = True,
69+
_extensions: Optional[Sequence[StreamTransformExtension]] = None,
70+
enable_write_result_caching: bool = False,
71+
cache_key_prefix: str = "",
72+
) -> None:
73+
# Build kwargs conditionally to support both PyTorch 2.6 and 2.7+
74+
kwargs: dict[str, Any] = dict()
75+
if _TORCH_DCP_FSWRITER_HAS_EXTENSIONS:
76+
kwargs["_extensions"] = _extensions
77+
super().__init__(
78+
path,
79+
single_file_per_rank=single_file_per_rank,
80+
sync_files=sync_files,
81+
thread_count=thread_count,
82+
per_thread_copy_ahead=per_thread_copy_ahead,
83+
cache_staged_state_dict=cache_staged_state_dict,
84+
overwrite=overwrite,
85+
**kwargs,
86+
)
87+
self._enable_write_result_caching = enable_write_result_caching
88+
self._cached_write_results_key = cache_key_prefix + self.__class__.__name__
89+
90+
def write_data(
91+
self,
92+
plan: SavePlan,
93+
planner: SavePlanner,
94+
) -> Future[list[WriteResult]]:
95+
all_writes_fut = super().write_data(plan, planner)
96+
97+
if self._enable_write_result_caching:
98+
all_writes_fut = self._get_write_future_with_caching(all_writes_fut)
99+
return all_writes_fut
100+
101+
def _get_write_future_with_caching(self, all_writes_fut):
102+
new_fut: Future[list[WriteResult]] = Future()
103+
all_writes_fut.wait()
104+
105+
if self._cached_write_results_key not in XtunerCacheWriter._cache_write_results:
106+
# Case 1: If the write results are not cached,.............
107+
XtunerCacheWriter._cache_write_results[self._cached_write_results_key] = all_writes_fut.value()
108+
new_fut.set_result(all_writes_fut.value())
109+
elif _compare_write_results(
110+
all_writes_fut.value(), XtunerCacheWriter._cache_write_results[self._cached_write_results_key]
111+
):
112+
# Case 2: equal
113+
new_fut.set_result([])
114+
else:
115+
# Case 3: not equal
116+
XtunerCacheWriter._cache_write_results[self._cached_write_results_key] = all_writes_fut.value()
117+
new_fut.set_result(all_writes_fut.value())
118+
119+
return new_fut
120+
121+
def finish(self, metadata: Metadata, results: list[list[WriteResult]]) -> None:
122+
if self._enable_write_result_caching:
123+
results = self._get_results_from_caching(results)
124+
125+
super().finish(metadata, results)
126+
127+
def _get_results_from_caching(self, results: list[list[WriteResult]]):
128+
if self._cached_write_results_key not in XtunerCacheWriter._cached_all_write_results:
129+
# Case 1:
130+
XtunerCacheWriter._cached_all_write_results[self._cached_write_results_key] = results
131+
elif not _contains_new_write_results(results):
132+
# Case 2: no new
133+
results = XtunerCacheWriter._cached_all_write_results[self._cached_write_results_key]
134+
else:
135+
# Case 3: not equal TODO: merge
136+
XtunerCacheWriter._cached_all_write_results[self._cached_write_results_key] = results
137+
138+
return results

0 commit comments

Comments
 (0)