從頭開始訓練因果語言模型

到目前為止，我們主要使用預訓練模型並透過重用預訓練權重來微調它們以適應新的用例。正如我們在第 1 章中看到的，這通常被稱為*遷移學習*，它是將 Transformer 模型應用於大多數標記資料稀缺的實際用例的非常成功的策略。在本章中，我們將採用不同的方法，從頭開始訓練一個全新的模型。如果你有大量資料並且它與現有模型使用的預訓練資料非常不同，這是一個很好的方法。然而，預訓練語言模型比僅微調現有模型需要更多的計算資源。訓練新模型可能很有意義的示例包括由音符、DNA 等分子序列或程式語言組成的資料集。後者最近因 TabNine 和 GitHub 的 Copilot 等工具而受到關注，這些工具由 OpenAI 的 Codex 模型提供支援，可以生成長串程式碼。這種文字生成任務最適合使用自迴歸或因果語言模型，例如 GPT-2。

在本節中，我們將構建一個程式碼生成模型的縮小版：我們將使用 Python 程式碼的一個子集，只關注單行補全，而不是完整的函式或類。在使用 Python 處理資料時，您會經常接觸 Python 資料科學堆疊，其中包括 matplotlib、seaborn、pandas 和 scikit-learn 庫。使用這些框架時，通常需要查詢特定命令，因此如果我們能使用模型為我們完成這些呼叫，那將是很好的。

在第 6 章中，我們建立了一個高效的 tokenizer 來處理 Python 原始碼，但我們仍然需要一個大規模資料集來預訓練模型。在這裡，我們將把我們的 tokenizer 應用到一個源自 GitHub 倉庫的 Python 程式碼語料庫中。然後，我們將使用 Trainer API 和 🤗 Accelerate 來訓練模型。我們開始吧！

這實際上展示了使用本節中顯示的程式碼訓練並上傳到 Hub 的模型。您可以在這裡找到它。請注意，由於文字生成中存在一些隨機性，您可能會得到略有不同的結果。

收集資料

Python 程式碼在 GitHub 等程式碼倉庫中非常豐富，我們可以透過抓取每個 Python 倉庫來建立資料集。這是Transformers 教科書中用於預訓練大型 GPT-2 模型的方法。作者使用一個約 180 GB 的 GitHub 轉儲，其中包含大約 2000 萬個名為 codeparrot 的 Python 檔案，構建了一個數據集，然後將其共享到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 個 token，而不是 GPT-2 或 GPT-3 中分別使用的 1,024 或 2,048 個 token。

大多數文件包含的 token 數量遠超過 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 個 token（分別為 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 個 token，總共約 21 億個 token。作為參考，OpenAI 的 GPT-3 和 Codex 模型分別在 300 億和 100 億個 token 上訓練，其中 Codex 模型是從 GPT-3 檢查點初始化的。我們本節的目標不是與這些能夠生成長而連貫文字的模型競爭，而是建立一個縮小版，為資料科學家提供一個快速自動補全功能。

現在我們已經準備好資料集，讓我們設定模型！

✏️ 試試看！ 丟棄所有小於上下文大小的塊在這裡不是一個大問題，因為我們使用的是小上下文視窗。隨著您增加上下文大小（或者如果您有一個短文件語料庫），被丟棄的塊的比例也會增加。一種更有效的資料準備方法是將批次中所有分詞後的樣本與一個 eos_token_id token 連線起來，然後對連線後的序列執行分塊。作為練習，修改 tokenize() 函式以利用這種方法。請注意，您需要設定 truncation=False 並從 tokenizer 中移除其他引數，以獲取完整的 token ID 序列。

初始化新模型

我們的第一步是重新初始化一個 GPT-2 模型。我們將為模型使用與小型 GPT-2 模型相同的配置，因此我們載入預訓練配置，確保 tokenizer 大小與模型詞彙表大小匹配，並傳遞 bos 和 eos（序列開始和結束）token 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()

訓練完成後，我們可以將模型和 tokenizer 推送到 Hub。

trainer.push_to_hub()

✏️ 試試看！ 從原始文字到訓練 GPT-2，我們只用了大約 30 行程式碼，外加 TrainingArguments。嘗試使用您自己的資料集，看看能否獲得好結果！

💡 如果您有多 GPU 機器，請嘗試在那裡執行程式碼。Trainer 會自動管理多臺機器，這可以極大地加快訓練速度。

使用管道進行程式碼生成

現在是檢驗成果的時候了：讓我們看看訓練好的模型實際表現如何！我們可以從日誌中看到損失穩步下降，但為了真正測試模型，讓我們看看它在一些提示上的表現。為此，我們將模型封裝在一個文字生成 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
)

我們從建立一個簡單的散點圖任務開始

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`。由於生成的 token 數量有限，接下來的 `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 最常用的匯入名稱，以及後者的擬合/預測模式。如果這些都表示為單個 token，我們可以輕鬆檢查它們是否出現在輸入序列中。token 可以有空格字首，所以我們也會檢查 tokenizer 詞彙表中這些版本。為了驗證其工作原理，我們將新增一個應該拆分為多個 token 的測試 token。

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 和我們剛剛選擇的關鍵 token 作為輸入。首先，我們需要對齊 logits 和輸入：向右移動一個位置的輸入序列形成標籤，因為下一個 token 是當前 token 的標籤。我們可以透過從輸入序列的第二個 token 開始標籤來實現這一點，因為模型無論如何都不會對第一個 token 進行預測。然後我們截斷最後一個 logit，因為我們沒有完整輸入序列之後 token 的標籤。有了這個，我們可以計算每個樣本的損失並計算每個樣本中所有關鍵詞的出現次數。最後，我們使用出現次數作為權重來計算所有樣本的加權平均值。由於我們不想丟棄所有沒有關鍵詞的樣本，我們給權重加 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_datasets.set_format("torch")
train_dataloader = DataLoader(tokenized_datasets["train"], batch_size=32, shuffle=True)
eval_dataloader = DataLoader(tokenized_datasets["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 等工具記錄重要指標是一個好主意。為訓練迴圈新增適當的日誌記錄，以便您始終可以檢查訓練的進展情況。