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

到目前為止，我們主要使用的是預訓練模型，並透過重用預訓練權重將其微調到新的用例。正如我們在第 1 章中看到的那樣，這通常被稱為遷移學習，它是一種非常成功的策略，可將變壓器模型應用於大多數實際用例，在這些用例中，標記資料稀疏。在本章中，我們將採用不同的方法，從頭開始訓練一個全新的模型。如果你有大量的資料並且它與用於可用模型的預訓練資料非常不同，那麼這是一個很好的方法。但是，與僅微調現有模型相比，預訓練語言模型需要更多的計算資源。可以考慮訓練新模型的示例包括由音樂音符、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 個 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 個片段。檢視塊長度，我們可以看到，兩個文件末尾的塊的 token 數少於 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 個 token，總共約 21 億個 token。作為參考，OpenAI 的 GPT-3 和 Codex 模型分別在 3000 億個和 1000 億個 token 上進行了訓練，其中 Codex 模型是從 GPT-3 檢查點初始化的。我們本節的目標不是與這些模型競爭，這些模型可以生成長而連貫的文字，而是建立一個縮小版本，為資料科學家提供快速的自動補全功能。

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

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

初始化一個新的模型

我們的第一步是全新地初始化一個 GPT-2 模型。我們將使用與小型 GPT-2 模型相同的配置，因此我們將載入預訓練的配置，確保分詞器大小與模型詞彙量大小匹配，並傳遞 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()

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

trainer.push_to_hub()

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

💡 如果您有權訪問具有多個 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最常見的匯入名稱，以及後者的擬合/預測模式。如果這些每個都表示為單個標記，我們可以輕鬆地檢查它們是否出現在輸入序列中。標記可以有空格字首，因此我們還將檢查標記器詞彙表中的這些版本。為了驗證它是否有效，我們將新增一個應被拆分為多個標記的測試標記。

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）的自定義訓練迴圈，您可以根據需要進一步自定義它。

✏️ 試一試！ 您可以建立自己的自定義損失函式，專門針對您的用例，或者在訓練迴圈中新增另一個自定義步驟。