TRL documentation
Настройка обучения
Настройка обучения
TRL спроектирован с учётом модульности, чтобы пользователи могли эффективно настраивать цикл обучения под свои задачи. Ниже приведены примеры применения и проверки различных техник.
Хотя в этих примерах используется DPOTrainer, описанные методы кастомизации применимы к большинству (если не ко всем) тренерам в TRL.
Используйте различные оптимизаторы и планировщики
По умолчанию DPOTrainer создаёт оптимизатор torch.optim.AdamW. Вы можете создать и определить другой оптимизатор и передать его в DPOTrainer следующим образом:
from datasets import load_dataset
from torch import optim
from transformers import AutoModelForCausalLM
from trl import DPOTrainer
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
optimizer = optim.SGD(model.parameters(), lr=1e-6)
trainer = DPOTrainer(
model=model,
train_dataset=dataset,
optimizers=(optimizer, None),
)
trainer.train()Добавление планировщика скорости обучения
Вы также можете добавить планировщики скорости обучения, передав и оптимизатор, и планировщик:
from torch import optim
optimizer = optim.AdamW(model.parameters(), lr=1e-6)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
trainer = DPOTrainer(..., optimizers=(optimizer, lr_scheduler))Ссылочные модели 8 бит
Поскольку trl поддерживает все именованные аргументы при загрузке модели из transformers с помощью from_pretrained, вы также можете использовать load_in_8bit из transformers для более эффективного по памяти дообучения.
Подробнее о загрузке моделей в 8-битном формате читайте в разделе Загрузка в 8-битном или 4-битном формате документации transformers.
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
ref_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", quantization_config=quantization_config)
trainer = DPOTrainer(..., ref_model=ref_model)Добавление пользовательских колбэков
Вы можете настроить цикл обучения, добавив колбэки для ведения журнала, мониторинга или раннего завершения. Колбэки позволяют выполнять собственный код в определённых точках процесса обучения.
from transformers import TrainerCallback
class CustomLoggingCallback(TrainerCallback):
def on_log(self, args, state, control, logs=None, **kwargs):
if logs is not None:
print(f"Step {state.global_step}: {logs}")
trainer = DPOTrainer(..., callbacks=[CustomLoggingCallback()])Добавление пользовательских метрик оценки
Вы можете определить собственные метрики оценки для отслеживания во время обучения. Это полезно для мониторинга производительности модели на конкретных задачах.
def compute_metrics(eval_preds):
logits, labels = eval_preds
# Add your metric computation here
return {"custom_metric": 0.0}
training_args = DPOConfig(..., eval_strategy="steps", eval_steps=100)
trainer = DPOTrainer(..., eval_dataset=eval_dataset, compute_metrics=compute_metrics)Использование обучения с смешанной точностью
Обучение с смешанной точностью может значительно ускорить процесс обучения и снизить потребление памяти. Вы можете включить его, установив bf16=True или fp16=True в конфигурации обучения.
# Use bfloat16 precision (recommended for modern GPUs)
training_args = DPOConfig(..., bf16=True)Примечание: используйте bf16=True для GPU Ampere (A100, RTX 30xx) и новее, а fp16=True — для более старых GPU.
Использование накопления градиентов
При обучении с ограниченной памятью GPU накопление градиентов позволяет имитировать более крупные размеры батчей, накапливая градиенты на протяжении нескольких шагов перед обновлением весов.
# Simulate a batch size of 32 with per_device_train_batch_size=4 and gradient_accumulation_steps=8
training_args = DPOConfig(
...,
per_device_train_batch_size=4,
gradient_accumulation_steps=8,
)Изменение цели обучения
Создание подкласса тренера — самый простой способ настроить обучение. Потери определяются в методе compute_loss, поэтому, чтобы обучать модель с другой целевой функцией, создайте подкласс тренера и переопределите этот метод. Подготовка данных, генерация, ведение журнала и сохранение контрольных точек наследуются, поэтому в подклассе остаётся только то, что действительно меняется.
import torch
from trl import DPOTrainer
from trl.trainer.utils import selective_log_softmax
class MyDPOTrainer(DPOTrainer):
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
# In this example: the hinge loss from https://huggingface.co/papers/2309.06657
input_ids, attention_mask = inputs["input_ids"], inputs["attention_mask"]
shift_labels, shift_completion_mask = input_ids[:, 1:], inputs["completion_mask"][:, 1:]
logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
with torch.no_grad():
ref_logits = self.ref_model(input_ids=input_ids, attention_mask=attention_mask).logits
# Sum the log-probs over the completion tokens, for the policy and for the reference model
logps = (selective_log_softmax(logits[:, :-1], shift_labels) * shift_completion_mask).sum(dim=1)
ref_logps = (selective_log_softmax(ref_logits[:, :-1], shift_labels) * shift_completion_mask).sum(dim=1)
chosen_logps, rejected_logps = logps.chunk(2, dim=0) # batch is [chosen, rejected]
ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0)
delta = (chosen_logps - ref_chosen_logps) - (rejected_logps - ref_rejected_logps)
return torch.relu(1 - self.beta * delta).mean()Добавление параметров в конфигурацию
Создайте подкласс конфигурации, чтобы объявить параметры, необходимые для вашей функции потерь:
from dataclasses import dataclass, field
from trl import DPOConfig
@dataclass
class MyDPOConfig(DPOConfig):
my_coef: float = field(default=0.1, metadata={"help": "Coefficient of the custom term."})Изменение формата батча
Если для функции потерь требуются входные данные, которые не создаёт стандартный коллаборатор, передайте собственный коллаборатор и переопределите _prepare_dataset, чтобы не изменять набор данных, если он уже находится в ожидаемом формате:
class MyDPOTrainer(DPOTrainer):
def _prepare_dataset(self, dataset, *args, **kwargs):
return dataset
trainer = MyDPOTrainer(..., data_collator=my_collator)Полный пример
Пример SFT с блок-диффузией examples/sft_diffusion_gemma/sft_diffusion_gemma.py объединяет все три: DiffusionGemmaSFTConfig расширяет SFTConfig параметрами canvas и corruption, а DiffusionGemmaSFTTrainer расширяет SFTTrainer и заменяет авторегрессионную кросс-энтропию на целевую функцию денойзинга блок-диффузии. Ни один из них не требует изменений в библиотеке.
Antidoom — открытый инструмент от Liquid AI, расширяющий DPOTrainer методом Final Token Preference Optimization (Antislop), который применяет функцию потерь предпочтений к позиции одного токена и снижает количество циклических повторов в рассуждающих моделях.
Update on GitHub