(self, save_path, tag)
| 4623 | ) |
| 4624 | |
| 4625 | def _save_zero_checkpoint(self, save_path, tag): |
| 4626 | zero_checkpoint_name = self._get_zero_ckpt_name(save_path, tag) |
| 4627 | zero_sd = dict(optimizer_state_dict=self.optimizer.state_dict(), ds_config=self.config, ds_version=version) |
| 4628 | self.checkpoint_engine.save(zero_sd, zero_checkpoint_name) |
| 4629 | |
| 4630 | if self.global_rank == 0: |
| 4631 | self._copy_recovery_script(save_path) |
| 4632 | ckpt_type = 'zero' if self.zero_optimization() else 'bf16_zero' |
| 4633 | #logger.info(f'{ckpt_type} checkpoint saved {zero_checkpoint_name}') |
| 4634 | |
| 4635 | def _replace_module_consolidated_state_dict(self): |
| 4636 | """ |
no test coverage detected