Busca de hiperparâmetros

Hiperparâmetros como taxa de aprendizado, tamanho do lote e número de épocas afetam significativamente os resultados do treinamento. Trainer.hyperparameter_search() encontra a melhor combinação executando múltiplas tentativas, cada uma com um conjunto diferente de valores, e retornando a melhor delas.

Cada tentativa inicializa um novo modelo com model_init, amostra novos hiperparâmetros, executa um loop de treinamento completo e relata um objetivo para o backend de busca. O backend utiliza cada objetivo para orientar a próxima tentativa. Após a conclusão de todas as tentativas, os melhores hiperparâmetros são retornados em um ~trainer.utils.BestRun.

Inicializando um modelo

Inicie cada tentativa com um modelo novo para evitar o estado das execuções anteriores. model_init é chamado no início de cada tentativa e retorna uma nova instância do modelo, de modo que cada tentativa começa a partir dos mesmos pesos iniciais.

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

Não passe model= e model_init= juntos, ou o Trainer lançará um erro.

Defina o espaço de busca

Crie uma função que define o espaço de busca. O formato depende do backend. Se você não definir uma função hp_space, a busca padrão cobre learning_rate, num_train_epochs e 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]},
        },
    }

Executar a busca

Forneça uma função opcional compute_objective para definir o alvo de otimização. O padrão é eval_loss, se presente, ou a soma de todos os valores de métrica, caso contrário. Passe uma função explícita para evitar a dependência desse recurso de fallback. O backend de busca otimiza o objetivo ao longo de n_trials execuções em uma 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() retorna um ~trainer.utils.BestRun contendo o valor do objetivo e a melhor combinação 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, ...}

Aplique os melhores hiperparâmetros a TrainingArguments e treine novamente no conjunto de dados completo.