Возможности Trainer

Каждый рецепт ниже демонстрирует конкретную возможность Trainer: пользовательские функции потерь, эффективная по памяти оценка, стратегии создания контрольных точек и многое другое.

Создайте issue, если вы хотите видеть здесь определённую функцию или рабочий процесс.

Пользовательская функция потерь

Передайте compute_loss_func в Trainer, чтобы заменить функцию потерь по умолчанию. Функция выполняется после прямого прохода и определяет только способ вычисления потерь по результатам. Чтобы изменить сам прямой проход, вместо этого создайте подкласс и переопределите compute_loss().

Функция потерь, определённая пользователем, должна иметь следующую сигнатуру:

import torch.nn.functional as F

def my_loss_fn(outputs, labels, num_items_in_batch):
    logits = outputs["logits"]
    loss = F.cross_entropy(logits, labels, reduction="sum")
    return loss / num_items_in_batch

• outputs — это необработанный вывод модели (outputs.logits имеет форму (batch, seq_len, vocab_size)).
• labels — это идентификаторы токенов, извлечённые из входной выборки Trainer перед прямым проходом.
• num_items_in_batch — это количество целевых значений предсказания во всей накопленной выборке. Для моделей каузального LM он учитывает сдвинутые метки (labels[..., 1:]), поскольку из-за сдвига меток позиция 0 каждой последовательности не имеет целевого значения. Подробности см. в разделе Масштабирование потерь. Trainer пропускает автоматическую нормализацию потерь, если предоставлена пользовательская функция потерь, поэтому ваша функция должна самостоятельно выполнять нормализацию.

trainer = Trainer(
    model=model,
    args=TrainingArguments(...),
    train_dataset=train_dataset,
    compute_loss_func=my_loss_fn,
)
trainer.train()

См. руководство по наследованию для дополнительных примеров переопределения compute_loss().

Оценка при запуске

Установите eval_on_start=True, чтобы выполнить полный проход оценки перед первым шагом обучения. Оценка до начала обучения позволяет выявить проблемы в конвейере оценки на раннем этапе, особенно при длительных запусках.

eval_on_start требует корректного eval_strategy (например, "epoch") и набора данных для оценки.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        eval_on_start=True,
    ),
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
)
trainer.train()

Полная оценка увеличивает время выполнения, поэтому она наиболее полезна при первом запуске или после изменения compute_metrics.

Оценка с эффективным использованием памяти

Во время оценки Trainer выполняет прямой проход для каждой партии и объединяет логиты в единый тензор на GPU. После полной обработки набора данных для оценки Trainer перемещает объединенные логиты на CPU и вызывает compute_metrics. Для больших моделей или наборов данных для оценки накопленные логиты могут исчерпать память GPU, даже если обучение на том же оборудовании проходит без проблем, поскольку при обучении в памяти одновременно удерживаются активации только одной партии.

eval_accumulation_steps

Раз в n батчей выгружайте накопленные предсказания с GPU на CPU. Меньшие значения снижают потребление памяти GPU за счёт более частых передач на CPU.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        eval_accumulation_steps=16,   # move predictions to CPU every 16 batches
    ),
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
)
trainer.train()

preprocess_logits_for_metrics

Вызывается один раз на каждый батч оценки на GPU, сразу после прямого прохода и до накопления логитов. Возвращаемое значение заменяет логиты в eval_pred.predictions. Выполнение вычислений на уровне батча уменьшает размер тензора на батч и позволяет eval_accumulation_steps выгружать в память тензор меньшего размера.

import evaluate
from transformers import Trainer, TrainingArguments

metric = evaluate.load("accuracy")

def preprocess_logits_for_metrics(logits, labels):
    if isinstance(logits, tuple):
        logits = logits[0]
    return logits.argmax(dim=-1)

def compute_metrics(eval_preds):
    preds, labels = eval_preds
    labels = labels[:, 1:].reshape(-1)
    preds = preds[:, :-1].reshape(-1)
    return metric.compute(predictions=preds, references=labels)

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        eval_accumulation_steps=16,
    ),
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
    preprocess_logits_for_metrics=preprocess_logits_for_metrics,
)
trainer.train()

Производительность загрузчика данных

По умолчанию Trainer создаёт dataloader с параметром dataloader_num_workers=0. Данные загружаются в основном процессе, пока GPU простаивает, что проявляется в низкой загрузке GPU между батчами.

Оба параметра dataloader_persistent_workers и dataloader_prefetch_factor требуют, чтобы dataloader_num_workers > 0.

• dataloader_persistent_workers сохраняет рабочие подпроцессы между эпохами, чтобы избежать повторной инициализации с нуля, ценой повышенного потребления памяти.
• dataloader_prefetch_factor определяет, сколько батчей каждый воркер готовит заранее. При dataloader_prefetch_factor=2 и num_workers=4 в памяти может находиться до 8 батчей, пока GPU обучается на текущем.

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="out",
    dataloader_num_workers=4,            # spawn 4 worker subprocesses
    dataloader_persistent_workers=True,  # keep them alive between epochs
    dataloader_prefetch_factor=2,        # each worker preloads 2 batches ahead
)

Группировка образцов по длине

Используйте train_sampling_strategy="group_by_length", чтобы группировать примеры с близкой длиной и уменьшать количество padding. Если вы не указываете заранее вычисленные длины, Trainer определяет их по первому входному значению модели в каждом элементе набора данных. Это также работает, когда мультимодальные наборы данных на основе процессоров возвращают объекты BatchFeature, поскольку они представляют собой контейнеры признаков, похожие на отображения.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="qwen3-vl-finetuned",
    train_sampling_strategy="group_by_length",
)

Если у Dataset уже есть столбец с заранее вычисленной длиной, Trainer использует именно его. Имя столбца по умолчанию — length. Если в вашем наборе данных используется другое имя, задайте его через length_column_name. Эта стратегия требует набора данных с известной длиной и игнорируется для IterableDataset.

Выборка с перераспределением по батчам

На наборах данных переменной длины дисбаланс внутри микропачки и между устройствами приводит к тому, что устройства тратят время на заполнение (padding) и простаивают на этапах синхронизации градиентов.

Установите train_sampling_strategy="batch_rebalance" в TrainingArguments, чтобы уменьшить оба эффекта. На каждом шаге оптимизатора сэмплер:

1. Сортирует примеры в батче по длине.
2. Разбивает отсортированную партию по устройствам так, чтобы стоимость токенов-заполнителей в каждом микро-пакете была сбалансирована. Микро-пакеты с длинными примерами получают меньше примеров, а микро-пакеты с короткими примерами — больше.

Это уменьшает количество заполнителей (padding) в каждом микробатче, и каждое устройство завершает микробатч примерно в одно и то же время, что сокращает время простоя при синхронизации и позволяет поддерживать пиковое потребление памяти ниже, чем при использовании "group_by_length". На данный момент эта стратегия поддерживается только для обучения с параллелизмом данных (тензорный параллелизм пока не поддерживается).

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        train_sampling_strategy="batch_rebalance",  # balance padding cost across devices
        per_device_train_batch_size=8,              # average samples per micro-batch
        length_column_name="length",                # optional: dataset column with precomputed lengths
    ),
    train_dataset=train_dataset,
)
trainer.train()

per_device_train_batch_size здесь является средним значением, а не точным количеством на каждый шаг: в некоторых микробатчах меньше примеров, в других — больше, но общее число примеров, на которых обучается модель за шаг, остаётся таким же, как при обычной стратегии распределённой выборки.

Сэмплеру требуется длина каждого примера для сортировки и балансировки батчей. По умолчанию Trainer один раз в начале обучения просматривает весь датасет, чтобы вычислить эти длины. Чтобы пропустить эту просканирование, заранее вычислите длины и сохраните их в отдельную колонку датасета (например, на этапе предобработки) и передайте её имя в параметр length_column_name (по умолчанию "length").

NEFTune

NEFTune добавляет случайный шум к эмбеддингам токенов на этапе прямого прохода. Шум выполняет роль регуляризации и может повысить качество при дообучении по инструкциям.

Включите NEFTune, задав neftune_noise_alpha в TrainingArguments. Типичные значения alpha находятся в диапазоне от 5 до 15.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        num_train_epochs=3,
        neftune_noise_alpha=5,
    ),
    train_dataset=train_dataset,
)
trainer.train()

NEFTune влияет только на процесс обучения, а исходный слой эмбеддингов восстанавливается после его завершения.

Логирование

Управляйте моментом и местом записи записей журнала Trainer с помощью параметров logging_strategy, logging_steps и report_to.

• logging_strategy="steps" ведёт журнал после каждого обновления оптимизатора, указанного в logging_steps(). Чтобы вести журнал в конце каждой эпохи, используйте "epoch".
• report_to направляет логи в трекер экспериментов, например Trackio.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        logging_strategy="steps",
        logging_steps=50,               # write a log entry every 50 optimizer updates
        report_to="trackio",            # stream to Trackio (or "wandb", "tensorboard", …)
        run_name="model-experiment-v1", # display name in the tracker
    ),
    train_dataset=train_dataset,
)
trainer.train()

Создание контрольных точек

Trainer сохраняет контрольную точку каждые save_steps() обновлений оптимизатора и хранит все контрольные точки (или последние ~TrainingArguments.save_total_limit).

save_strategy="best" сохраняет только один лучший контрольный пункт по значению метрики. Новый контрольный пункт сохраняется только при улучшении отслеживаемой метрики, что экономит место на диске и предотвращает накопление устаревших контрольных пунктов.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        save_strategy="best",
        metric_for_best_model="perplexity",   # save when eval perplexity improves
        greater_is_better=False,              # lower perplexity is better
        load_best_model_at_end=True,          # load the best weights after training finishes
    ),
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,          # must return {"perplexity": ...}
)
trainer.train()

Возобновление обучения

Передайте resume_from_checkpoint=True в train(), если обучение было прервано и вы хотите продолжить его без потери прогресса. Обучение возобновится с последней контрольной точки в output_dir.

trainer.train(resume_from_checkpoint=True)

Укажите путь к контрольной точке для возобновления с определённого момента.

trainer.train(resume_from_checkpoint="out/checkpoint-1000")

При возобновлении Trainer восстанавливает состояние оптимизатора, состояние планировщика и состояние генератора случайных чисел.

Возобновление с контрольной точки требует наличия файлов состояния оптимизатора и планировщика в каталоге контрольной точки. Если эти файлы отсутствуют (например, при save_only_model=True), оптимизатор начинает работу с нуля.

JIT-чекпоинтинг

При периодическом сохранении контрольных точек (save_strategy=“steps” или “epoch”) вы теряете весь прогресс обучения, накопленный между последней сохранённой контрольной точкой и моментом прерывания. На общих кластерах с прерываемыми нагрузками, такими как Kueue, задания могут быть завершены в любой момент, поэтому этот интервал может означать часы потраченных впустую вычислительных ресурсов.

Функция JIT (Just-In-Time) checkpointing закрывает этот пробел. Когда тренер получает сигнал SIGTERM, он сохраняет контрольную точку в точной точке, где было прервано обучение, что позволяет возобновить процесс с минимальной потерей прогресса. Она работает в связке с периодическим сохранением контрольных точек. Периодические сохранения защищают от сбоев и аппаратных отказов, а JIT-сохранения — от преэмации и корректного завершения работы.

Включите его, установив enable_jit_checkpoint=True в TrainingArguments.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="your-model",
    enable_jit_checkpoint=True,
)

При получении SIGTERM Trainer дожидается завершения текущего шага обучения, сохраняет контрольную точку и корректно останавливает процесс. Файл-маркер (checkpoint-is-incomplete.txt) создаётся в начале сохранения и удаляется после того, как контрольная точка полностью записана. Если в каталоге контрольной точки этот файл всё ещё присутствует, это означает, что сохранение было прервано до завершения. Trainer не проверяет его наличие автоматически, поэтому перед возобновлением обучения проверьте это самостоятельно.

Восстановление из контрольной точки JIT выполняется так же, как и из любой другой контрольной точки.

trainer.train(resume_from_checkpoint=True)

Необходимо настроить оркестратор так, чтобы он выделял достаточно времени для завершения создания контрольной точки. По умолчанию период корректного завершения работы в Kubernetes составляет всего 30 секунд, чего обычно недостаточно для больших моделей.

spec:
  template:
    spec:
      terminationGracePeriodSeconds: 300

#SBATCH --signal=TERM@300

Рассчитайте необходимый период ожидания как сумму максимального времени выполнения шага обучения, времени сохранения контрольной точки и задержки kill_wait в 3 секунды перед началом сохранения контрольной точки. Например, если шаг обучения занимает до 2 минут, а сохранение контрольной точки — 2 минуты, установите период ожидания не менее 243 секунд.