(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False)
| 5159 | return success |
| 5160 | |
| 5161 | def _save_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False): |
| 5162 | |
| 5163 | save_path = self._get_ckpt_name(save_dir, tag) |
| 5164 | |
| 5165 | zero_optimizer_state = self.zero_optimization() |
| 5166 | |
| 5167 | save_frozen_param = self.zero_optimization_partition_gradients() and not exclude_frozen_parameters |
| 5168 | |
| 5169 | # A hack to save the checkpointing directory. Pipeline parallelism overrides |
| 5170 | # module_state_dict() and uses this path to save the model. module_state_dict() |
| 5171 | # then instead just returns None. The module_state_dict() implementation in |
| 5172 | # PipelineEngine expects the save path to be set in self._curr_ckpt_path. |
| 5173 | self._curr_ckpt_path = os.path.join(save_dir, tag) |
| 5174 | module = self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters) |
| 5175 | self._curr_ckpt_path = None |
| 5176 | |
| 5177 | state = self._common_checkpoint_state(module, zero_optimizer_state, save_frozen_param) |
| 5178 | autotp_uc_info = getattr(self.module, UNIVERSAL_CHECKPOINT_INFO, None) |
| 5179 | if autotp_uc_info is not None: |
| 5180 | state[UNIVERSAL_CHECKPOINT_INFO] = autotp_uc_info |
| 5181 | state.update(client_state) |
| 5182 | log_dist(message=f'Saving model checkpoint: {save_path}', ranks=[0]) |
| 5183 | |
| 5184 | if self.save_non_zero_checkpoint: |
| 5185 | self.checkpoint_engine.save(state_dict=state, path=save_path) |
| 5186 | |
| 5187 | def _get_buffer_names(self): |
| 5188 | buffer_names = [] |
no test coverage detected