Diff Coverage

            "Rank%s Start training. Global step: %s. Train iters: %s. Start epoch: %s. Train epochs: %s.",
            self.local_rank, self.state.global_step, self.train_iters, self.state.epoch, self.train_epochs,
        )

        try:
            start_epoch = self.state.epoch
            for epoch in range(start_epoch, self.train_epochs):
                if epoch != start_epoch:
                    self.train_dataloader.set_epoch(epoch)
                    self.data_iterator = HyperIter(
                        self.train_dataloader, use_background_prefetcher=config.dataloader.use_background_prefetcher
                    )
                self.state.epoch = epoch

                self.on_epoch_begin()

                start_step = self.state.global_step - epoch * self.train_steps
                train_steps = min(self.train_steps, self.train_iters - epoch * self.train_steps)
                for _ in range(start_step, train_steps):
                    try:
                        self.train_step(self.data_iterator)
                    except StopIteration:
                        logger.info(
                            "epoch:%s Dataloader finished with drop_last %s",
                            epoch,
                            config.dataloader.drop_last,
                        )
                        break

                self.on_epoch_end()
                self.state.epoch = epoch + 1

                print_device_mem_info(f"VRAM usage after epoch {epoch + 1}")

                if config.dataloader.use_background_prefetcher:
                    self.data_iterator.stop()

            self.on_train_end()
        finally:
            if config.dataloader.use_background_prefetcher:
                self.data_iterator.stop()

        synchronize()

        self.destroy_distributed()

            while not self.stop_event.is_set():
                try:
                    item = next(self.iterator)
                except StopIteration:
                    self._put_result((StopIteration, None))
                    break

                # A stop request cannot interrupt next(), so discard the item once
                # that call returns instead of retaining another batch and state.
                if self.stop_event.is_set():
                    break

                # Ensure we capture the state so that subsequent dataloader advances
                # don't mutate the captured state in-place. The underlying dataloader's

                # Ensure we capture the state so that subsequent dataloader advances
                # don't mutate the captured state in-place. The underlying dataloader's
                # state_dict() should handle deepcopying if necessary.
                state = self.original_state_dict() if self.original_state_dict else None
                if not self._put_result((item, state)):
                    break
        # The worker must transfer any producer failure back to the training thread.
        except Exception as exc:  # pylint: disable=broad-exception-caught
            self._put_result((exc, None))
        finally:
            # A timed-out stop cannot cancel next(), so the worker must release
            # its own references when that call eventually returns.
            if self.stop_event.is_set():
                self._cleanup_after_stop()

    def _put_result(self, result: tuple[Any, Any]) -> bool:
        """Put a worker result while allowing a stop request to cancel the write."""
        while not self.stop_event.is_set():
            try:
                self.queue.put(result, timeout=0.1)
                return True
            except queue.Full:
                continue
        return False

    def _drain_queue(self) -> None:
        """Release every result currently retained by the prefetch queue."""
        while True:
            try:
                self.queue.get_nowait()
            except queue.Empty:
                return

    def _cleanup_after_stop(self) -> None:
        """Release queued results and live dataloader references after stopping."""
        self._drain_queue()
        self.iterator = None
        self.dataloader = None
        self.original_state_dict = None
        # The final checkpoint may be collected after stop(), so keep the
        # consumed-batch snapshot while releasing live dataloader references.

    def __iter__(self) -> "BackgroundPrefetcher":

        Returns:
            Whether the worker terminated within the requested timeout.
        """
        self.stop_event.set()
        self._drain_queue()
        if self.thread.is_alive():
            self.thread.join(timeout=timeout)
        if self.thread.is_alive():
            logger.warning("BackgroundPrefetcher worker thread did not terminate within timeout.")
            return False

        self._cleanup_after_stop()
        return True


class HyperIter:
    """

        Returns:
            Whether the worker is stopped, or ``True`` when prefetching is disabled.
        """
        if self.use_background_prefetcher and hasattr(self.iterator, "stop"):
            return self.iterator.stop(timeout=timeout)
        return True

    def state_dict(self) -> Dict[str, Any]:
        """Return the underlying dataloader or prefetcher state for checkpointing."""
        if self.use_background_prefetcher and hasattr(self.iterator, "state_dict"):