| 4075 | moe_layer_id += 1 |
| 4076 | |
| 4077 | def load_module_state_dict(self, |
| 4078 | checkpoint, |
| 4079 | strict=True, |
| 4080 | custom_load_fn=None, |
| 4081 | fetch_z3_params=False, |
| 4082 | z3_params_to_fetch=None, |
| 4083 | allowed_missing_keys=None): |
| 4084 | if z3_params_to_fetch is not None: |
| 4085 | params_to_fetch = [ |
| 4086 | p for p in z3_params_to_fetch if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE |
| 4087 | ] |
| 4088 | elif fetch_z3_params: |
| 4089 | params_to_fetch = [ |
| 4090 | p for p in self.module.parameters() |
| 4091 | if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE |
| 4092 | ] |
| 4093 | else: |
| 4094 | params_to_fetch = [] |
| 4095 | |
| 4096 | with deepspeed.zero.GatheredParameters(params_to_fetch, modifier_rank=0): |
| 4097 | module_state_dict = checkpoint['module'] |
| 4098 | if custom_load_fn: |
| 4099 | custom_load_fn(src=module_state_dict, dst=self.module) |
| 4100 | else: |
| 4101 | load_result = self.module.load_state_dict( |
| 4102 | module_state_dict, # TODO |
| 4103 | strict=strict and allowed_missing_keys is None) |
| 4104 | # The expert-key allowance only tightens strict loads; a caller |
| 4105 | # passing strict=False keeps the usual non-strict semantics. |
| 4106 | if strict and allowed_missing_keys is not None: |
| 4107 | missing_keys = set(load_result.missing_keys) |
| 4108 | unexpected_keys = set(load_result.unexpected_keys) |
| 4109 | unexpected_missing = missing_keys - set(allowed_missing_keys) |
| 4110 | if unexpected_missing or unexpected_keys: |
| 4111 | raise RuntimeError("Checkpoint module state did not match the model outside AutoEP expert " |
| 4112 | f"parameters: missing={sorted(unexpected_missing)}, " |
| 4113 | f"unexpected={sorted(unexpected_keys)}") |
| 4114 | |
| 4115 | if checkpoint.get(FROZEN_PARAM_FRAGMENTS, None) is not None: |
| 4116 | saved_frozen_params = checkpoint[FROZEN_PARAM_FRAGMENTS] |
| 4117 | for param in self.module.parameters(): |
| 4118 | if param.requires_grad: |
| 4119 | continue |
| 4120 | if param not in self.param_names: |
| 4121 | raise ValueError(f"failed to find frozen {param} in named params") |
| 4122 | name = self.param_names[param] |
| 4123 | if hasattr(param, 'ds_id'): |
| 4124 | param.ds_tensor.data.copy_(saved_frozen_params[name].data) |
| 4125 | else: |
| 4126 | param.data.copy_(saved_frozen_params[name].data) |
| 4127 | |
| 4128 | def _get_zero_ckpt_prefix(self, dp_rank, bf16_mode): |
| 4129 | return f'{"bf16_" if bf16_mode else ""}zero_pp_rank_{dp_rank}' |