Hugging Face 的徽标

NLP 课程文档

从头开始训练因果语言模型

Hugging Face's logo
加入 Hugging Face 社区

并获得增强文档体验的访问权限

开始使用

Pytorch TensorFlow

从头开始训练因果语言模型

Ask a Question Open In Colab Open In Studio Lab

到目前为止，我们主要使用预训练模型并通过重复使用预训练的权重来微调它们，以便用于新的用例。正如我们在 第 1 章 中看到的，这通常被称为迁移学习，它是一种非常成功的策略，可以将 Transformer 模型应用于大多数实际用例，在这些用例中，标记数据稀少。在本章中，我们将采用不同的方法，从头开始训练一个全新的模型。如果你有大量数据，并且它与用于可用模型的预训练数据非常不同，那么这种方法是比较好的选择。但是，与仅仅微调现有模型相比，从头开始预训练语言模型需要相当多的计算资源。可以考虑训练新模型的示例包括由音乐音符、DNA 等分子序列或编程语言组成的数据集。最近，由于 TabNine 和 GitHub 的 Copilot 等工具的出现，后者获得了发展，这些工具由 OpenAI 的 Codex 模型驱动，可以生成长代码序列。文本生成的任务最适合使用自回归或因果语言模型，例如 GPT-2。

在本节中，我们将构建一个缩减版本的代码生成模型：我们将专注于单行补全而不是完整的函数或类，并使用 Python 代码的子集。在使用 Python 中的数据时，你经常会接触到 Python 数据科学栈，它由 matplotlib、seaborn、pandas 和 scikit-learn 库组成。使用这些框架时，通常需要查找特定命令，因此如果我们可以使用模型来为我们完成这些调用，那就太好了。

在 第 6 章 中，我们创建了一个高效的分词器来处理 Python 源代码，但我们仍然需要一个大型数据集来对其进行预训练。在这里，我们将把我们的分词器应用于从 GitHub 存储库派生的 Python 代码语料库。然后，我们将使用 Trainer API 和 🤗 Accelerate 来训练模型。让我们开始吧！

这实际上展示了使用本节中显示的代码在 Hub 上训练并上传的模型。你可以在 这里 找到它。请注意，由于文本生成中存在一些随机化，你可能会得到略微不同的结果。

收集数据

Python 代码从 GitHub 等代码存储库中广泛可用，我们可以通过抓取每个 Python 存储库来创建数据集。这是在 Transformers 教科书 中使用的方法来预训练大型 GPT-2 模型。使用大约 180 GB 的 GitHub 转储，其中包含大约 2000 万个 Python 文件，称为 codeparrot，作者构建了一个数据集，然后在 Hugging Face Hub 上共享。

但是，在整个语料库上进行训练既耗时又耗费计算资源，我们只需要与 Python 数据科学栈相关的子集数据集。因此，让我们从过滤 codeparrot 数据集开始，以获取包含此栈中任何库的所有文件。由于数据集的大小，我们希望避免下载它；相反，我们将使用流式功能来即时过滤它。为了帮助我们使用前面提到的库过滤代码样本，我们将使用以下函数

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)

然后，我们可以简单地将此函数应用于流式传输的数据集

# This cell will take a very long time to execute, so you should skip it and go to
# the next one!
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 GB，包含 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 个令牌，而不是 GPT-2 或 GPT-3 分别使用的 1,024 或 2,048 个令牌。

大多数文档包含的令牌数量远远超过 128 个，因此简单地将输入截断到最大长度会导致我们数据集的大部分数据被丢弃。相反，我们将使用 `return_overflowing_tokens` 选项来对整个输入进行分词，并将其拆分为多个块，就像我们在 第 6 章 中所做的那样。我们还将使用 `return_length` 选项来自动返回每个创建的块的长度。通常最后一个块的长度会小于上下文大小，我们将丢弃这些片段以避免填充问题；因为我们已经拥有足够多的数据，所以我们并不真的需要它们。

Chunking a large texts in several pieces.

让我们通过查看前两个示例来看看这究竟是如何工作的。

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` 字段，我们也可以重建哪些块属于哪些输入样本。

通过此操作，我们利用了 🤗 Datasets 中 `Dataset.map()` 函数的一个便利特性，即它不需要一对一映射；正如我们在 第 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
    })
})

我们现在有 1670 万个样本，每个样本包含 128 个令牌，总计约 21 亿个令牌。作为参考，OpenAI 的 GPT-3 和 Codex 模型分别在 300 亿个和 100 亿个令牌上进行了训练，其中 Codex 模型从 GPT-3 检查点初始化。本节的目标不是与这些能够生成长篇连贯文本的模型竞争，而是创建一个缩小版本的模型，为数据科学家提供快速自动完成功能。

现在我们已经准备好数据集，让我们开始设置模型吧！

✏️ **试试看！** 由于我们使用的是较小的上下文窗口，所以丢弃所有小于上下文大小的块并没有什么大问题。随着上下文大小的增加（或者如果你有一个包含短文档的语料库），被丢弃的块的比例也会增加。准备数据的更有效方法是将批次中的所有分词后的样本连接在一起，在它们之间使用 `eos_token_id` 令牌，然后对连接后的序列进行分块。作为一个练习，修改 `tokenize()` 函数以使用这种方法。请注意，你需要将 `truncation` 设置为 `False`，并从分词器中删除其他参数以获取完整的令牌 ID 序列。

初始化一个新模型

我们的第一步是全新地初始化一个 GPT-2 模型。我们将使用与小型 GPT-2 模型相同的配置，因此我们加载预训练的配置，确保分词器大小与模型词汇量大小匹配，并传递 `bos` 和 `eos`（序列的开头和结尾）令牌 ID。

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

我们的模型有 1.24 亿个参数需要调整。在开始训练之前，我们需要设置一个数据整理器来负责创建批次。我们可以使用 `DataCollatorForLanguageModeling` 整理器，它专门为语言建模而设计（正如其名称微妙地暗示的那样）。除了堆叠和填充批次之外，它还负责创建语言模型标签——在因果语言建模中，输入也充当标签（只是向后移位一个元素），这个数据整理器在训练期间动态地创建标签，因此我们不需要复制 `input_ids`。

请注意，`DataCollatorForLanguageModeling` 支持掩码语言建模 (MLM) 和因果语言建模 (CLM)。默认情况下，它为 MLM 准备数据，但我们可以通过设置参数 `mlm=False` 来切换到 CLM。

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`。我们将使用余弦学习率调度，并进行一些预热，有效批次大小为 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()

✏️ **试试看！** 我们只用了大约 30 行代码，外加 `TrainingArguments`，就从原始文本到训练 GPT-2。试试你自己的数据集，看看是否能得到好的结果！

💡 如果你有一台拥有多个 GPU 的机器，请尝试在那里运行代码。`Trainer` 会自动管理多台机器，这可以极大地加快训练速度。

使用管道进行代码生成

现在是见证真相的时刻了：让我们看看训练后的模型实际效果如何！我们在日志中可以看到损失一直在稳步下降，但为了对模型进行测试，让我们看看它在一些提示上的表现。为此，我们将模型包装在一个文本生成 `pipeline` 中，并在有 GPU 可用时将其放到 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
)

让我们从创建散点图的简单任务开始。

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` 并设置一个随机森林模型。

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 数据科学库的语法（当然，我们需要更深入地评估它，然后再将其部署到现实世界中）。不过，有时需要对模型训练进行更多定制才能实现特定用例所需的性能。例如，如果我们想动态地更新批次大小，或者有一个条件训练循环可以动态地跳过错误的示例，该怎么办？一种选择是子类化 `Trainer` 并添加必要的更改，但有时直接编写训练循环更简单。这就是 🤗 Accelerate 的用武之地。

使用 🤗 Accelerate 进行训练

我们已经了解了如何使用 `Trainer` 训练模型，它可以进行一些定制。但是，有时我们希望完全控制训练循环，或者我们希望进行一些奇特的更改。在这种情况下，🤗 Accelerate 是一个不错的选择，在本节中，我们将介绍使用它来训练模型的步骤。为了让事情更有趣，我们还将为训练循环添加一个转折。

由于我们主要对数据科学库的合理自动完成功能感兴趣，因此给那些更多地使用这些库的训练样本赋予更多权重是有意义的。我们可以通过使用诸如 `plt`、`pd`、`sk`、`fit` 和 `predict` 之类的关键字来轻松识别这些示例，这些关键字是 `matplotlib.pyplot`、`pandas` 和 `sklearn` 的最常见导入名称，也是后者的拟合/预测模式。如果这些关键字每个都表示一个单独的令牌，那么我们可以轻松地检查它们是否出现在输入序列中。令牌可能带有一个空格前缀，因此我们也会在分词器词汇表中检查这些版本。为了验证它是否有效，我们将添加一个测试令牌，该令牌应该被拆分为多个令牌。

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'

太好了，看起来效果不错！现在我们可以编写一个自定义损失函数，该函数以输入序列、logits 和我们刚刚选择的关键标记作为输入。首先，我们需要对齐 logits 和输入：输入序列向右移动一位形成标签，因为下一个标记是当前标记的标签。我们可以通过从输入序列的第二个标记开始标签来实现这一点，因为模型无论如何都不会对第一个标记进行预测。然后，我们截断最后一个 logit，因为我们没有针对完整输入序列之后标记的标签。有了这些，我们就可以计算每个样本的损失，并统计每个样本中所有关键字出现的次数。最后，我们使用出现次数作为权重，计算所有样本的加权平均值。由于我们不想丢弃所有没有关键字的样本，因此我们在权重中加 1。

from torch.nn import CrossEntropyLoss
import torch


def keytoken_weighted_loss(inputs, logits, keytoken_ids, alpha=1.0):
    # Shift so that tokens < n predict n
    shift_labels = inputs[..., 1:].contiguous()
    shift_logits = logits[..., :-1, :].contiguous()
    # Calculate per-token loss
    loss_fct = CrossEntropyLoss(reduce=False)
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
    # Resize and average loss per sample
    loss_per_sample = loss.view(shift_logits.size(0), shift_logits.size(1)).mean(axis=1)
    # Calculate and scale weighting
    weights = torch.stack([(inputs == kt).float() for kt in keytoken_ids]).sum(
        axis=[0, 2]
    )
    weights = alpha * (1.0 + weights)
    # Calculate weighted average
    weighted_loss = (loss_per_sample * weights).mean()
    return weighted_loss

在我们开始使用这个很棒的新损失函数进行训练之前，我们需要准备几件事。

  • 我们需要数据加载器来批量加载数据。
  • 我们需要设置权重衰减参数。
  • 我们希望不时地进行评估，因此将评估代码封装在函数中是有意义的。

让我们从数据加载器开始。我们只需要将数据集的格式设置为 "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)

接下来，我们将参数分组，以便优化器知道哪些参数将获得额外的权重衰减。通常，所有偏差和 LayerNorm 权重项都不受此影响；以下是我们可以执行此操作的方式。

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
)

🚨 如果你在 TPU 上进行训练，则需要将从上面单元格开始的所有代码移到专用训练函数中。有关更多详细信息，请参见 第 3 章。

现在我们已将 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 对象。首先登录 Hugging Face Hub，如果你还没有登录。我们将从我们想给模型的模型 ID 中确定存储库名称（随意使用你自己的选择替换 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)

现在，我们可以通过调用 repo.push_to_hub() 方法上传在 output_dir 中保存的任何内容。这将有助于我们在每个 epoch 结束时上传中间模型。

在我们进行训练之前，让我们进行快速测试以查看评估功能是否正常工作。

evaluate()
(10.934126853942871, 56057.14453125)

损失和困惑度的值非常高，但这并不奇怪，因为我们还没有训练模型。有了这些，我们就准备好编写训练脚本的核心部分：训练循环。在训练循环中，我们迭代数据加载器并将批次传递给模型。有了 logits，我们就可以评估我们的自定义损失函数。我们根据梯度累积步骤的数量对损失进行缩放，这样在聚合更多步骤时不会创建更大的损失。在我们进行优化之前，我们还会剪裁梯度以获得更好的收敛性。最后，每隔几步，我们就会使用新的 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）的自定义训练循环，你可以根据需要进一步定制它。

✏️ 试试看！ 或者创建你自己的定制损失函数，使其适应你的用例，或者在训练循环中添加另一个定制步骤。

✏️ 试试看！ 在运行长时间的训练实验时，最好使用 TensorBoard 或 Weights & Biases 等工具记录重要指标。在训练循环中添加正确的日志记录，这样你就可以始终检查训练进展情况。

←摘要 问答→