Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / api.py: 87%
132 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd. All rights reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Hyper Parallel Checkpoint API"""
17import itertools
18import os
19import threading
20import multiprocessing as mp
21from collections.abc import Callable
22from concurrent.futures import Future
23from contextlib import nullcontext
24from pathlib import Path
25from typing import Any, Optional, Union
27import torch.distributed as dist
29from hyper_parallel.core.distributed_checkpoint.async_persist import (
30 AsyncSaveResponse,
31 FileType,
32 build_staged_state_dict,
33 execute_async_persist_with_gloo,
34 execute_async_persist,
35 resolve_async_persist_result,
36 gather_all_results_from_storage,
37)
38from hyper_parallel.core.distributed_checkpoint.filesystem_storage import (
39 FileSystemReader,
40 FileSystemWriter,
41)
42from hyper_parallel.core.distributed_checkpoint.metadata import (
43 Metadata,
44)
45from hyper_parallel.core.distributed_checkpoint.planner import (
46 SavePlan,
47 SavePlanner,
48 LoadPlanner,
49)
50from hyper_parallel.core.distributed_checkpoint.standard_planner import (
51 StandardSavePlanner,
52 StandardLoadPlanner,
53)
54from hyper_parallel.core.distributed_checkpoint.storage import (
55 StorageReader,
56 StorageWriter,
57)
58from hyper_parallel.core.distributed_checkpoint.broadcast import (
59 broadcast_groups_for_load,
60 DEFAULT_BROADCAST_BATCH_BYTES,
61)
62from hyper_parallel.core.distributed_checkpoint.utils import (
63 all_gather_object,
64 dcp_timer_decorator,
65 logger,
66)
68# Names the gloo persist process of each async save, so several in one job stay apart.
69_async_save_counter = itertools.count()
72@dcp_timer_decorator
73def _save_impl(
74 state_dict: dict[str, Any],
75 *,
76 checkpoint_id: Optional[Union[Path, str]] = None,
77 storage_writer: Optional[StorageWriter] = None,
78 planner: Optional[SavePlanner] = None,
79 no_dist: bool = False,
80 use_collectives: bool = True,
81 use_storage_comm: bool = False,
82) -> Metadata:
83 """Distributed checkpoint save implement(shared by :func:`save` and :func:`async_save`)."""
84 # Convert checkpoint_id to Path if it's a string
85 checkpoint_id = Path(checkpoint_id) if isinstance(checkpoint_id, str) else checkpoint_id
87 # Determine if we're in distributed mode
88 use_collectives = False if no_dist else use_collectives
89 use_storage_comm = False if not use_collectives else use_storage_comm
91 # Set up storage writer
92 if storage_writer is None:
93 if checkpoint_id is None:
94 raise ValueError("Either storage_writer or checkpoint_id must be provided")
95 storage_writer = FileSystemWriter(checkpoint_id)
96 else:
97 if checkpoint_id:
98 storage_writer.initialize_writer(checkpoint_id)
99 else:
100 checkpoint_id = getattr(storage_writer, "checkpoint_dir", None)
102 if use_storage_comm and checkpoint_id is None:
103 raise ValueError(
104 "Coordinating through the storage needs a checkpoint directory to exchange the "
105 "plans and write results in: pass checkpoint_id, or a storage_writer exposing "
106 "checkpoint_dir."
107 )
109 # Set up planner
110 planner = StandardSavePlanner() if planner is None else planner
112 # Get rank and coordinator info
113 rank = dist.get_rank()
114 world_size = dist.get_world_size()
115 is_coordinator = rank == 0
117 # Configure planner
118 planner.configure_planner(
119 state_dict=state_dict,
120 is_coordinator=is_coordinator,
121 rank=rank,
122 use_collectives=use_collectives
123 )
125 # Configure storage writer (use_collectives for rank-local metadata when False)
126 storage_writer.configure_writer(
127 is_coordinator=is_coordinator,
128 rank=rank,
129 use_collectives=use_collectives
130 )
132 @dcp_timer_decorator
133 def generate_final_plan_and_metadata() -> tuple[SavePlan, Metadata]:
134 """Build this rank's final plan and the checkpoint metadata, caching both.
136 Returns:
137 tuple[SavePlan, Metadata]: The finalized local plan and the global metadata.
138 """
139 # Build local plan
140 local_plan = planner.build_local_plan()
141 local_plan = storage_writer.optimize_local_plan(local_plan)
143 #Gather all local plans and build global plan
144 all_local_plans = gather_all_results_from_storage(checkpoint_id,
145 is_coordinator,
146 world_size,
147 local_plan,
148 FileType.LOCAL_PLAN,
149 rank) if use_storage_comm \
150 else all_gather_object(local_plan, world_size, use_collectives)
151 # build_global_plan returns this rank's plan only: every rank runs it over the same
152 # gathered plans and writes just its own shards.
153 central_plan, global_metadata = planner.build_global_plan(all_local_plans)
154 central_plan = storage_writer.optimize_global_plan(central_plan)
156 # Finalize and cache plan
157 finalized_plan = planner.finalize_plan(central_plan)
158 # Add final plan and metadata to the cache
159 if hasattr(planner, 'cache_result'):
160 planner.cache_result(finalized_plan, global_metadata)
161 return finalized_plan, global_metadata
163 cached_res = planner.get_cached() if hasattr(planner, 'get_cached') else None
164 if cached_res:
165 # Get final plan and metadata from cache
166 logger.info("Hit final plan and metadata cache.")
167 final_plan, metadata = cached_res.final_plan, cached_res.metadata
168 else:
169 # First time generating a plan and metadata.
170 final_plan, metadata = generate_final_plan_and_metadata()
172 # Write data
173 write_results = storage_writer.execute_write(final_plan, planner)
175 # Finalize checkpoint
176 all_write_results = gather_all_results_from_storage(checkpoint_id,
177 is_coordinator,
178 world_size,
179 write_results,
180 FileType.STORAGE_DATA,
181 rank) if use_storage_comm \
182 else all_gather_object(write_results, world_size, use_collectives)
183 storage_writer.finalize_checkpoint(metadata, all_write_results)
185 return metadata
188@dcp_timer_decorator
189def save(
190 state_dict: dict[str, Any],
191 *,
192 checkpoint_id: Optional[Union[Path, str]] = None,
193 storage_writer: Optional[StorageWriter] = None,
194 planner: Optional[SavePlanner] = None,
195 no_dist: bool = False,
196 use_collectives: bool = True,
197) -> Metadata:
198 """
199 Save a distributed checkpoint in SPMD style.
201 This function saves a state_dict containing DTensors, where each rank
202 only saves their local shards.
204 Args:
205 state_dict (dict[str, Any]): The state_dict to save.
206 checkpoint_id (Optional[Union[Path, str]]): The ID/path of this checkpoint instance (can be Path or str).
207 Default None.
208 storage_writer (Optional[StorageWriter]): Instance of StorageWriter. If None, FileSystemWriter
209 will be created based on checkpoint_id. Default None.
210 planner (Optional[SavePlanner]): Instance of SavePlanner. If None, StandardSavePlanner will be used.
211 Default None.
212 no_dist (bool): If True, save in single process mode. Default False.
213 use_collectives (bool): If True, use collective communication for coordination.
214 If False, each rank saves its own shard data and rank-local metadata (.metadata_rank{rank}),
215 with no cross-rank interaction. Default True.
217 Returns:
218 Metadata: Metadata object for the saved checkpoint.
219 """
220 metadata = _save_impl(
221 state_dict,
222 checkpoint_id=checkpoint_id,
223 storage_writer=storage_writer,
224 planner=planner,
225 no_dist=no_dist,
226 use_collectives=use_collectives,
227 )
228 dist.barrier()
229 return metadata
232def _create_persist_process(
233 result_queue: mp.Queue,
234 staged: dict[str, Any],
235 checkpoint_id: Optional[Union[Path, str]],
236 storage_writer: Optional[StorageWriter],
237 planner: Optional[SavePlanner],
238 no_dist: bool,
239 use_collectives: bool,
240 use_gloo: bool,
241) -> mp.Process:
242 """
243 Build the child process persisting the staged state_dict for :func:`async_save`.
245 Depending on the caller, the child either runs the plan and write result collectives on a
246 dedicated gloo group it reinitialises itself, or exchanges those results through the
247 storage, or runs without any cross-rank interaction at all.
249 Args:
250 result_queue (mp.Queue): IPC queue the child returns its status and metadata on.
251 staged (dict[str, Any]): Host resident staged state_dict to persist.
252 checkpoint_id (Optional[Union[Path, str]]): Same as :func:`async_save`.
253 storage_writer (Optional[StorageWriter]): Same as :func:`async_save`.
254 planner (Optional[SavePlanner]): Same as :func:`async_save`.
255 no_dist (bool): Same as :func:`async_save`.
256 use_collectives (bool): Same as :func:`async_save`.
257 use_gloo (bool): Same as :func:`async_save`.
259 Returns:
260 mp.Process: The persist process, not started yet.
262 Raises:
263 AssertionError: If gloo communication is requested without MASTER_ADDR or MASTER_PORT.
264 """
265 world_size = dist.get_world_size()
266 need_comm = use_collectives and not no_dist and world_size > 1
267 if not need_comm or not use_gloo:
268 use_storage_comm = need_comm and not use_gloo
269 if use_storage_comm:
270 logger.info("use storage communication for dcp async save.")
271 return mp.Process(
272 target=execute_async_persist,
273 args=(
274 result_queue,
275 staged,
276 checkpoint_id,
277 storage_writer,
278 planner,
279 no_dist,
280 use_collectives,
281 use_storage_comm,
282 ),
283 name="AsyncCheckpointPersist",
284 )
286 logger.info("use Gloo communication for dcp async save.")
287 rank = dist.get_rank()
288 master_addr = os.environ.get("MASTER_ADDR", None)
289 master_port = os.environ.get("MASTER_PORT", None)
290 if master_addr is None:
291 raise AssertionError("Async DCP needs MASTER_ADDR to use prefix store")
292 if master_port is None:
293 raise AssertionError("Async DCP needs MASTER_PORT to use prefix store")
294 master_port = int(master_port)
295 save_index = next(_async_save_counter)
296 proc = mp.Process(
297 target=execute_async_persist_with_gloo,
298 args=(
299 result_queue,
300 staged,
301 checkpoint_id,
302 storage_writer,
303 planner,
304 rank,
305 world_size,
306 master_addr,
307 master_port,
308 ),
309 name=f"AsyncCheckpointPersistGloo_{save_index}",
310 )
311 return proc
314@dcp_timer_decorator
315def async_save(
316 state_dict: dict[str, Any],
317 *,
318 checkpoint_id: Optional[Union[Path, str]] = None,
319 storage_writer: Optional[StorageWriter] = None,
320 planner: Optional[SavePlanner] = None,
321 no_dist: bool = False,
322 use_collectives: bool = True,
323 use_gloo: bool = False,
324 callback: Optional[Callable[[], None]] = None,
325) -> AsyncSaveResponse:
326 """
327 Asynchronous version of :func:`save` using a **background child process** for persistence.
329 **Staging** (tensor / DTensor → host copy) runs **synchronously in the caller
330 process** via :func:`build_staged_state_dict`, so no process pool is used for
331 staging and the training stack sees a normal Python call path. When this
332 function returns successfully, host staging is done and the original
333 ``state_dict`` may be mutated.
335 **Persistence** (plan, collectives, disk I/O) runs in **one** background
336 :class:`multiprocessing.Process` that executes :func:`_save_impl` on the staged
337 dict. A small background **thread** only joins that process and fills
338 ``persist_completion``; it does not perform tensor work.
340 Neither the thread nor the child process is a daemon, so a caller that exits with a
341 checkpoint still in flight waits for it to land rather than losing it. Making only the
342 thread a daemon would not shorten that wait either, since ``multiprocessing`` joins the
343 non-daemon child at interpreter exit anyway; it would just drop ``persist_completion``
344 and ``callback`` on the floor.
346 The staged dict and ``storage_writer`` / ``planner`` must be picklable for the
347 persist child process (same constraints as before for the worker path).
349 .. warning::
350 Experimental API. Always wait on ``persist_completion`` for a fully persisted checkpoint.
352 Args:
353 state_dict (dict[str, Any]): The state_dict to save.
354 checkpoint_id (Optional[Union[Path, str]]): Same as :func:`save`.
355 storage_writer (Optional[StorageWriter]): Same as :func:`save`.
356 planner (Optional[SavePlanner]): Same as :func:`save`.
357 no_dist (bool): Same as :func:`save`.
358 use_collectives (bool): Same as :func:`save`.
359 use_gloo(bool): Whether to use Gloo. Default False.
360 callback(Callable): Async save callback function.
362 Returns:
363 AsyncSaveResponse: Contains ``persist_completion`` only; staging is synchronous.
364 """
365 # Copy the state_dict to cpu memory for async save
366 staged = build_staged_state_dict(state_dict)
368 persist_completion: Future[Metadata] = Future()
369 result_queue: mp.Queue = mp.Queue(maxsize=1)
370 proc = _create_persist_process(
371 result_queue,
372 staged,
373 checkpoint_id,
374 storage_writer,
375 planner,
376 no_dist,
377 use_collectives,
378 use_gloo,
379 )
381 # After the async save is completed, the user callback will be executed.
382 def async_callback() -> None:
383 """Run the user callback once the save has been persisted, if one was given."""
384 if callback is not None:
385 callback()
387 # Deliberately not a daemon: it has to outlive the training loop long enough to resolve
388 # persist_completion and run callback, and killing it early would not let the process exit
389 # any sooner because the persist child is joined at interpreter exit regardless.
390 join_thread = threading.Thread(
391 target=resolve_async_persist_result,
392 args=(proc, result_queue, persist_completion, async_callback),
393 name="AsyncCheckpointPersistJoin",
394 )
395 proc.start()
396 join_thread.start()
397 ret = AsyncSaveResponse(persist_completion=persist_completion)
398 return ret
401def _open_checkpoint(
402 checkpoint_id: Optional[Path],
403 storage_reader: Optional[StorageReader],
404 rank: int,
405) -> tuple[StorageReader, Metadata, bool]:
406 """
407 The reader a load goes through, and what that reader says the checkpoint holds.
409 Args:
410 checkpoint_id (Optional[Path]): Where the checkpoint is, if the caller named it.
411 storage_reader (Optional[StorageReader]): The reader to read through, or None to
412 open a file system one on ``checkpoint_id``, which the caller has checked is
413 there to open one on.
414 rank (int): This rank, which names the rank-local metadata of a checkpoint that
415 has no shared one.
417 Returns:
418 tuple[StorageReader, Metadata, bool]: The reader, the metadata, and whether that
419 metadata came from the shared ``.metadata``. A checkpoint written without
420 collectives has only a ``{rank}.metadata`` apiece, and comes back False: its ranks
421 never agreed on one description of it, so a load of it has none to plan against
422 together either.
423 """
424 if storage_reader is None:
425 storage_reader = FileSystemReader(checkpoint_id)
426 elif checkpoint_id:
427 storage_reader.initialize_reader(checkpoint_id)
429 try:
430 return storage_reader, storage_reader.load_metadata(), True
431 except FileNotFoundError:
432 return storage_reader, storage_reader.load_metadata(rank=rank), False
435def load(
436 state_dict: dict[str, Any],
437 *,
438 checkpoint_id: Optional[Union[Path, str]] = None,
439 storage_reader: Optional[StorageReader] = None,
440 planner: Optional[LoadPlanner] = None,
441 no_dist: bool = False,
442 use_collectives: bool = True,
443 broadcast_replicated_tensors: bool = True,
444 broadcast_groups: Optional[dict[tuple, Any]] = None,
445 broadcast_batch_bytes: int = DEFAULT_BROADCAST_BATCH_BYTES,
446) -> None:
447 """
448 Load a distributed checkpoint into state_dict in SPMD style.
450 Each rank will try to read the least amount of data necessary
451 to fulfill the requested state_dict. When loading DTensor instances,
452 each rank only reads data for their local shards.
454 Args:
455 state_dict (dict[str, Any]): The state_dict to load the checkpoint into (modified in-place).
456 checkpoint_id (Optional[Union[Path, str]]): The ID/path of this checkpoint instance (can be Path or str).
457 Default None.
458 storage_reader (Optional[StorageReader]): Instance of StorageReader. If None, FileSystemReader
459 will be created based on checkpoint_id. Default None.
460 planner (Optional[LoadPlanner]): Instance of LoadPlanner. If None, StandardLoadPlanner will be used.
461 Default None.
462 no_dist (bool): If True, load without cross-rank synchronization. Default False.
463 use_collectives (bool): If False, each rank loads `.metadata` or `rank.metadata`,
464 there is no communication between ranks, and each rank loads its own data. Default True,
465 matching :func:`save`. A checkpoint written without it has no shared `.metadata`, and a
466 load falls back to the rank-local one and turns this off by itself, so the default costs
467 nothing to a checkpoint that cannot use it. What it does ask, as :func:`save` already
468 does, is that **every rank of the default process group call this together**: the plan
469 exchange runs on that group, so a rank that stays out of it leaves the others waiting on
470 a gather that never completes. A rank with nothing to load still takes part, passing an
471 empty state dict. Pass False to load on some of the ranks only.
472 broadcast_replicated_tensors (bool): If True, a tensor that several ranks load identically is
473 read by one of them and sent to the rest instead of every rank reading it. Which rank
474 reads which tensor is chosen so that the reads spread evenly over the group. Default
475 True, which divides what a load takes off storage by the number of ranks holding the
476 same copy. Only consulted together with ``use_collectives``: without the plan exchange
477 there is no way to tell which ranks load the same shard, so a load that asks for this
478 without it quietly reads every copy instead, which is correct but no faster.
479 broadcast_groups (dict): The Communication groups for broadcast, keyed by rank tuple.
480 Only valid together with broadcast_replicated_tensors; built on demand when not given.
481 broadcast_batch_bytes (int): Shards smaller than this are gathered into one buffer and
482 broadcast together, so that the cost of starting a broadcast is paid once for many
483 of them rather than once each. Larger shards are sent on their own. Zero turns the
484 gathering off. Only consulted together with broadcast_replicated_tensors.
485 Every rank must pass the same value: what a rank gathers together is worked
486 out from this and from the plan, without asking anybody, so ranks that
487 disagreed would broadcast in different shapes and wait on each other.
489 Returns:
490 None. The state_dict is modified in-place.
491 """
492 # Convert checkpoint_id to Path if it's a string
493 checkpoint_id = Path(checkpoint_id) if isinstance(checkpoint_id, str) else checkpoint_id
495 # Determine if we're in distributed mode
496 use_collectives = False if no_dist else use_collectives
498 # Check the arguments
499 if broadcast_groups and not broadcast_replicated_tensors:
500 raise ValueError("broadcast_groups is only used when broadcast_replicated_tensors is True.")
501 if storage_reader is None and checkpoint_id is None:
502 raise ValueError("Either storage_reader or checkpoint_id must be provided")
504 # Set up planner
505 planner = StandardLoadPlanner() if planner is None else planner
507 # Get rank and coordinator info
508 rank = dist.get_rank()
509 world_size = dist.get_world_size()
510 is_coordinator = rank == 0
512 # Open the checkpoint. One saved without collectives has no shared metadata to plan
513 # against, so a load of it cannot exchange plans however it was asked to.
514 storage_reader, metadata, shared_metadata = _open_checkpoint(checkpoint_id, storage_reader, rank)
515 use_collectives = use_collectives and shared_metadata
517 # Configure planner
518 planner.configure_planner(
519 state_dict=state_dict,
520 metadata=metadata,
521 is_coordinator=is_coordinator,
522 rank=rank,
523 broadcast_replicated_tensors=broadcast_replicated_tensors,
524 use_collectives=use_collectives,
525 )
527 # Configure storage reader
528 storage_reader.configure_reader(
529 metadata=metadata,
530 is_coordinator=is_coordinator,
531 rank=rank,
532 )
534 # Build local plan
535 local_plan = planner.build_local_plan()
536 local_plan = storage_reader.optimize_local_plan(local_plan)
538 # Gather all local plans and build global plan. Only the identity of each read is sent:
539 # the global plan matches reads across ranks on that alone, so where they come from in
540 # the checkpoint stays with the rank that will do them. This rank puts its own full plan
541 # back over what it just sent itself, since that is the one it goes on to execute.
542 all_local_plans = all_gather_object(local_plan.identity(), world_size, use_collectives)
543 all_local_plans[rank if len(all_local_plans) > 1 else 0] = local_plan
544 # build_global_plan returns this rank's plan only: every rank runs it over the same
545 # gathered plans and reads just its own shards.
546 central_plan = planner.build_global_plan(all_local_plans)
547 central_plan = storage_reader.optimize_global_plan(central_plan)
549 # Finalize plan
550 final_plan = planner.finalize_plan(central_plan)
552 # Build the groups the broadcasts go through, and destroy them once the read is
553 # done. Creating one is collective, so what decides whether to do it has to be
554 # something every rank agrees on before anybody reads -- these three are, where
555 # "does this rank hold a shared shard" is not.
556 scope = (
557 broadcast_groups_for_load(final_plan.items, broadcast_groups)
558 if broadcast_replicated_tensors and use_collectives and world_size > 1
559 else nullcontext(broadcast_groups)
560 )
562 # Execute read
563 with scope as groups:
564 storage_reader.execute_read(final_plan, planner, groups, broadcast_batch_bytes)