oumi.core.trainers

Core trainers module for the Oumi (Open Universal Machine Intelligence) library.

This module provides various trainer implementations for use in the Oumi framework. These trainers are designed to facilitate the training process for different types of models and tasks.

Example

>>> from oumi.core.trainers import Trainer
>>> trainer = Trainer(model=my_model, dataset=my_dataset)
>>> trainer.train()

Note

For detailed information on each trainer, please refer to their respective

class documentation.

class oumi.core.trainers.BaseTrainer

Bases: ABC

abstractmethod get_last_eval_metrics() → dict[str, Any]

Returns the last evaluation metrics from the trained model.

Returns:

dict with the evaluation metrics from the last trained model step.

abstractmethod save_model(config: TrainingConfig, final: bool = True) → None

Saves the model’s state dictionary to the specified output directory.

Parameters:

• config (TrainingConfig) – The Oumi training config.

• final (bool) – Whether this is the final model being saved during training.

Returns:

None

abstractmethod save_state() → None

Saves the Trainer state.

Under distributed environment this is done only for a process with rank 0.

abstractmethod train(resume_from_checkpoint: str | None = None) → None

Trains a model.

class oumi.core.trainers.HuggingFaceTrainer(hf_trainer: Trainer, processor: BaseProcessor | None = None)

Bases: BaseTrainer

get_last_eval_metrics() → dict

Gets the last evaluation metrics from the trainer.

Returns:

A dictionary containing the last evaluation metrics.

save_model(config: TrainingConfig, final: bool = True) → None

Saves the model’s weights to the specified output directory.

Parameters:

• config – The Oumi training config.

• final – Whether this is the final model being saved during training. - Applies optimizations for the final model checkpoint. - In the case of FSDP, this will always save the FULL_STATE_DICT instead of the default STATE_DICT.

Returns:

None

save_state() → None

See base class.

Saves the Trainer state, since Trainer.save_model saves only the tokenizer with the model.

HuggingFace normally writes state into “trainer_state.json” under output_dir.

train(resume_from_checkpoint: str | None = None) → None

Trains a model.

class oumi.core.trainers.Trainer(model: Module, processing_class: PreTrainedTokenizerBase | None, args: TrainingParams, train_dataset: Dataset, processor: BaseProcessor | None = None, eval_dataset: Dataset | None = None, callbacks: list[TrainerCallback] | None = None, data_collator: Callable | None = None, config: TrainingConfig | None = None, **kwargs)

Bases: BaseTrainer

evaluate() → dict[str, float]

Evaluates the model on the evaluation dataset.

get_last_eval_metrics() → dict[str, Any]

TODO: This should be implemented in the future to enavle Oumi Trainer in the hyper parameters tuning.

log(message: str)

Logs a message if the process is the local process zero.

log_metrics(metrics: dict[str, Any], step: int) → None

Logs metrics to wandb and tensorboard.

save_model(config: TrainingConfig, final: bool = True) → None

Saves the model.

save_state()

Saves the training state.

train(resume_from_checkpoint: str | None = None)

Trains the model.

class oumi.core.trainers.TrlDpoTrainer(*args, **kwargs)

Bases: DPOTrainer

Light wrapper around the DPOTrainer to handle vision models.

class oumi.core.trainers.VerlGrpoTrainer(processing_class: PreTrainedTokenizerBase | None, config: TrainingConfig, reward_funcs: list[Callable], train_dataset: Dataset, eval_dataset: Dataset, processor: BaseProcessor | None = None, cache_dir: str | Path | None = None, **kwargs)

Bases: BaseTrainer

verl GRPO Trainer.

This class wraps verl’s RayPPOTrainer. This class’ name is misleading as it supports other RL algorithms as well, including GRPO, which we use here.

For documentation on the underlying verl RayPPOTrainer, see https://verl.readthedocs.io/en/latest/examples/config.html.

get_last_eval_metrics() → dict[str, Any]

TODO: This should be implemented in the future to enable verl grpo in the hyper parameters tuning loops.

save_model(config: TrainingConfig, final: bool = True) → None

Saves the model.

Parameters:

• config – The Oumi training config.

• final – Whether this is the final model being saved during training.

save_state() → None

Saves the training state.

train() → None

Trains the model using verl’s RayPPOTrainer.