| 459 | |
| 460 | |
| 461 | class DistributedCheckpointer(AbstractCheckpointer): |
| 462 | KEYS_TO_SAVE = ["model", "optim", "scheduler", "trainer"] |
| 463 | |
| 464 | def __init__( |
| 465 | self, |
| 466 | config_checkpoint: CheckpointConfig, |
| 467 | config_job: JobConfig, |
| 468 | callbacks: Optional[callback.CallBackGroup] = None, |
| 469 | disable_async: bool = False, |
| 470 | ): |
| 471 | super().__init__(config_checkpoint, config_job, callbacks) |
| 472 | self.config_checkpoint = config_checkpoint |
| 473 | if config_checkpoint.dcp_async_mode_enabled: |
| 474 | self.async_mode = AsyncMode.ASYNC_WITH_PINNED_MEM |
| 475 | else: |
| 476 | self.async_mode = AsyncMode.DISABLED |
| 477 | |
| 478 | if disable_async: |
| 479 | self.async_mode = AsyncMode.DISABLED |
| 480 | |
| 481 | if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM: |
| 482 | ctx = get_context("spawn") |
| 483 | self.mp_queue_send = ctx.Queue() |
| 484 | self.mp_queue_recv = ctx.Queue() |
| 485 | self.mp = ctx.Process( |
| 486 | target=save_checkpoint_in_background, |
| 487 | args=( |
| 488 | self.mp_queue_send, |
| 489 | self.mp_queue_recv, |
| 490 | config_checkpoint, |
| 491 | config_job, |
| 492 | ), |
| 493 | daemon=True, |
| 494 | ) |
| 495 | self.mp.start() |
| 496 | self.cpu_offload_state_dict = None |
| 497 | self.staging = False |
| 498 | self.staging_ckpt_file = None |
| 499 | self.staging_stream = torch.cuda.Stream() |
| 500 | |
| 501 | def keys_to_resume_during_load(self) -> Tuple[Set, Union[str, None]]: |
| 502 | latest_checkpoint_file = self._read_latest_checkpoint_file() |
| 503 | |
| 504 | resume_keys = [] |
| 505 | |
| 506 | if latest_checkpoint_file is not None: |
| 507 | # 1. Resume training from latest_checkpoint.txt under the same name. |
| 508 | checkpoint_path = os.path.join(self.load_dirname, latest_checkpoint_file) |
| 509 | resume_keys.extend(self.KEYS_TO_SAVE) |
| 510 | else: |
| 511 | if self.load_path and not str(self.load_path).endswith(".pt"): |
| 512 | # 2. Load the module weights specified by config_checkpoint.path. |
| 513 | checkpoint_path = self.load_path |
| 514 | if self.load_s3_backend_key: |
| 515 | checkpoint_path = f"s3://{self.config_checkpoint.load_from_object_store.bucket}/{checkpoint_path}" |
| 516 | if not re.search(r"/checkpoints/iter_\d{9}/?$", checkpoint_path): |
| 517 | old_ckpt_path = checkpoint_path |
| 518 | # If path doesn't end with specific checkpoint, read latest checkpoint file |
no outgoing calls
no test coverage detected