Train the model.
(self)
| 102 | await self.engine.teardown_weight_sync_group() |
| 103 | |
| 104 | async def train(self) -> str: |
| 105 | """Train the model.""" |
| 106 | while self.train_step_num < self.total_steps: |
| 107 | try: |
| 108 | metrics = {} |
| 109 | # sample may be blocked due to explorer does not generate enough data |
| 110 | self.logger.info(f"Sample data for step {self.train_step_num + 1} started.") |
| 111 | sample_task = asyncio.create_task(self._sample_data()) |
| 112 | while not sample_task.done(): |
| 113 | # sync weight to make sure the explorer can continue to explore and generate enough data |
| 114 | if await self.need_sync(): |
| 115 | metrics.update(await self.sync_weight()) |
| 116 | await asyncio.sleep(1) |
| 117 | exps, sample_metrics, repr_samples = await sample_task |
| 118 | metrics.update(sample_metrics) |
| 119 | self.logger.info(f"Sample data for step {self.train_step_num + 1} finished.") |
| 120 | metrics.update(await self.train_step(exps)) |
| 121 | need_sync = await self.need_sync() |
| 122 | need_save = self.need_save() |
| 123 | # For CHECKPOINT sync, save_checkpoint is a superset of |
| 124 | # save_state_dict — skip the latter to avoid redundant writes |
| 125 | # to the same directory. |
| 126 | if need_sync and not (need_save and self.sync_method == SyncMethod.CHECKPOINT): |
| 127 | metrics.update(await self.sync_weight()) |
| 128 | if need_save: |
| 129 | # Only block for the final step (when total_steps is |
| 130 | # finite) to avoid stalling the training loop on earlier |
| 131 | # saves. The HF-model write for save_hf_checkpoint="last" |
| 132 | # is handled unconditionally by the post-loop save below. |
| 133 | is_final_step = self.train_step_num >= self.total_steps |
| 134 | save_as_hf = self.save_hf_checkpoint == "always" or ( |
| 135 | is_final_step and self.save_hf_checkpoint == "last" |
| 136 | ) |
| 137 | metrics.update( |
| 138 | await self.save_checkpoint( |
| 139 | block_until_saved=is_final_step, |
| 140 | save_as_hf=save_as_hf, |
| 141 | ) |
| 142 | ) |
| 143 | if need_sync: |
| 144 | # Update sync bookkeeping even though sync_weight was |
| 145 | # skipped — save_checkpoint already wrote the weights |
| 146 | # and updated latest_state_dict_iteration.txt. |
| 147 | self.last_sync_step = self.train_step_num |
| 148 | self.last_sync_time = time.time() |
| 149 | if self.config.trainer.enable_preview: |
| 150 | self._log_experiences(repr_samples) |
| 151 | self.monitor.log(metrics, self.train_step_num) |
| 152 | except StopAsyncIteration: |
| 153 | self.logger.info("No more samples to train. Stopping training.") |
| 154 | break |
| 155 | except Exception: |
| 156 | self.logger.error(f"Error in Trainer:\n{traceback.format_exc()}") |
| 157 | break |
| 158 | |
| 159 | # Always perform a final save to guarantee: |
| 160 | # 1. A checkpoint exists even when the loop never triggered |
| 161 | # need_save() at the final step. |
no test coverage detected