(self, save_path, tag)
| 5201 | ) |
| 5202 | |
| 5203 | def _save_zero_checkpoint(self, save_path, tag): |
| 5204 | zero_checkpoint_name = self._get_zero_ckpt_name(save_path, tag) |
| 5205 | zero_sd = dict(optimizer_state_dict=self.optimizer.state_dict(), ds_config=self.config, ds_version=version) |
| 5206 | folding_spec = getattr(self, "_autoep_folding_spec", None) |
| 5207 | if folding_spec is not None and folding_spec.tp_size > 1: |
| 5208 | ep_group_name = f"ep_size_{folding_spec.ep_size}" |
| 5209 | zero_sd[FOLDING_METADATA_KEY] = DeepSpeedEngine._make_autoep_folding_metadata( |
| 5210 | folding_spec, |
| 5211 | family="zero_optimizer_state", |
| 5212 | ep_rank=groups._get_expert_parallel_rank(ep_group_name), |
| 5213 | zero_partition_group="per_family", |
| 5214 | zero_partition_rank=dist.get_rank(group=self.optimizer.dp_process_group), |
| 5215 | zero_partition_count=dist.get_world_size(group=self.optimizer.dp_process_group), |
| 5216 | param_families=self._autoep_zero_optimizer_param_families(), |
| 5217 | ) |
| 5218 | self.checkpoint_engine.save(zero_sd, zero_checkpoint_name) |
| 5219 | |
| 5220 | if self.global_rank == 0: |
| 5221 | self._copy_recovery_script(save_path) |
| 5222 | ckpt_type = 'zero' if self.zero_optimization() else 'bf16_zero' |
| 5223 | #logger.info(f'{ckpt_type} checkpoint saved {zero_checkpoint_name}') |
| 5224 | |
| 5225 | def _replace_module_consolidated_state_dict(self): |
| 5226 | """ |
no test coverage detected