(self, step_num: Optional[int] = None)
| 181 | await asyncio.gather(*refs) |
| 182 | |
| 183 | async def _checkpoint_weights_update(self, step_num: Optional[int] = None) -> int: |
| 184 | self.logger.info(f"Start to update model weights from checkpoint at step {step_num}.") |
| 185 | if step_num is None: |
| 186 | step_num = await self.synchronizer.get_latest_model_version.remote() |
| 187 | if step_num is None or step_num <= self.model_version: |
| 188 | self.logger.warning( |
| 189 | f"No new checkpoint found for step {step_num}. Current model version: {self.model_version}." |
| 190 | ) |
| 191 | return self.model_version |
| 192 | await asyncio.gather( |
| 193 | *[ |
| 194 | model.sync_model_weights( |
| 195 | step_num, |
| 196 | self.config.synchronizer.sync_method, |
| 197 | timeout=self.config.synchronizer.sync_timeout, |
| 198 | ) |
| 199 | for model in self.models |
| 200 | ] |
| 201 | ) |
| 202 | self.logger.info(f"Model weights updated to checkpoint at step {step_num}.") |
| 203 | return step_num # type: ignore |
| 204 | |
| 205 | async def _pull_latest_weights(self): |
| 206 | self.logger.info("Start to pull latest model weights.") |
no test coverage detected