Recherche d'hyperparamètres

Les hyperparamètres tels que le taux d'apprentissage, la taille du lot et le nombre d'époques influencent considérablement les résultats de l'entraînement. Trainer.hyperparameter_search() trouve la meilleure combinaison en exécutant plusieurs essais, chacun avec un ensemble de valeurs différent, et en renvoyant le meilleur.

Chaque essai initialise un nouveau modèle avec model_init, échantillonne de nouveaux hyperparamètres, exécute une boucle d’entraînement complète et signale un objectif au backend de recherche. Le backend utilise chaque objectif pour orienter l’essai suivant. Une fois tous les essais terminés, les meilleurs hyperparamètres sont retournés dans un ~trainer.utils.BestRun.

Initialisation d'un modèle

Commencez chaque essai avec un modèle neuf afin d’éviter l’état des exécutions précédentes. model_init est appelé au début de chaque essai et renvoie une nouvelle instance de modèle, de sorte que chaque essai démarre à partir des mêmes poids initiaux.

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,
)

Ne passez pas model= et model_init= ensemble, sinon Trainer lève une erreur.

Définir l'espace de recherche

Créez une fonction qui définit l’espace de recherche. Le format dépend du backend. Si vous ne définissez pas de fonction hp_space, la recherche par défaut couvre learning_rate, num_train_epochs et 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]},
        },
    }

Lancer la recherche

Fournissez une fonction compute_objective facultative pour définir l'objectif d'optimisation. Par défaut, elle est eval_loss si celle-ci est présente, ou la somme de toutes les valeurs de métriques sinon. Transmettez une fonction explicite pour éviter de vous fier à ce repli. L'arrière-plan de recherche backend optimise l'objectif sur n_trials exécutions dans une direction donnée.

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() renvoie un ~trainer.utils.BestRun contenant la valeur de l'objectif et la meilleure combinaison d'hyperparamètres.

best_run = trainer.hyperparameter_search(...)

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

Appliquez les meilleurs hyperparamètres à TrainingArguments et réentraînez sur l'ensemble de données complet.