Búsqueda de hiperparámetros

Los hiperparámetros, como la tasa de aprendizaje, el tamaño del lote y el número de épocas, afectan significativamente los resultados del entrenamiento. Trainer.hyperparameter_search() encuentra la mejor combinación ejecutando múltiples ensayos, cada uno con un conjunto diferente de valores, y devolviendo el mejor.

Cada ensayo inicializa un modelo nuevo con model_init, muestrea nuevos hiperparámetros, ejecuta un bucle de entrenamiento completo e informa un objetivo al backend de búsqueda. El backend utiliza cada objetivo para informar el siguiente ensayo. Una vez que todos los ensayos se completan, los mejores hiperparámetros se devuelven en un ~trainer.utils.BestRun.

Inicialización de un modelo

Inicie cada prueba con un modelo nuevo para evitar el estado de las ejecuciones anteriores. model_init se llama al inicio de cada prueba y devuelve una nueva instancia de modelo, por lo que cada prueba comienza con los mismos pesos iniciales.

from transformers import AutoModelForCausalLM

def model_init(trial):
    return AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B")

trainer = Trainer(
    model_init=model_init,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

No pases model= y model_init= a la vez, o Trainer lanzará un error.

Definir el espacio de búsqueda

Cree una función que defina el espacio de búsqueda. El formato depende del backend. Si no define una función hp_space, la búsqueda predeterminada cubre learning_rate, num_train_epochs y per_device_train_batch_size.

# install one of these hyperparam search backends
pip install optuna
pip install wandb
pip install ray[tune]

def hp_space(trial):
    return {
        "learning_rate": trial.suggest_float("learning_rate", 1e-6, 1e-4, log=True),
        "per_device_train_batch_size": trial.suggest_categorical("per_device_train_batch_size", [16, 32, 64, 128]),
    }

from ray import tune

def hp_space(trial):
    return {
        "learning_rate": tune.loguniform(1e-6, 1e-4),
        "per_device_train_batch_size": tune.choice([16, 32, 64, 128]),
    }

def hp_space(trial):
    return {
        "method": "random",
        "metric": {"name": "objective", "goal": "minimize"},
        "parameters": {
            "learning_rate": {"distribution": "uniform", "min": 1e-6, "max": 1e-4},
            "per_device_train_batch_size": {"values": [16, 32, 64, 128]},
        },
    }

Ejecutar la búsqueda

Proporcione una función opcional compute_objective para definir el objetivo de optimización. Por defecto, se usa eval_loss si está presente, o la suma de todos los valores de métricas en caso contrario. Pase una función explícita para evitar depender de esta alternativa. El backend de búsqueda optimiza el objetivo en n_trials ejecuciones en una direction dada.

def compute_objective(metrics):
    return metrics["eval_loss"]

best_run = trainer.hyperparameter_search(
    hp_space=hp_space,
    compute_objective=compute_objective,
    n_trials=30,               # how many trials to run
    direction="minimize",      # or "maximize" for metrics like accuracy/F1
    backend="optuna",          # "optuna", "ray", or "wandb"
)

hyperparameter_search() devuelve un ~trainer.utils.BestRun que contiene el valor del objetivo y la mejor combinación de hiperparámetros.

best_run = trainer.hyperparameter_search(...)

best_run.objective        # 0.38  (best eval loss)
best_run.hyperparameters  # {"learning_rate": 5e-5, "num_train_epochs": 4, ...}

Aplica los mejores hiperparámetros a TrainingArguments y vuelve a entrenar con el conjunto de datos completo.