(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False)
| 5034 | return success |
| 5035 | |
| 5036 | def _save_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False): |
| 5037 | |
| 5038 | save_path = self._get_ckpt_name(save_dir, tag) |
| 5039 | |
| 5040 | zero_optimizer_state = self.zero_optimization() |
| 5041 | |
| 5042 | save_frozen_param = self.zero_optimization_partition_gradients() and not exclude_frozen_parameters |
| 5043 | |
| 5044 | # A hack to save the checkpointing directory. Pipeline parallelism overrides |
| 5045 | # module_state_dict() and uses this path to save the model. module_state_dict() |
| 5046 | # then instead just returns None. The module_state_dict() implementation in |
| 5047 | # PipelineEngine expects the save path to be set in self._curr_ckpt_path. |
| 5048 | self._curr_ckpt_path = os.path.join(save_dir, tag) |
| 5049 | module = self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters) |
| 5050 | self._curr_ckpt_path = None |
| 5051 | |
| 5052 | state = self._common_checkpoint_state(module, zero_optimizer_state, save_frozen_param) |
| 5053 | autotp_uc_info = getattr(self.module, UNIVERSAL_CHECKPOINT_INFO, None) |
| 5054 | if autotp_uc_info is not None: |
| 5055 | state[UNIVERSAL_CHECKPOINT_INFO] = autotp_uc_info |
| 5056 | state.update(client_state) |
| 5057 | log_dist(message=f'Saving model checkpoint: {save_path}', ranks=[0]) |
| 5058 | |
| 5059 | if self.save_non_zero_checkpoint: |
| 5060 | self.checkpoint_engine.save(state_dict=state, path=save_path) |
| 5061 | |
| 5062 | def _get_buffer_names(self): |
| 5063 | buffer_names = [] |
no test coverage detected