| 262 | self._sample_exps_to_log.clear() |
| 263 | |
| 264 | async def save_checkpoint( |
| 265 | self, block_until_saved: bool = False, save_as_hf: bool = False |
| 266 | ) -> Dict: |
| 267 | metrics = {} |
| 268 | with Timer(metrics, "time/save_checkpoint"): |
| 269 | self.logger.info(f"Saving checkpoint at step {self.train_step_num}...") |
| 270 | await self.engine.save_checkpoint( |
| 271 | block_until_saved=block_until_saved, save_as_hf=save_as_hf |
| 272 | ) |
| 273 | self.state.save_trainer( |
| 274 | current_step=self.train_step_num, |
| 275 | sample_strategy_state=self.sample_strategy.state_dict(), |
| 276 | ) |
| 277 | return metrics |
| 278 | |
| 279 | async def shutdown(self) -> None: |
| 280 | self.monitor.close() |