(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)
| 4233 | return names |
| 4234 | |
| 4235 | def _load_checkpoint(self, |
| 4236 | load_dir, |
| 4237 | tag, |
| 4238 | load_module_strict=True, |
| 4239 | load_optimizer_states=True, |
| 4240 | load_lr_scheduler_states=True, |
| 4241 | load_module_only=False, |
| 4242 | custom_load_fn=None): |
| 4243 | |
| 4244 | from deepspeed.runtime.state_dict_factory import SDLoaderFactory |
| 4245 | |
| 4246 | ckpt_list = self._get_all_ckpt_names(load_dir, tag) |
| 4247 | sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine=self.checkpoint_engine) |
| 4248 | |
| 4249 | is_pipe_parallel = isinstance(self.module, PipelineModule) |
| 4250 | |
| 4251 | mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank() |
| 4252 | load_path, checkpoint, _ = sd_loader.load(self.mp_world_size, mp_rank, is_pipe_parallel=is_pipe_parallel) |
| 4253 | |
| 4254 | if checkpoint is None: |
| 4255 | return None, None |
| 4256 | |
| 4257 | folding_spec = getattr(self, "_autoep_folding_spec", None) |
| 4258 | folded_autoep_tp = folding_spec is not None and folding_spec.tp_size > 1 |
| 4259 | ep_group_name = f"ep_size_{folding_spec.ep_size}" if folded_autoep_tp else None |
| 4260 | DeepSpeedEngine._validate_autoep_folding_checkpoint_metadata( |
| 4261 | checkpoint, |
| 4262 | folding_spec=folding_spec, |
| 4263 | family="dense", |
| 4264 | zero_partition_group="dense_dp", |
| 4265 | zero_partition_count=folding_spec.dp_size if folded_autoep_tp else None, |
| 4266 | tp_rank=groups.get_tensor_model_parallel_rank() if folded_autoep_tp else None) |
| 4267 | |
| 4268 | fetch_z3_params = False |
| 4269 | z3_params_to_fetch = None |
| 4270 | autoep_partitioned_experts = False |
| 4271 | allowed_missing_keys = None |
| 4272 | if self.zero_optimization_partition_weights() and not load_optimizer_states and not self.has_moe_layers: |
| 4273 | checkpoint['module'] = get_fp32_state_dict_from_zero_checkpoint(load_dir) |
| 4274 | fetch_z3_params = True |
| 4275 | |
| 4276 | if is_pipe_parallel: |
| 4277 | # Pipeline parallelism uses this to load its own checkpoint files. |
| 4278 | self._curr_ckpt_path = os.path.join(load_dir, tag) |
| 4279 | |
| 4280 | # Universal Checkpoint restores parameters from the zero/ layout, so |
| 4281 | # do not require regular MoE expert checkpoint files in that path. |
| 4282 | if self.has_moe_layers and not self.load_universal_checkpoint(): |
| 4283 | # print(checkpoint.keys()) |
| 4284 | old_moe_load = False |
| 4285 | if not isinstance(checkpoint['num_experts'], list): |
| 4286 | old_moe_load = True |
| 4287 | from deepspeed.checkpoint.constants import AUTOEP_LAYERS_KEY, AUTOEP_LAYERS_KEY_LEGACY |
| 4288 | autoep_layers = checkpoint.get(AUTOEP_LAYERS_KEY) |
| 4289 | if autoep_layers is None: |
| 4290 | autoep_layers = checkpoint.get(AUTOEP_LAYERS_KEY_LEGACY) |
| 4291 | autoep_partitioned_experts = (self.zero_optimization_partition_weights() |
| 4292 | and DeepSpeedEngine._uses_autoep_zero3_partitioned_experts(autoep_layers)) |
no test coverage detected