| 3950 | moe_layer_id += 1 |
| 3951 | |
| 3952 | def load_module_state_dict(self, |
| 3953 | checkpoint, |
| 3954 | strict=True, |
| 3955 | custom_load_fn=None, |
| 3956 | fetch_z3_params=False, |
| 3957 | z3_params_to_fetch=None, |
| 3958 | allowed_missing_keys=None): |
| 3959 | if z3_params_to_fetch is not None: |
| 3960 | params_to_fetch = [ |
| 3961 | p for p in z3_params_to_fetch if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE |
| 3962 | ] |
| 3963 | elif fetch_z3_params: |
| 3964 | params_to_fetch = [ |
| 3965 | p for p in self.module.parameters() |
| 3966 | if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE |
| 3967 | ] |
| 3968 | else: |
| 3969 | params_to_fetch = [] |
| 3970 | |
| 3971 | with deepspeed.zero.GatheredParameters(params_to_fetch, modifier_rank=0): |
| 3972 | module_state_dict = checkpoint['module'] |
| 3973 | if custom_load_fn: |
| 3974 | custom_load_fn(src=module_state_dict, dst=self.module) |
| 3975 | else: |
| 3976 | load_result = self.module.load_state_dict( |
| 3977 | module_state_dict, # TODO |
| 3978 | strict=strict and allowed_missing_keys is None) |
| 3979 | # The expert-key allowance only tightens strict loads; a caller |
| 3980 | # passing strict=False keeps the usual non-strict semantics. |
| 3981 | if strict and allowed_missing_keys is not None: |
| 3982 | missing_keys = set(load_result.missing_keys) |
| 3983 | unexpected_keys = set(load_result.unexpected_keys) |
| 3984 | unexpected_missing = missing_keys - set(allowed_missing_keys) |
| 3985 | if unexpected_missing or unexpected_keys: |
| 3986 | raise RuntimeError("Checkpoint module state did not match the model outside AutoEP expert " |
| 3987 | f"parameters: missing={sorted(unexpected_missing)}, " |
| 3988 | f"unexpected={sorted(unexpected_keys)}") |
| 3989 | |
| 3990 | if checkpoint.get(FROZEN_PARAM_FRAGMENTS, None) is not None: |
| 3991 | saved_frozen_params = checkpoint[FROZEN_PARAM_FRAGMENTS] |
| 3992 | for param in self.module.parameters(): |
| 3993 | if param.requires_grad: |
| 3994 | continue |
| 3995 | if param not in self.param_names: |
| 3996 | raise ValueError(f"failed to find frozen {param} in named params") |
| 3997 | name = self.param_names[param] |
| 3998 | if hasattr(param, 'ds_id'): |
| 3999 | param.ds_tensor.data.copy_(saved_frozen_params[name].data) |
| 4000 | else: |
| 4001 | param.data.copy_(saved_frozen_params[name].data) |
| 4002 | |
| 4003 | def _get_zero_ckpt_prefix(self, dp_rank, bf16_mode): |
| 4004 | return f'{"bf16_" if bf16_mode else ""}zero_pp_rank_{dp_rank}' |