(self, load_dir, tag, load_optimizer_states=True)
| 4418 | return load_path, client_state |
| 4419 | |
| 4420 | def _load_zero_checkpoint(self, load_dir, tag, load_optimizer_states=True): |
| 4421 | |
| 4422 | load_serial = None |
| 4423 | # When use loading checkpoint serial, checkpoint loading start from local rank 0, |
| 4424 | # all other local rank would be paused, waiting for its rank-1 peer ready and its notification. |
| 4425 | if self._config.zero_config.pipeline_loading_checkpoint: |
| 4426 | assert self.zero_optimization_stage( |
| 4427 | ) == ZeroStageEnum.weights, "Only stage3 support for pipeline checkpoint loading" |
| 4428 | load_serial = torch.zeros(1).to(self.device) |
| 4429 | if dist.get_local_rank() != 0: |
| 4430 | dist.recv(tensor=load_serial, src=dist.get_rank() - 1) |
| 4431 | if self.load_universal_checkpoint(): |
| 4432 | zero_sd_list = None |
| 4433 | checkpoint_folder = f'{os.path.join(load_dir, tag)}' |
| 4434 | else: |
| 4435 | if load_optimizer_states and self.seq_dp_world_size != self.loaded_checkpoint_dp_world_size: |
| 4436 | raise ZeRORuntimeException("The checkpoint being loaded used a DP " \ |
| 4437 | f"world size of {self.loaded_checkpoint_dp_world_size} but the " \ |
| 4438 | f"current world size is {self.seq_dp_world_size}. Automatic adjustment " \ |
| 4439 | "of ZeRO's optimizer state partitioning with a new world size is not " \ |
| 4440 | "currently supported.") |
| 4441 | checkpoint_folder = None |
| 4442 | zero_sd_list = self._get_all_zero_checkpoints(load_dir, tag) |
| 4443 | if zero_sd_list is None: |
| 4444 | return False |
| 4445 | |
| 4446 | param_shapes = self._get_zero_param_shapes() |
| 4447 | self.optimizer.load_state_dict(state_dict_list=zero_sd_list, |
| 4448 | load_optimizer_states=load_optimizer_states, |
| 4449 | load_from_fp32_weights=self.zero_load_from_fp32_weights(), |
| 4450 | checkpoint_folder=checkpoint_folder, |
| 4451 | load_serial=load_serial, |
| 4452 | param_shapes=param_shapes) |
| 4453 | |
| 4454 | if self.load_universal_checkpoint(): |
| 4455 | logger.info(f'loaded universal zero checkpoints from {checkpoint_folder} for rank {self.global_rank}') |
| 4456 | else: |
| 4457 | logger.info(f"loading {len(zero_sd_list)} zero partition checkpoints for rank {self.global_rank}") |
| 4458 | return True |
| 4459 | |
| 4460 | def _get_mp_rank_zero_checkpoint_names(self, load_dir, tag, mp_rank, dp_world_size, bf16_mode): |
| 4461 | zero_ckpt_names = [] |
no test coverage detected