Repository: oumi-ai/oumi
File: src/oumi/core/configs/params/training_params.py (Python)

Code before:
            bf16=self.mixed_precision_dtype == MixedPrecisionDtype.BF16,
            torch_compile=self.compile,
            save_steps=self.save_steps,
            save_strategy=save_strategy,
            logging_first_step=self.logging_first_step,
            torch_empty_cache_steps=self.empty_device_cache_steps,
            resume_from_checkpoint=self.resume_from_checkpoint,
            eval_strategy=self.eval_strategy,
            eval_steps=self.eval_steps,
            dataloader_num_workers=dataloader_num_workers,
            dataloader_prefetch_factor=(
                self.dataloader_prefetch_factor if dataloader_num_workers > 0 else None
            ),
            dataloader_pin_memory=True,  # Set it to True to be explicit.
            ddp_find_unused_parameters=self.ddp_find_unused_parameters,
            max_grad_norm=self.max_grad_norm,  # type: ignore
            dispatch_batches=dispatch_batches,
            # TODO Switch to `accelerator_config` for `dispatch_batches`
            # accelerator_config={  # accelerator config for multi-device training
            #    "split_batches": False,
            #    "dispatch_batches": dispatch_batches,
            #    "even_batches": True,
            #    "use_seedable_sampler": True,
            # },
            seed=self.seed,
            **self.trainer_kwargs,
        )
        assert isinstance(result, transformers.TrainingArguments)
        return result

    def _get_hf_report_to(self) -> list[str]:
        """Gets the list of reporting tools enabled for the current instance.

        Returns:
            list: A list of reporting tools enabled.
                Possible values are "wandb", "tensorboard", or "none".
        """
        report_to = []
        if self.enable_wandb:
            report_to.append("wandb")
        if self.enable_tensorboard:
            report_to.append("tensorboard")
        if len(report_to) == 0:
            report_to.append("none")
        return report_to

    def __post_init__(self):
        """Verifies params."""
        self.run_name = sanitize_run_name(self.run_name)

        if isinstance(self.dataloader_num_workers, str) and not (

Review comment (suggestion):
Let's keep  TODO for data_seed. Ideally, let's file a task to handle data_seed properly, and quote its number here.

Code after:
            bf16=self.mixed_precision_dtype == MixedPrecisionDtype.BF16,
            torch_compile=self.compile,
            save_steps=self.save_steps,
            save_strategy=save_strategy,
            logging_first_step=self.logging_first_step,
            torch_empty_cache_steps=self.empty_device_cache_steps,
            resume_from_checkpoint=self.resume_from_checkpoint,
            eval_strategy=self.eval_strategy,
            eval_steps=self.eval_steps,
            dataloader_num_workers=dataloader_num_workers,
            dataloader_prefetch_factor=(
                self.dataloader_prefetch_factor if dataloader_num_workers > 0 else None
            ),
            dataloader_pin_memory=True,  # Set it to True to be explicit.
            ddp_find_unused_parameters=self.ddp_find_unused_parameters,
            max_grad_norm=self.max_grad_norm,  # type: ignore
            dispatch_batches=dispatch_batches,
            # TODO Switch to `accelerator_config` for `dispatch_batches`
            # accelerator_config={  # accelerator config for multi-device training
            #    "split_batches": False,
            #    "dispatch_batches": dispatch_batches,
            #    "even_batches": True,
            #    "use_seedable_sampler": True,
            # },
            seed=self.seed,
            # TODO: OPE-891 - Support setting a data seed.
            # By default, HF will use the global seed for data loading.
            **self.trainer_kwargs,
        )
        assert isinstance(result, transformers.TrainingArguments)
        return result

    def _get_hf_report_to(self) -> list[str]:
        """Gets the list of reporting tools enabled for the current instance.

        Returns:
            list: A list of reporting tools enabled.
                Possible values are "wandb", "tensorboard", or "none".
        """
        report_to = []
        if self.enable_wandb:
            report_to.append("wandb")
        if self.enable_tensorboard:
            report_to.append("tensorboard")
        if len(report_to) == 0:
            report_to.append("none")
        return report_to

    def __post_init__(self):
        """Verifies params."""
        self.run_name = sanitize_run_name(self.run_name)