Repository: oumi-ai/oumi
File: src/oumi/core/trainers/oumi_trainer.py (Python)

Code before:
                elif hasattr(self.train_dataset, "datapipe"):
                    # Hacky way to get examples count from
                    # MapToIterConverterIterDataPipe.
                    # FIXME Remove DataPipes OPE-811
                    num_dataset_examples = len(self.train_dataset.datapipe)  # type: ignore
            except Exception:
                num_dataset_examples = 0

            if num_dataset_examples > 0:
                device_info = get_device_rank_info()
                return int(
                    self.params.num_train_epochs
                    * max(
                        1,
                        math.ceil(
                            float(num_dataset_examples)
                            / (
                                self.params.per_device_train_batch_size
                                * device_info.world_size
                            )
                        ),
                    )
                )

        # Return a positive number (otherwise, LR scheduler may be completely off).
        _DEFAULT_TOTAL_STEPS: Final[int] = 1000
        logger.warning(
            f"Unable to estimate `total_training_steps`, using {_DEFAULT_TOTAL_STEPS}"
        )
        return _DEFAULT_TOTAL_STEPS

    def _set_sampler_epoch(self, epoch: int) -> None:
        """Sets the current epoch on sampler, if it exists and supports it."""
        if self._sampler and hasattr(self._sampler, "set_epoch"):
            self.log(f"Setting sampler epoch to {epoch}.")
            self._sampler.set_epoch(epoch)

    #
    # Handle callbacks
    #
    def _process_callbacks(
        self, event: str, logs: Optional[dict[str, Any]] = None
    ) -> dict[str, Any]:
        """Process callbacks.

        Extremely hacky way to handle HF callbacks.
        Just here to unblock debugging with our MfuCallback
        """
        logs = logs or {}

        for callback in self.callbacks:

Review comment (bug):
I think we should raise an error here asking the user to set `max_steps`. Using 1k will lead to wrong results (this is what the other trainers usually do)

Code after:
                    num_dataset_examples = len(self.train_dataset)  # type: ignore
                elif hasattr(self.train_dataset, "datapipe"):
                    # Hacky way to get examples count from
                    # MapToIterConverterIterDataPipe.
                    # FIXME Remove DataPipes OPE-811
                    num_dataset_examples = len(self.train_dataset.datapipe)  # type: ignore
            except Exception:
                num_dataset_examples = 0

            if num_dataset_examples > 0:
                world_size = get_device_rank_info().world_size
                batch_size = self.params.per_device_train_batch_size
                steps_per_epoch_per_device = math.ceil(
                    float(num_dataset_examples) / (batch_size * world_size)
                )
                return int(num_epochs * max(steps_per_epoch_per_device, 1))

        raise ValueError(
            "Unable to estimate `total_training_steps` "
            + (f"in {num_epochs} epochs" if num_epochs > 0 else "")
            + ". Please define `max_steps` training parameter!"
        )

    def _set_sampler_epoch(self, epoch: int) -> None:
        """Sets the current epoch on sampler, if it exists and supports it."""
        if self._sampler and hasattr(self._sampler, "set_epoch"):
            self.log(f"Setting sampler epoch to {epoch}.")
            self._sampler.set_epoch(epoch)

    #
    # Handle callbacks
    #
    def _process_callbacks(
        self, event: str, logs: Optional[dict[str, Any]] = None
    ) -> dict[str, Any]:
        """Process callbacks.

        Extremely hacky way to handle HF callbacks.
        Just here to unblock debugging with our MfuCallback
        """
        logs = logs or {}

        for callback in self.callbacks:
            if hasattr(callback, event):
                action = getattr(callback, event)
                action(args=self.params, state=None, control=None, logs=logs)

        return logs