(self, load_dir, load_module_strict=True, tag=None)
| 411 | return ckpt_name |
| 412 | |
| 413 | def _load_checkpoint(self, load_dir, load_module_strict=True, tag=None): |
| 414 | is_pipe_parallel = isinstance(self.module, PipelineModule) |
| 415 | if is_pipe_parallel: |
| 416 | raise RuntimeError('pipeline parallelism is currently not supported in inference.') |
| 417 | if not isinstance(load_dir, dict) and os.path.isdir(load_dir): |
| 418 | if tag is None: |
| 419 | latest_path = os.path.join(load_dir, "latest") |
| 420 | if os.path.isfile(latest_path): |
| 421 | with open(latest_path, "r") as fd: |
| 422 | tag = fd.read().strip() |
| 423 | |
| 424 | ckpt_list = self._get_all_ckpt_names(load_dir, tag) |
| 425 | sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, self.checkpoint_engine) |
| 426 | else: |
| 427 | sd_loader = SDLoaderFactory.get_sd_loader_json(load_dir, self.checkpoint_engine) |
| 428 | |
| 429 | checkpoint = sd_loader['checkpoints'] |
| 430 | |
| 431 | if type(checkpoint) is list: |
| 432 | self.sd = torch.load(checkpoint[0], map_location='cpu', weights_only=False) |
| 433 | self.key_list = list(self.sd.keys()) |
| 434 | |
| 435 | self.load_model_with_checkpoint(self.module) |
| 436 | |
| 437 | for i in range(1, len(checkpoint)): |
| 438 | if not dist.is_initialized() or dist.get_rank() == 0: |
| 439 | print(f"loading checkpoint ({i})") |
| 440 | self.sd = torch.load(checkpoint[i], map_location=get_accelerator().device_name(), weights_only=False) |
| 441 | self.key_list = list(self.sd.keys()) |
| 442 | self.load_model_with_checkpoint(self.module) |
| 443 | else: |
| 444 | mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() |
| 445 | |
| 446 | load_path, checkpoint, quantize_config = sd_loader.load(self._config.tensor_parallel.tp_size, |
| 447 | mp_rank, |
| 448 | is_pipe_parallel=is_pipe_parallel, |
| 449 | quantize=(self._config.dtype is torch.int8), |
| 450 | quantize_groups=self.quantize_groups, |
| 451 | mlp_extra_grouping=self.mlp_extra_grouping) |
| 452 | |
| 453 | self.quantization_scales, self.quantize_merge_count = quantize_config |
| 454 | |
| 455 | moe, _ = has_moe_layers(self.module) |
| 456 | if moe: |
| 457 | from deepspeed.runtime.engine import DeepSpeedEngine |
| 458 | old_moe_load = False |
| 459 | if not isinstance(checkpoint['num_experts'], list): |
| 460 | old_moe_load = True |
| 461 | DeepSpeedEngine.load_moe_state_dict(load_dir, |
| 462 | tag, |
| 463 | state_dict=checkpoint[self._choose_module_key(checkpoint)], |
| 464 | old_moe_load=old_moe_load, |
| 465 | model=self.module, |
| 466 | mpu=self.mpu, |
| 467 | checkpoint_engine=self.checkpoint_engine, |
| 468 | autoep_layers=None) |
| 469 | |
| 470 | self.module.load_state_dict(state_dict=checkpoint[self._choose_module_key(checkpoint)], |
nothing calls this directly
no test coverage detected