| |
| |
|
|
| |
|
|
| import os |
| import torch |
|
|
| from deepspeed.utils import logger |
| from deepspeed.utils.tensor_fragment import map_to_flat_opt_states |
| from deepspeed.runtime.utils import bwc_tensor_model_parallel_rank |
|
|
|
|
| class DeepSpeedOptimizer(object): |
| pass |
|
|
|
|
| class ZeROOptimizer(DeepSpeedOptimizer): |
|
|
| def load_hp_checkpoint_state_from_checkpoint_dir(self, lp_groups_name: str, checkpoint_dir: str) -> None: |
| checkpoint_dir = os.path.join(checkpoint_dir, "zero") |
| optim_state_path = os.path.join(checkpoint_dir, "optimizer_state.pt") |
| assert os.path.isfile( |
| optim_state_path), f'{optim_state_path} containing optimizer global state is missing! Cannot proceed.' |
| optim_sd = torch.load(optim_state_path, weights_only=False) |
|
|
| self._load_global_state(optim_sd) |
|
|
| tp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu) |
| if self.mpu is None: |
| logger.warning("MPU is not provided, setting tp size to 1 in checkpoint loading.") |
| tp_world_size = 1 |
| else: |
| tp_world_size = self.mpu.get_slice_parallel_world_size() if hasattr(self.mpu, "get_slice_parallel_world_size") \ |
| else self.mpu.get_tensor_model_parallel_world_size() |
|
|
| for i, (param_group, |
| loaded_param_group) in enumerate(zip(self.optimizer.param_groups, optim_sd['param_groups'])): |
| |
| opt_keys = set() |
| steps = [] |
|
|
| lp_groups = getattr(self, lp_groups_name) |
| for lp in lp_groups[i]: |
| if lp._hp_mapping is not None: |
| |
| step = lp.load_hp_checkpoint_state(os.path.join(checkpoint_dir, self.param_names[lp]), tp_rank, |
| tp_world_size) |
| for key in lp._hp_mapping.get_optim_state_keys(): |
| opt_keys.add(key) |
| steps.append(step) |
|
|
| hp_param = param_group['params'][0] |
| assert all(step == steps[0] for step in steps), f"Steps {steps} are not equal" |
| if steps[0] is not None: |
| self.optimizer.state[hp_param]['step'] = steps[0] |
|
|
| map_to_flat_opt_states(hp_param, lp_groups[i], self.optimizer.state, opt_keys) |
|
|
| for key, value in loaded_param_group.items(): |
| if key == 'params': |
| continue |
| param_group[key] = value |
|
|
