Обучение каузальной языковой модели с нуля

До сих пор мы в основном использовали предварительно обученные модели и осуществляли их дообучение для новых случаев использования, повторно используя веса, полученные в ходе предварительного обучения. Как мы видели в Главе 1, это принято называть трансферным обучением, и это очень успешная стратегия применения моделей Transformer для большинства реальных случаев использования, когда размеченных данных недостаточно. В этой главе мы применим другой подход и обучим совершенно новую модель с нуля. Это хороший подход, если у вас много данных, и они сильно отличаются от данных предварительного обучения, используемых для имеющихся моделей. Однако предварительное обучение языковой модели требует значительно больше вычислительных ресурсов, чем дообучение существующей. Примеры, когда имеет смысл обучать новую модель, включают датасеты, состоящие из музыкальных нот, молекулярных последовательностей, таких как ДНК, или языков программирования. Последние в недавнее время получили широкое распространение благодаря таким инструментам, как TabNine и GitHub’s Copilot, работающим на основе модели Codex от OpenAI, которые могут генерировать длинные последовательности кода. Для решения этой задачи генерации текста лучше всего подходят авторегрессионные или каузальные языковые модели, такие как GPT-2.

В этом разделе мы построим уменьшенную версию модели генерации кода: мы сосредоточимся на однострочных завершениях вместо полных функций или классов, используя подмножество кода Python. Работая с данными на Python, вы часто сталкиваетесь со стеком Data Science на Python, состоящем из библиотек matplotlib, seaborn, pandas и scikit-learn. При использовании этих фреймворков часто возникает необходимость поиска определенных команд, поэтому было бы неплохо, если бы мы могли использовать модель выполненяющую эти вызововы за нас.

В Главе 6 мы создали эффективный токенизатор для обработки исходного кода Python, но нам все еще нужен крупный датасет для предварительного обучения модели. Здесь мы применим наш токенизатор к корпусу кода Python, полученному из розиториев GitHub. Затем мы воспользуемся API Trainer и 🤗 Accelerate для обучения модели. Приступим!

На самом деле это демонстрация модели, которая была обучена и загружена в Hub с помощью кода, приведенного в этом разделе. Вы можете найти ее здесь. Обратите внимание, что поскольку при генерации текста происходит некоторая рандомизация, вы, вероятно, получите немного другой результат.

Сбор данных

Код на Python в изобилии доступен в таких репозиториях кода, как GitHub, и мы можем использовать его для создания датасета путем поиска каждого розитория Python. Именно такой подход был использован в книге Transformers textbook для предварительного обучения большой модели GPT-2. Используя дамп GitHub объемом около 180 ГБ, содержащий примерно 20 миллионов файлов Python под названием codeparrot, авторы создали датасет, которым затем поделились на Hugging Face Hub.

Однако обучение на полном корпусе требует много времени и вычислений, а нам нужно только подмножество датасетов, связанных со стеком data science на Python. Итак, давайте начнем с фильтрации датасета codeparrot по всем файлам, включающим любую из библиотек из этого стека. Из-за большого размера датасета мы хотим избежать его загрузки; вместо этого мы будем использовать функцию потоковой передачи (streaming), чтобы фильтровать его на лету. Чтобы помочь нам отфильтровать примеры кода с использованием библиотек, о которых мы говорили ранее, мы воспользуемся следующей функцией:

def any_keyword_in_string(string, keywords):
    for keyword in keywords:
        if keyword in string:
            return True
    return False

Давайте протестируем ее на двух примерах:

filters = ["pandas", "sklearn", "matplotlib", "seaborn"]
example_1 = "import numpy as np"
example_2 = "import pandas as pd"

print(
    any_keyword_in_string(example_1, filters), any_keyword_in_string(example_2, filters)
)

False True

Мы можем использовать ее для создания функции, которая будет передавать датасет и отфильтровывать нужные нам элементы:

from collections import defaultdict
from tqdm import tqdm
from datasets import Dataset


def filter_streaming_dataset(dataset, filters):
    filtered_dict = defaultdict(list)
    total = 0
    for sample in tqdm(iter(dataset)):
        total += 1
        if any_keyword_in_string(sample["content"], filters):
            for k, v in sample.items():
                filtered_dict[k].append(v)
    print(f"{len(filtered_dict['content'])/total:.2%} of data after filtering.")
    return Dataset.from_dict(filtered_dict)

Затем мы можем просто применить эту функцию к потоковому датасету:

# Выполнение этой ячейки займет очень много времени, поэтому ее следует пропустить и перейти к
# следующей!
from datasets import load_dataset

split = "train"  # "valid"
filters = ["pandas", "sklearn", "matplotlib", "seaborn"]

data = load_dataset(f"transformersbook/codeparrot-{split}", split=split, streaming=True)
filtered_data = filter_streaming_dataset(data, filters)

3.26% of data after filtering.

Таким образом, у нас остается около 3% от исходного датасета, что все равно довольно много - результирующий датасет занимает 6 ГБ и состоит из 600 000 Python-скриптов!

Фильтрация полного датасета может занять 2-3 часа в зависимости от вашей машины и пропускной способности сети. Если вы не хотите выполнять этот длительный процесс самостоятельно, мы предоставляем отфильтрованный датасет на Hub для загрузки:

from datasets import load_dataset, DatasetDict

ds_train = load_dataset("huggingface-course/codeparrot-ds-train", split="train")
ds_valid = load_dataset("huggingface-course/codeparrot-ds-valid", split="validation")

raw_datasets = DatasetDict(
    {
        "train": ds_train,  # .shuffle().select(range(50000)),
        "valid": ds_valid,  # .shuffle().select(range(500))
    }
)

raw_datasets

DatasetDict({
    train: Dataset({
        features: ['repo_name', 'path', 'copies', 'size', 'content', 'license'],
        num_rows: 606720
    })
    valid: Dataset({
        features: ['repo_name', 'path', 'copies', 'size', 'content', 'license'],
        num_rows: 3322
    })
})

Давайте рассмотрим пример из датасета. Мы покажем только первые 200 символов каждого поля:

for key in raw_datasets["train"][0]:
    print(f"{key.upper()}: {raw_datasets['train'][0][key][:200]}")

'REPO_NAME: kmike/scikit-learn'
'PATH: sklearn/utils/__init__.py'
'COPIES: 3'
'SIZE: 10094'
'''CONTENT: """
The :mod:`sklearn.utils` module includes various utilites.
"""

from collections import Sequence

import numpy as np
from scipy.sparse import issparse
import warnings

from .murmurhash import murm
LICENSE: bsd-3-clause'''

Мы видим, что поле content содержит код, на котором мы хотим обучить нашу модель. Теперь, когда у нас есть датасет, нам нужно подготовить тексты, чтобы они были в формате, подходящем для предварительного обучения.

Подготовка датасета

Первым шагом будет токенизация данных, чтобы мы могли использовать их для обучения. Поскольку наша цель - автозаполнение коротких вызовов функций, мы можем оставить размер контекста относительно небольшим. Благодаря этому мы сможем обучать модель гораздо быстрее, и она будет занимать значительно меньше памяти. Если для вашего приложения важно, чтобы контекст был больше (например, если вы хотите, чтобы модель писала юнит-тесты на основе файла с определением функции), обязательно увеличьте это число, но не забывайте, что это приведет к увеличению объема памяти GPU. Пока что давайте зафиксируем размер контекста на 128 токенов, в отличие от 1 024 или 2 048, используемых в GPT-2 или GPT-3 соответственно.

Большинство документов содержит гораздо больше 128 токенов, поэтому простое обрезание входных данных до максимальной длины приведет к тому, что большая часть нашего датасета будет удалена. Вместо этого мы используем параметр return_overflowing_tokens для токенизации всего ввода и разбиения его на части, как мы делали в Главе 6. Мы также будем использовать параметр return_length, чтобы автоматически возвращать длину каждого созданного фрагмента. Часто последний фрагмент будет меньше размера контекста, и мы избавимся от этих фрагментов, чтобы избежать проблем с дополнением; на самом деле они нам не нужны, поскольку у нас и так много данных.

Давайте посмотрим, как это работает, на первых двух примерах:

from transformers import AutoTokenizer

context_length = 128
tokenizer = AutoTokenizer.from_pretrained("huggingface-course/code-search-net-tokenizer")

outputs = tokenizer(
    raw_datasets["train"][:2]["content"],
    truncation=True,
    max_length=context_length,
    return_overflowing_tokens=True,
    return_length=True,
)

print(f"Input IDs length: {len(outputs['input_ids'])}")
print(f"Input chunk lengths: {(outputs['length'])}")
print(f"Chunk mapping: {outputs['overflow_to_sample_mapping']}")

Input IDs length: 34
Input chunk lengths: [128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 117, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 41]
Chunk mapping: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]

Из этих двух примеров видно, что в общей сложности мы получили 34 сегмента. Взглянув на длину фрагментов, мы видим, что фрагменты в конце обоих документов содержат менее 128 токенов (117 и 41, соответственно). Это лишь малая часть всех имеющихся у нас фрагментов, поэтому мы можем смело отбросить их. С помощью поля overflow_to_sample_mapping мы также можем восстановить, какие фрагменты принадлежали каким входным примерам.

В этой операции мы используем удобную особенность функции Dataset.map() из 🤗 Datasets, которая заключается в том, что она не требует отображения один к одному; как мы видели в разделе 3, мы можем создавать батч с большим или меньшим количеством элементов, чем входной батч. Это полезно при выполнении таких операций, как аугментация или фильтрация данных, которые изменяют количество элементов. В нашем случае при токенизации каждого элемента на фрагменты заданного размера контекста мы создаем много примеров из каждого документа. Нам просто нужно убедиться, что мы удалили существующие столбцы, поскольку они имеют противоречивый размер. Если бы мы хотели их сохранить, то могли бы повторить их соответствующим образом и вернуть в рамках вызова Dataset.map():

def tokenize(element):
    outputs = tokenizer(
        element["content"],
        truncation=True,
        max_length=context_length,
        return_overflowing_tokens=True,
        return_length=True,
    )
    input_batch = []
    for length, input_ids in zip(outputs["length"], outputs["input_ids"]):
        if length == context_length:
            input_batch.append(input_ids)
    return {"input_ids": input_batch}


tokenized_datasets = raw_datasets.map(
    tokenize, batched=True, remove_columns=raw_datasets["train"].column_names
)
tokenized_datasets

DatasetDict({
    train: Dataset({
        features: ['input_ids'],
        num_rows: 16702061
    })
    valid: Dataset({
        features: ['input_ids'],
        num_rows: 93164
    })
})

Теперь у нас есть 16,7 миллиона примеров со 128 токенами в каждом, что соответствует примерно 2,1 миллиарда токенов в общей сложности. Для сравнения, модели OpenAI GPT-3 и Codex обучены на 300 и 100 миллиардах токенов соответственно, причем модели Codex инициализируются из контрольных точек GPT-3. Наша цель в этом разделе - не конкурировать с этими моделями, которые могут генерировать длинные, связные тексты, а создать уменьшенную версию, обеспечивающую быструю функцию автозаполнения для специалистов по обработке данных.

Теперь, когда у нас есть готовый датасет, давайте создадим модель!

Инициализация новой модели

Наш первый шаг - инициализация свежей модели GPT-2. Мы будем использовать ту же конфигурацию для нашей модели, что и для маленькой модели GPT-2, поэтому загрузим предварительно обученную конфигурацию, убедимся, что размер токенизатора соответствует размеру словарного запаса модели, и передадим идентификаторы токенов bos и eos (начало и конец последовательности):

from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig

config = AutoConfig.from_pretrained(
    "gpt2",
    vocab_size=len(tokenizer),
    n_ctx=context_length,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
)

С этой конфигурацией мы можем загрузить новую модель. Обратите внимание, что впервые мы не используем функцию from_pretrained(), поскольку фактически инициализируем модель самостоятельно:

model = GPT2LMHeadModel(config)
model_size = sum(t.numel() for t in model.parameters())
print(f"GPT-2 size: {model_size/1000**2:.1f}M parameters")

GPT-2 size: 124.2M parameters

Наша модель имеет 124 миллиона параметров, которые нам предстоит дообучить. Прежде чем начать обучение, нам нужно настроить коллатор данных, который займется созданием батчей. Мы можем использовать коллатор DataCollatorForLanguageModeling, который разработан специально для языкового моделирования (о чем недвусмысленно говорит его название). Помимо стекирования и дополнения батчей, он также заботится о создании меток языковой модели - в каузальном языковом моделировании входы тоже служат метками (просто сдвинутыми на один элемент), и этот коллатор данных создает их на лету во время обучения, так что нам не нужно дублировать input_ids.

Обратите внимание, что DataCollatorForLanguageModeling поддерживает как маскированное языковое моделирование (masked language modeling - MLM), так и каузальное языковое моделирование (causal language modeling - CLM). По умолчанию он подготавливает данные для MLM, но мы можем переключиться на CLM, задав аргумент mlm=False:

from transformers import DataCollatorForLanguageModeling

tokenizer.pad_token = tokenizer.eos_token
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

Давайте посмотрим на пример:

out = data_collator([tokenized_datasets["train"][i] for i in range(5)])
for key in out:
    print(f"{key} shape: {out[key].shape}")

input_ids shape: torch.Size([5, 128])
attention_mask shape: torch.Size([5, 128])
labels shape: torch.Size([5, 128])

Мы видим, что примеры были стекированы, и все тензоры имеют одинаковую форму.

Теперь у нас есть все необходимое для обучения нашей модели - в конце концов, это было не так уж и сложно! Прежде чем приступить к обучению, мы должны войти в Hugging Face. Если вы работаете в блокноте, вы можете сделать это с помощью следующей служебной функции:

from huggingface_hub import notebook_login

notebook_login()

Появится виджет, в котором вы можете ввести свои учетные данные для входа в Hugging Face.

Если вы работаете не в блокноте, просто введите следующую строку в терминале:

huggingface-cli login

Осталось только настроить аргументы обучения и запустить Trainer. Мы будем использовать косинусный график скорости обучения с некоторым разогревом (warmup) и эффективным размером батча в 256 (per_device_train_batch_size * gradient_accumulation_steps). Аккумулирование градиента используется, когда один батч не помещается в память, и инкрементально накапливает градиент за несколько проходов вперед/назад. Мы увидим это в действии, когда создадим цикл обучения с использованием 🤗 Accelerate.

from transformers import Trainer, TrainingArguments

args = TrainingArguments(
    output_dir="codeparrot-ds",
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    evaluation_strategy="steps",
    eval_steps=5_000,
    logging_steps=5_000,
    gradient_accumulation_steps=8,
    num_train_epochs=1,
    weight_decay=0.1,
    warmup_steps=1_000,
    lr_scheduler_type="cosine",
    learning_rate=5e-4,
    save_steps=5_000,
    fp16=True,
    push_to_hub=True,
)

trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=args,
    data_collator=data_collator,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["valid"],
)

Теперь мы можем просто запустить Trainer и дождаться окончания обучения. В зависимости от того, запустите ли вы его на полном или на подмножестве обучающего набора, это займет 20 или 2 часа соответственно, так что захватите несколько чашек кофе и хорошую книгу для чтения!

trainer.train()

После завершения обучения мы можем отправить модель и токенизатор в Hub:

trainer.push_to_hub()

Генерация кода с помощью конвейера

Настал момент истины: давайте посмотрим, насколько хорошо работает обученная модель! В логах мы видим, что потери постоянно снижаются, но чтобы проверить модель на практике, давайте посмотрим, насколько хорошо она работает на некоторых подсказках. Для этого мы обернем модель в pipeline для генерации текста и поместим ее на GPU для быстрой генерации, если таковой доступен:

import torch
from transformers import pipeline

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
pipe = pipeline(
    "text-generation", model="huggingface-course/codeparrot-ds", device=device
)

Давайте начнем с простой задачи - создания диаграммы рассеивания (scatter plot):

txt = """\
# create some data
x = np.random.randn(100)
y = np.random.randn(100)

# create scatter plot with x, y
"""
print(pipe(txt, num_return_sequences=1)[0]["generated_text"])

# create some data
x = np.random.randn(100)
y = np.random.randn(100)

# create scatter plot with x, y
plt.scatter(x, y)

# create scatter

Результат выглядит корректно. Работает ли это также для pandas операции? Давайте посмотрим, сможем ли мы создать DataFrame из двух массивов:

txt = """\
# create some data
x = np.random.randn(100)
y = np.random.randn(100)

# create dataframe from x and y
"""
print(pipe(txt, num_return_sequences=1)[0]["generated_text"])

# create some data
x = np.random.randn(100)
y = np.random.randn(100)

# create dataframe from x and y
df = pd.DataFrame({'x': x, 'y': y})
df.insert(0,'x', x)
for

Отлично, это правильный ответ — хотя затем он снова вставляет столбец x. Поскольку количество генерируемых токенов ограничено, следующий цикл for обрывается. Давайте посмотрим, сможем ли мы сделать что-то более сложное, и чтобы модель помогла нам использовать операцию groupby:

txt = """\
# dataframe with profession, income and name
df = pd.DataFrame({'profession': x, 'income':y, 'name': z})

# calculate the mean income per profession
"""
print(pipe(txt, num_return_sequences=1)[0]["generated_text"])

# dataframe with profession, income and name
df = pd.DataFrame({'profession': x, 'income':y, 'name': z})

# calculate the mean income per profession
profession = df.groupby(['profession']).mean()

# compute the

Неплохо; это правильный способ сделать это. Наконец, давайте посмотрим, сможем ли мы также использовать его для scikit-learn и создать модель Random Forest:

txt = """
# import random forest regressor from scikit-learn
from sklearn.ensemble import RandomForestRegressor

# fit random forest model with 300 estimators on X, y:
"""
print(pipe(txt, num_return_sequences=1)[0]["generated_text"])

# import random forest regressor from scikit-learn
from sklearn.ensemble import RandomForestRegressor

# fit random forest model with 300 estimators on X, y:
rf = RandomForestRegressor(n_estimators=300, random_state=random_state, max_depth=3)
rf.fit(X, y)
rf

Глядя на эти несколько примеров, кажется, что модель усвоила часть синтаксиса стека Python Data Science (конечно, нам нужно будет оценить это более тщательно, прежде чем разворачивать модель в реальном мире). Однако иногда требуется более тщательная настройка обучения модели, чтобы добиться необходимого качества работы для конкретного случая использования. Например, если мы хотим динамически обновлять размер батча или иметь условный цикл обучения, который пропускает плохие примеры на лету? Одним из вариантов может быть подкласс Trainer и добавление необходимых изменений, но иногда проще написать цикл обучения с нуля. Вот тут-то и приходит на помощь 🤗 Accelerate.

Обучение с 🤗 Accelerate

Мы уже видели, как обучать модель с помощью Trainer, который позволяет сделать некоторые настройки. Однако иногда нам нужен полный контроль над циклом обучения, или мы хотим внести некоторые экзотические изменения. В этом случае 🤗 Accelerate - отличный выбор, и в этом разделе мы рассмотрим шаги по его использованию для обучения нашей модели. Чтобы сделать все более интересным, мы также добавим изюминку в цикл обучения.

Поскольку нас в основном интересует разумное автодополнение для библиотек data science, имеет смысл придать больший вес обучающим примерам, в которых чаще используются эти библиотеки. Мы можем легко определить эти примеры по использованию таких ключевых слов, как plt, pd, sk, fit и predict, которые являются наиболее частыми именами для импорта matplotlib.pyplot, pandas и sklearn, а также шаблонам fit/predict для последней. Если каждый из них представлен в виде одного токена, мы можем легко проверить, встречаются ли они во входной последовательности. Токены могут иметь пробельный префикс, поэтому мы также проверим наличие таких версий в словаре токенизатора. Чтобы убедиться, что все работает, мы добавим один тестовый токен, который должен быть разбит на несколько токенов:

keytoken_ids = []
for keyword in [
    "plt",
    "pd",
    "sk",
    "fit",
    "predict",
    " plt",
    " pd",
    " sk",
    " fit",
    " predict",
    "testtest",
]:
    ids = tokenizer([keyword]).input_ids[0]
    if len(ids) == 1:
        keytoken_ids.append(ids[0])
    else:
        print(f"Keyword has not single token: {keyword}")

'Keyword has not single token: testtest'

Отлично, похоже, это прекрасно работает! Теперь мы можем написать пользовательскую функцию потерь, которая принимает входную последовательность, логиты и ключевые токены, которые мы только что выбрали в качестве входных данных. Сначала нам нужно выровнять логиты и входные данные: входная последовательность, сдвинутая на единицу вправо, формирует метки, поскольку следующий токен является меткой для текущего токена. Мы можем добиться этого, начиная метки со второго токена входной последовательности, поскольку модель все равно не делает предсказания для первого токена. Затем мы отсекаем последний логит, поскольку у нас нет метки для токена, который следует за всей входной последовательностью. Таким образом, мы можем вычислить потери для каждого примера и подсчитать количество вхождений всех ключевых слов в каждом примере. Наконец, мы вычисляем средневзвешенное значение по всем примерам, используя вхождения в качестве весов. Поскольку мы не хотим отбрасывать все выборки, в которых нет ключевых слов, мы добавляем 1 к весам:

from torch.nn import CrossEntropyLoss
import torch


def keytoken_weighted_loss(inputs, logits, keytoken_ids, alpha=1.0):
    # Сдвигаем так, чтобы токены < n предсказывали n
    shift_labels = inputs[..., 1:].contiguous()
    shift_logits = logits[..., :-1, :].contiguous()
    # Вычисляем потери на каждый токен
    loss_fct = CrossEntropyLoss(reduce=False)
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
    # Изменение размера и средние потери на каждый пример
    loss_per_sample = loss.view(shift_logits.size(0), shift_logits.size(1)).mean(axis=1)
    # Расчет и масштабирование весов
    weights = torch.stack([(inputs == kt).float() for kt in keytoken_ids]).sum(
        axis=[0, 2]
    )
    weights = alpha * (1.0 + weights)
    # Расчет взвешенного среднего
    weighted_loss = (loss_per_sample * weights).mean()
    return weighted_loss

Прежде чем приступить к обучению с этой потрясающей новой функцией потерь, нам нужно подготовить несколько вещей:

• Нам нужны загрузчики данных, чтобы загружать данные батчами.
• Нам нужно настроить параметры затухания веса (weight decay).
• Время от времени мы хотим проводить оценку, поэтому имеет смысл обернуть код оценки в функцию.

Давайте начнем с загрузчиков данных. Нам нужно только задать для датасета формат "torch", а затем мы можем передать его в PyTorch DataLoader с соответствующим размером батча:

from torch.utils.data.dataloader import DataLoader

tokenized_dataset.set_format("torch")
train_dataloader = DataLoader(tokenized_dataset["train"], batch_size=32, shuffle=True)
eval_dataloader = DataLoader(tokenized_dataset["valid"], batch_size=32)

Далее мы группируем параметры, чтобы оптимизатор знал, какие из них получат дополнительное затухание веса. Обычно все смещения (bias) и весовые коэффициенты LayerNorm (LayerNorm weights) исключаются из этого правила; вот как мы можем это реализовать:

weight_decay = 0.1


def get_grouped_params(model, no_decay=["bias", "LayerNorm.weight"]):
    params_with_wd, params_without_wd = [], []
    for n, p in model.named_parameters():
        if any(nd in n for nd in no_decay):
            params_without_wd.append(p)
        else:
            params_with_wd.append(p)
    return [
        {"params": params_with_wd, "weight_decay": weight_decay},
        {"params": params_without_wd, "weight_decay": 0.0},
    ]

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

def evaluate():
    model.eval()
    losses = []
    for step, batch in enumerate(eval_dataloader):
        with torch.no_grad():
            outputs = model(batch["input_ids"], labels=batch["input_ids"])

        losses.append(accelerator.gather(outputs.loss))
    loss = torch.mean(torch.cat(losses))
    try:
        perplexity = torch.exp(loss)
    except OverflowError:
        perplexity = float("inf")
    return loss.item(), perplexity.item()

С помощью функции evaluate() мы можем сообщать о потерях и перплексии через регулярные промежутки времени. Далее мы переопределим нашу модель, чтобы убедиться, что мы снова обучаемся с нуля:

model = GPT2LMHeadModel(config)

Затем мы можем определить наш оптимизатор, используя предыдущую функцию для части параметров для затухания веса:

from torch.optim import AdamW

optimizer = AdamW(get_grouped_params(model), lr=5e-4)

Теперь давайте подготовим модель, оптимизатор и загрузчики данных, чтобы начать обучение:

from accelerate import Accelerator

accelerator = Accelerator(fp16=True)

model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader
)

Теперь, когда мы отправили наш train_dataloader в accelerator.prepare(), мы можем использовать его длину для вычисления количества шагов обучения. Помните, что это всегда нужно делать после подготовки загрузчика данных, так как этот метод изменит его длину. Мы используем классический линейный график скорости обучения до 0:

from transformers import get_scheduler

num_train_epochs = 1
num_update_steps_per_epoch = len(train_dataloader)
num_training_steps = num_train_epochs * num_update_steps_per_epoch

lr_scheduler = get_scheduler(
    name="linear",
    optimizer=optimizer,
    num_warmup_steps=1_000,
    num_training_steps=num_training_steps,
)

Наконец, чтобы отправить нашу модель в Hub, нам нужно создать объект Repository в рабочей папке. Сначала войдите в Hub Hugging Face, если вы еще не вошли в него. Мы определим имя розитория по идентификатору модели, который мы хотим присвоить нашей модели (не стесняйтесь заменить repo_name на свой собственный вариант; он просто должен содержать ваше имя пользователя, что и делает функция get_full_repo_name()):

from huggingface_hub import Repository, get_full_repo_name

model_name = "codeparrot-ds-accelerate"
repo_name = get_full_repo_name(model_name)
repo_name

'sgugger/codeparrot-ds-accelerate'

Затем мы можем клонировать этот розиторий в локальную папку. Если она уже существует, эта локальная папка должна быть существующим клоном розитория, с которым мы работаем:

output_dir = "codeparrot-ds-accelerate"
repo = Repository(output_dir, clone_from=repo_name)

Теперь мы можем загрузить все, что сохранили в output_dir, вызвав метод repo.push_to_hub(). Это поможет нам загружать промежуточные модели в конце каждой эпохи.

Перед обучением давайте проведем быстрый тест, чтобы проверить, правильно ли работает функция оценки:

evaluate()

(10.934126853942871, 56057.14453125)

Это очень высокие значения для потерь и перплексии, но это неудивительно, ведь мы еще не обучили модель. Итак, у нас все готово для написания основной части скрипта обучения: цикла обучения. В цикле обучения мы выполняем итерации по загрузчику данных и передаем батчи в модель. С помощью логитов мы можем оценить нашу пользовательскую функцию потерь. Мы масштабируем потери по количеству шагов накопления градиента, чтобы не создавать больших потерь при агрегировании большего количества шагов. Перед оптимизацией мы также обрезаем градиенты для лучшей сходимости. Наконец, каждые несколько шагов мы оцениваем модель на оценочном наборе с помощью нашей новой функции evaluate():

from tqdm.notebook import tqdm

gradient_accumulation_steps = 8
eval_steps = 5_000

model.train()
completed_steps = 0
for epoch in range(num_train_epochs):
    for step, batch in tqdm(
        enumerate(train_dataloader, start=1), total=num_training_steps
    ):
        logits = model(batch["input_ids"]).logits
        loss = keytoken_weighted_loss(batch["input_ids"], logits, keytoken_ids)
        if step % 100 == 0:
            accelerator.print(
                {
                    "samples": step * samples_per_step,
                    "steps": completed_steps,
                    "loss/train": loss.item() * gradient_accumulation_steps,
                }
            )
        loss = loss / gradient_accumulation_steps
        accelerator.backward(loss)
        if step % gradient_accumulation_steps == 0:
            accelerator.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
            completed_steps += 1
        if (step % (eval_steps * gradient_accumulation_steps)) == 0:
            eval_loss, perplexity = evaluate()
            accelerator.print({"loss/eval": eval_loss, "perplexity": perplexity})
            model.train()
            accelerator.wait_for_everyone()
            unwrapped_model = accelerator.unwrap_model(model)
            unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
            if accelerator.is_main_process:
                tokenizer.save_pretrained(output_dir)
                repo.push_to_hub(
                    commit_message=f"Training in progress step {step}", blocking=False
                )

Вот и все — теперь у вас есть свой собственный цикл обучения для каузальных языковых моделей, таких как GPT-2, который вы можете дополнительно настроить под свои нужды.