(self, load_dir, tag, load_optimizer_states=True)
| 4721 | return load_path, client_state |
| 4722 | |
| 4723 | def _load_zero_checkpoint(self, load_dir, tag, load_optimizer_states=True): |
| 4724 | |
| 4725 | load_serial = None |
| 4726 | # When use loading checkpoint serial, checkpoint loading start from local rank 0, |
| 4727 | # all other local rank would be paused, waiting for its rank-1 peer ready and its notification. |
| 4728 | if self._config.zero_config.pipeline_loading_checkpoint: |
| 4729 | assert self.zero_optimization_stage( |
| 4730 | ) == ZeroStageEnum.weights, "Only stage3 support for pipeline checkpoint loading" |
| 4731 | load_serial = torch.zeros(1).to(self.device) |
| 4732 | if dist.get_local_rank() != 0: |
| 4733 | dist.recv(tensor=load_serial, src=dist.get_rank() - 1) |
| 4734 | if self.load_universal_checkpoint(): |
| 4735 | zero_sd_list = None |
| 4736 | checkpoint_folder = f'{os.path.join(load_dir, tag)}' |
| 4737 | else: |
| 4738 | if load_optimizer_states and self.seq_dp_world_size != self.loaded_checkpoint_dp_world_size: |
| 4739 | raise ZeRORuntimeException("The checkpoint being loaded used a DP " \ |
| 4740 | f"world size of {self.loaded_checkpoint_dp_world_size} but the " \ |
| 4741 | f"current world size is {self.seq_dp_world_size}. Automatic adjustment " \ |
| 4742 | "of ZeRO's optimizer state partitioning with a new world size is not " \ |
| 4743 | "currently supported.") |
| 4744 | checkpoint_folder = None |
| 4745 | zero_sd_list = self._get_all_zero_checkpoints(load_dir, tag) |
| 4746 | if zero_sd_list is None: |
| 4747 | return False |
| 4748 | |
| 4749 | param_shapes = self._get_zero_param_shapes() |
| 4750 | self.optimizer.load_state_dict(state_dict_list=zero_sd_list, |
| 4751 | load_optimizer_states=load_optimizer_states, |
| 4752 | load_from_fp32_weights=self.zero_load_from_fp32_weights(), |
| 4753 | checkpoint_folder=checkpoint_folder, |
| 4754 | load_serial=load_serial, |
| 4755 | param_shapes=param_shapes) |
| 4756 | |
| 4757 | if self.load_universal_checkpoint(): |
| 4758 | logger.info(f'loaded universal zero checkpoints from {checkpoint_folder} for rank {self.global_rank}') |
| 4759 | else: |
| 4760 | logger.info(f"loading {len(zero_sd_list)} zero partition checkpoints for rank {self.global_rank}") |
| 4761 | return True |
| 4762 | |
| 4763 | def _get_mp_rank_zero_checkpoint_names(self, load_dir, tag, mp_rank, dp_world_size, bf16_mode): |
| 4764 | zero_ckpt_names = [] |
no test coverage detected