Method sync_weight

trinity/trainer/trainer.py:223–254

Sync the model weight.

(self)

        return self.save_interval > 0 and self.train_step_num % self.save_interval == 0

    async def sync_weight(self) -> Dict:
        """Sync the model weight."""
        self.logger.info(f"Trainer sync_weights at step {self.train_step_num} started.")
        metrics = {}
        if self.last_sync_time is not None:
            metrics["time/trainer_sync_interval"] = time.time() - self.last_sync_time
        with Timer(metrics, "time/sync_weight"):
            if self.sync_method == SyncMethod.NCCL:
                result = await self.synchronizer.ready_to_nccl_sync.remote(
                    "trainer", self.train_step_num
                )
                if result is None:
                    self.logger.warning(
                        "NCCL weight sync skipped: Explorer has stopped or is unreachable."
                    )
                else:
                    try:
                        self.engine.sync_weight_nccl()
                    except Exception:
                        self.logger.warning(
                            "NCCL weight sync failed (Explorer may have exited);"
                            f" continuing with stale weights:\n{traceback.format_exc()}"
                        )
            elif self.train_step_num > 0:
                if self.sync_method == SyncMethod.CHECKPOINT:
                    await self.engine.save_state_dict()
                elif self.sync_method == SyncMethod.MEMORY:
                    await self.engine.upload_state_dict()
            self.last_sync_step = self.train_step_num
            self.last_sync_time = time.time()
        self.logger.info(f"Trainer sync_weights at step {self.train_step_num} finished.")
        return metrics

    def _log_experiences(self, samples: List[Dict]) -> None:
        self._sample_exps_to_log.extend(samples)