(self,
load_dir,
tag,
load_module_strict=True,
load_optimizer_states=True,
load_lr_scheduler_states=True,
load_module_only=False,
custom_load_fn=None)
| 4358 | return names |
| 4359 | |
| 4360 | def _load_checkpoint(self, |
| 4361 | load_dir, |
| 4362 | tag, |
| 4363 | load_module_strict=True, |
| 4364 | load_optimizer_states=True, |
| 4365 | load_lr_scheduler_states=True, |
| 4366 | load_module_only=False, |
| 4367 | custom_load_fn=None): |
| 4368 | |
| 4369 | from deepspeed.runtime.state_dict_factory import SDLoaderFactory |
| 4370 | |
| 4371 | ckpt_list = self._get_all_ckpt_names(load_dir, tag) |
| 4372 | sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine=self.checkpoint_engine) |
| 4373 | |
| 4374 | is_pipe_parallel = isinstance(self.module, PipelineModule) |
| 4375 | |
| 4376 | mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() |
| 4377 | load_path, checkpoint, _ = sd_loader.load(self.mp_world_size, mp_rank, is_pipe_parallel=is_pipe_parallel) |
| 4378 | |
| 4379 | if checkpoint is None: |
| 4380 | return None, None |
| 4381 | |
| 4382 | folding_spec = getattr(self, "_autoep_folding_spec", None) |
| 4383 | folded_autoep_tp = folding_spec is not None and folding_spec.tp_size > 1 |
| 4384 | ep_group_name = f"ep_size_{folding_spec.ep_size}" if folded_autoep_tp else None |
| 4385 | DeepSpeedEngine._validate_autoep_folding_checkpoint_metadata( |
| 4386 | checkpoint, |
| 4387 | folding_spec=folding_spec, |
| 4388 | family="dense", |
| 4389 | zero_partition_group="dense_dp", |
| 4390 | zero_partition_count=folding_spec.dp_size if folded_autoep_tp else None, |
| 4391 | tp_rank=groups.get_tensor_model_parallel_rank() if folded_autoep_tp else None) |
| 4392 | |
| 4393 | fetch_z3_params = False |
| 4394 | z3_params_to_fetch = None |
| 4395 | autoep_partitioned_experts = False |
| 4396 | allowed_missing_keys = None |
| 4397 | if self.zero_optimization_partition_weights() and not load_optimizer_states and not self.has_moe_layers: |
| 4398 | checkpoint['module'] = get_fp32_state_dict_from_zero_checkpoint(load_dir) |
| 4399 | fetch_z3_params = True |
| 4400 | |
| 4401 | if is_pipe_parallel: |
| 4402 | # Pipeline parallelism uses this to load its own checkpoint files. |
| 4403 | self._curr_ckpt_path = os.path.join(load_dir, tag) |
| 4404 | |
| 4405 | # Universal Checkpoint restores parameters from the zero/ layout, so |
| 4406 | # do not require regular MoE expert checkpoint files in that path. |
| 4407 | if self.has_moe_layers and not self.load_universal_checkpoint(): |
| 4408 | # print(checkpoint.keys()) |
| 4409 | old_moe_load = False |
| 4410 | if not isinstance(checkpoint['num_experts'], list): |
| 4411 | old_moe_load = True |
| 4412 | from deepspeed.checkpoint.constants import AUTOEP_LAYERS_KEY, AUTOEP_LAYERS_KEY_LEGACY |
| 4413 | autoep_layers = checkpoint.get(AUTOEP_LAYERS_KEY) |
| 4414 | if autoep_layers is None: |
| 4415 | autoep_layers = checkpoint.get(AUTOEP_LAYERS_KEY_LEGACY) |
| 4416 | autoep_partitioned_experts = (self.zero_optimization_partition_weights() |
| 4417 | and DeepSpeedEngine._uses_autoep_zero3_partitioned_experts(autoep_layers)) |
no test coverage detected