(self, load_dir, tag, load_optimizer_states=True)
| 4543 | return load_path, client_state |
| 4544 | |
| 4545 | def _load_zero_checkpoint(self, load_dir, tag, load_optimizer_states=True): |
| 4546 | |
| 4547 | load_serial = None |
| 4548 | # When use loading checkpoint serial, checkpoint loading start from local rank 0, |
| 4549 | # all other local rank would be paused, waiting for its rank-1 peer ready and its notification. |
| 4550 | if self._config.zero_config.pipeline_loading_checkpoint: |
| 4551 | assert self.zero_optimization_stage( |
| 4552 | ) == ZeroStageEnum.weights, "Only stage3 support for pipeline checkpoint loading" |
| 4553 | load_serial = torch.zeros(1).to(self.device) |
| 4554 | if dist.get_local_rank() != 0: |
| 4555 | dist.recv(tensor=load_serial, src=dist.get_rank() - 1) |
| 4556 | if self.load_universal_checkpoint(): |
| 4557 | zero_sd_list = None |
| 4558 | checkpoint_folder = f'{os.path.join(load_dir, tag)}' |
| 4559 | else: |
| 4560 | if load_optimizer_states and self.seq_dp_world_size != self.loaded_checkpoint_dp_world_size: |
| 4561 | raise ZeRORuntimeException("The checkpoint being loaded used a DP " \ |
| 4562 | f"world size of {self.loaded_checkpoint_dp_world_size} but the " \ |
| 4563 | f"current world size is {self.seq_dp_world_size}. Automatic adjustment " \ |
| 4564 | "of ZeRO's optimizer state partitioning with a new world size is not " \ |
| 4565 | "currently supported.") |
| 4566 | checkpoint_folder = None |
| 4567 | zero_sd_list = self._get_all_zero_checkpoints(load_dir, tag) |
| 4568 | if zero_sd_list is None: |
| 4569 | return False |
| 4570 | |
| 4571 | param_shapes = self._get_zero_param_shapes() |
| 4572 | self.optimizer.load_state_dict(state_dict_list=zero_sd_list, |
| 4573 | load_optimizer_states=load_optimizer_states, |
| 4574 | load_from_fp32_weights=self.zero_load_from_fp32_weights(), |
| 4575 | checkpoint_folder=checkpoint_folder, |
| 4576 | load_serial=load_serial, |
| 4577 | param_shapes=param_shapes) |
| 4578 | |
| 4579 | if self.load_universal_checkpoint(): |
| 4580 | logger.info(f'loaded universal zero checkpoints from {checkpoint_folder} for rank {self.global_rank}') |
| 4581 | else: |
| 4582 | logger.info(f"loading {len(zero_sd_list)} zero partition checkpoints for rank {self.global_rank}") |
| 4583 | return True |
| 4584 | |
| 4585 | def _get_mp_rank_zero_checkpoint_names(self, load_dir, tag, mp_rank, dp_world_size, bf16_mode): |
| 4586 | zero_ckpt_names = [] |
no test coverage detected