6. Передтренування та завантаження моделей

Reading time: 22 minutes

Генерація тексту

Щоб навчити модель, нам потрібно, щоб ця модель могла генерувати нові токени. Потім ми порівняємо згенеровані токени з очікуваними, щоб навчити модель вчитися генерувати необхідні токени.

Як і в попередніх прикладах, ми вже передбачили деякі токени, тому можливо повторно використати цю функцію для цієї мети.

Оцінка тексту

Щоб виконати правильне навчання, потрібно виміряти прогнози, отримані для очікуваного токена. Мета навчання полягає в максимізації ймовірності правильного токена, що передбачає збільшення його ймовірності відносно інших токенів.

Щоб максимізувати ймовірність правильного токена, потрібно змінити ваги моделі так, щоб ця ймовірність була максимізована. Оновлення ваг здійснюється через зворотне поширення. Це вимагає функцію втрат для максимізації. У цьому випадку функцією буде різниця між виконаним прогнозом і бажаним.

Однак, замість роботи з сирими прогнозами, буде працювати з логарифмом з основою n. Тож, якщо поточний прогноз очікуваного токена становив 7.4541e-05, натуральний логарифм (основа e) 7.4541e-05 приблизно дорівнює -9.5042.
Отже, для кожного запису з довжиною контексту 5 токенів, наприклад, модель повинна передбачити 5 токенів, причому перші 4 токени - це останній з вхідних, а п'ятий - передбачений. Таким чином, для кожного запису ми отримаємо 5 прогнозів у цьому випадку (навіть якщо перші 4 були у вхідних даних, модель цього не знає) з 5 очікуваними токенами і, отже, 5 ймовірностями для максимізації.

Отже, після виконання натурального логарифму для кожного прогнозу, обчислюється середнє, знак мінус видаляється (це називається крос-ентропійні втрати) і це число потрібно зменшити якомога ближче до 0, оскільки натуральний логарифм 1 дорівнює 0:

Інший спосіб виміряти, наскільки хороша модель, називається перплексія. Перплексія - це метрика, що використовується для оцінки того, наскільки добре ймовірнісна модель прогнозує зразок. У мовному моделюванні вона представляє невизначеність моделі при прогнозуванні наступного токена в послідовності.
Наприклад, значення перплексії 48725 означає, що коли потрібно передбачити токен, модель не впевнена, який з 48,725 токенів у словнику є правильним.

Приклад передтренування

Це початковий код, запропонований у https://github.com/rasbt/LLMs-from-scratch/blob/main/ch05/01_main-chapter-code/ch05.ipynb, іноді трохи модифікований

"""
This is code explained before so it won't be exaplained
"""

import tiktoken
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader


class GPTDatasetV1(Dataset):
def __init__(self, txt, tokenizer, max_length, stride):
self.input_ids = []
self.target_ids = []

# Tokenize the entire text
token_ids = tokenizer.encode(txt, allowed_special={"<|endoftext|>"})

# Use a sliding window to chunk the book into overlapping sequences of max_length
for i in range(0, len(token_ids) - max_length, stride):
input_chunk = token_ids[i:i + max_length]
target_chunk = token_ids[i + 1: i + max_length + 1]
self.input_ids.append(torch.tensor(input_chunk))
self.target_ids.append(torch.tensor(target_chunk))

def __len__(self):
return len(self.input_ids)

def __getitem__(self, idx):
return self.input_ids[idx], self.target_ids[idx]


def create_dataloader_v1(txt, batch_size=4, max_length=256,
stride=128, shuffle=True, drop_last=True, num_workers=0):
# Initialize the tokenizer
tokenizer = tiktoken.get_encoding("gpt2")

# Create dataset
dataset = GPTDatasetV1(txt, tokenizer, max_length, stride)

# Create dataloader
dataloader = DataLoader(
dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers)

return dataloader


class MultiHeadAttention(nn.Module):
def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
super().__init__()
assert d_out % num_heads == 0, "d_out must be divisible by n_heads"

self.d_out = d_out
self.num_heads = num_heads
self.head_dim = d_out // num_heads  # Reduce the projection dim to match desired output dim

self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
self.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs
self.dropout = nn.Dropout(dropout)
self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))

def forward(self, x):
b, num_tokens, d_in = x.shape

keys = self.W_key(x)  # Shape: (b, num_tokens, d_out)
queries = self.W_query(x)
values = self.W_value(x)

# We implicitly split the matrix by adding a `num_heads` dimension
# Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
values = values.view(b, num_tokens, self.num_heads, self.head_dim)
queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

# Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
keys = keys.transpose(1, 2)
queries = queries.transpose(1, 2)
values = values.transpose(1, 2)

# Compute scaled dot-product attention (aka self-attention) with a causal mask
attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head

# Original mask truncated to the number of tokens and converted to boolean
mask_bool = self.mask.bool()[:num_tokens, :num_tokens]

# Use the mask to fill attention scores
attn_scores.masked_fill_(mask_bool, -torch.inf)

attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
attn_weights = self.dropout(attn_weights)

# Shape: (b, num_tokens, num_heads, head_dim)
context_vec = (attn_weights @ values).transpose(1, 2)

# Combine heads, where self.d_out = self.num_heads * self.head_dim
context_vec = context_vec.reshape(b, num_tokens, self.d_out)
context_vec = self.out_proj(context_vec)  # optional projection

return context_vec


class LayerNorm(nn.Module):
def __init__(self, emb_dim):
super().__init__()
self.eps = 1e-5
self.scale = nn.Parameter(torch.ones(emb_dim))
self.shift = nn.Parameter(torch.zeros(emb_dim))

def forward(self, x):
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
norm_x = (x - mean) / torch.sqrt(var + self.eps)
return self.scale * norm_x + self.shift


class GELU(nn.Module):
def __init__(self):
super().__init__()

def forward(self, x):
return 0.5 * x * (1 + torch.tanh(
torch.sqrt(torch.tensor(2.0 / torch.pi)) *
(x + 0.044715 * torch.pow(x, 3))
))


class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
GELU(),
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
)

def forward(self, x):
return self.layers(x)


class TransformerBlock(nn.Module):
def __init__(self, cfg):
super().__init__()
self.att = MultiHeadAttention(
d_in=cfg["emb_dim"],
d_out=cfg["emb_dim"],
context_length=cfg["context_length"],
num_heads=cfg["n_heads"],
dropout=cfg["drop_rate"],
qkv_bias=cfg["qkv_bias"])
self.ff = FeedForward(cfg)
self.norm1 = LayerNorm(cfg["emb_dim"])
self.norm2 = LayerNorm(cfg["emb_dim"])
self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

def forward(self, x):
# Shortcut connection for attention block
shortcut = x
x = self.norm1(x)
x = self.att(x)   # Shape [batch_size, num_tokens, emb_size]
x = self.drop_shortcut(x)
x = x + shortcut  # Add the original input back

# Shortcut connection for feed-forward block
shortcut = x
x = self.norm2(x)
x = self.ff(x)
x = self.drop_shortcut(x)
x = x + shortcut  # Add the original input back

return x


class GPTModel(nn.Module):
def __init__(self, cfg):
super().__init__()
self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
self.drop_emb = nn.Dropout(cfg["drop_rate"])

self.trf_blocks = nn.Sequential(
*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])

self.final_norm = LayerNorm(cfg["emb_dim"])
self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

def forward(self, in_idx):
batch_size, seq_len = in_idx.shape
tok_embeds = self.tok_emb(in_idx)
pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
x = self.drop_emb(x)
x = self.trf_blocks(x)
x = self.final_norm(x)
logits = self.out_head(x)
return logits

# Download contents to train the data with
import os
import urllib.request

file_path = "the-verdict.txt"
url = "https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/main/ch02/01_main-chapter-code/the-verdict.txt"

if not os.path.exists(file_path):
with urllib.request.urlopen(url) as response:
text_data = response.read().decode('utf-8')
with open(file_path, "w", encoding="utf-8") as file:
file.write(text_data)
else:
with open(file_path, "r", encoding="utf-8") as file:
text_data = file.read()

total_characters = len(text_data)
tokenizer = tiktoken.get_encoding("gpt2")
total_tokens = len(tokenizer.encode(text_data))

print("Data downloaded")
print("Characters:", total_characters)
print("Tokens:", total_tokens)

# Model initialization
GPT_CONFIG_124M = {
"vocab_size": 50257,   # Vocabulary size
"context_length": 256, # Shortened context length (orig: 1024)
"emb_dim": 768,        # Embedding dimension
"n_heads": 12,         # Number of attention heads
"n_layers": 12,        # Number of layers
"drop_rate": 0.1,      # Dropout rate
"qkv_bias": False      # Query-key-value bias
}

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.eval()
print ("Model initialized")


# Functions to transform from tokens to ids and from to ids to tokens
def text_to_token_ids(text, tokenizer):
encoded = tokenizer.encode(text, allowed_special={'<|endoftext|>'})
encoded_tensor = torch.tensor(encoded).unsqueeze(0) # add batch dimension
return encoded_tensor

def token_ids_to_text(token_ids, tokenizer):
flat = token_ids.squeeze(0) # remove batch dimension
return tokenizer.decode(flat.tolist())



# Define loss functions
def calc_loss_batch(input_batch, target_batch, model, device):
input_batch, target_batch = input_batch.to(device), target_batch.to(device)
logits = model(input_batch)
loss = torch.nn.functional.cross_entropy(logits.flatten(0, 1), target_batch.flatten())
return loss


def calc_loss_loader(data_loader, model, device, num_batches=None):
total_loss = 0.
if len(data_loader) == 0:
return float("nan")
elif num_batches is None:
num_batches = len(data_loader)
else:
# Reduce the number of batches to match the total number of batches in the data loader
# if num_batches exceeds the number of batches in the data loader
num_batches = min(num_batches, len(data_loader))
for i, (input_batch, target_batch) in enumerate(data_loader):
if i < num_batches:
loss = calc_loss_batch(input_batch, target_batch, model, device)
total_loss += loss.item()
else:
break
return total_loss / num_batches


# Apply Train/validation ratio and create dataloaders
train_ratio = 0.90
split_idx = int(train_ratio * len(text_data))
train_data = text_data[:split_idx]
val_data = text_data[split_idx:]

torch.manual_seed(123)

train_loader = create_dataloader_v1(
train_data,
batch_size=2,
max_length=GPT_CONFIG_124M["context_length"],
stride=GPT_CONFIG_124M["context_length"],
drop_last=True,
shuffle=True,
num_workers=0
)

val_loader = create_dataloader_v1(
val_data,
batch_size=2,
max_length=GPT_CONFIG_124M["context_length"],
stride=GPT_CONFIG_124M["context_length"],
drop_last=False,
shuffle=False,
num_workers=0
)


# Sanity checks
if total_tokens * (train_ratio) < GPT_CONFIG_124M["context_length"]:
print("Not enough tokens for the training loader. "
"Try to lower the `GPT_CONFIG_124M['context_length']` or "
"increase the `training_ratio`")

if total_tokens * (1-train_ratio) < GPT_CONFIG_124M["context_length"]:
print("Not enough tokens for the validation loader. "
"Try to lower the `GPT_CONFIG_124M['context_length']` or "
"decrease the `training_ratio`")

print("Train loader:")
for x, y in train_loader:
print(x.shape, y.shape)

print("\nValidation loader:")
for x, y in val_loader:
print(x.shape, y.shape)

train_tokens = 0
for input_batch, target_batch in train_loader:
train_tokens += input_batch.numel()

val_tokens = 0
for input_batch, target_batch in val_loader:
val_tokens += input_batch.numel()

print("Training tokens:", train_tokens)
print("Validation tokens:", val_tokens)
print("All tokens:", train_tokens + val_tokens)


# Indicate the device to use
if torch.cuda.is_available():
device = torch.device("cuda")
elif torch.backends.mps.is_available():
device = torch.device("mps")
else:
device = torch.device("cpu")

print(f"Using {device} device.")

model.to(device) # no assignment model = model.to(device) necessary for nn.Module classes



# Pre-calculate losses without starting yet
torch.manual_seed(123) # For reproducibility due to the shuffling in the data loader

with torch.no_grad(): # Disable gradient tracking for efficiency because we are not training, yet
train_loss = calc_loss_loader(train_loader, model, device)
val_loss = calc_loss_loader(val_loader, model, device)

print("Training loss:", train_loss)
print("Validation loss:", val_loss)


# Functions to train the data
def train_model_simple(model, train_loader, val_loader, optimizer, device, num_epochs,
eval_freq, eval_iter, start_context, tokenizer):
# Initialize lists to track losses and tokens seen
train_losses, val_losses, track_tokens_seen = [], [], []
tokens_seen, global_step = 0, -1

# Main training loop
for epoch in range(num_epochs):
model.train()  # Set model to training mode

for input_batch, target_batch in train_loader:
optimizer.zero_grad() # Reset loss gradients from previous batch iteration
loss = calc_loss_batch(input_batch, target_batch, model, device)
loss.backward() # Calculate loss gradients
optimizer.step() # Update model weights using loss gradients
tokens_seen += input_batch.numel()
global_step += 1

# Optional evaluation step
if global_step % eval_freq == 0:
train_loss, val_loss = evaluate_model(
model, train_loader, val_loader, device, eval_iter)
train_losses.append(train_loss)
val_losses.append(val_loss)
track_tokens_seen.append(tokens_seen)
print(f"Ep {epoch+1} (Step {global_step:06d}): "
f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")

# Print a sample text after each epoch
generate_and_print_sample(
model, tokenizer, device, start_context
)

return train_losses, val_losses, track_tokens_seen


def evaluate_model(model, train_loader, val_loader, device, eval_iter):
model.eval()
with torch.no_grad():
train_loss = calc_loss_loader(train_loader, model, device, num_batches=eval_iter)
val_loss = calc_loss_loader(val_loader, model, device, num_batches=eval_iter)
model.train()
return train_loss, val_loss


def generate_and_print_sample(model, tokenizer, device, start_context):
model.eval()
context_size = model.pos_emb.weight.shape[0]
encoded = text_to_token_ids(start_context, tokenizer).to(device)
with torch.no_grad():
token_ids = generate_text(
model=model, idx=encoded,
max_new_tokens=50, context_size=context_size
)
decoded_text = token_ids_to_text(token_ids, tokenizer)
print(decoded_text.replace("\n", " "))  # Compact print format
model.train()


# Start training!
import time
start_time = time.time()

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.0004, weight_decay=0.1)

num_epochs = 10
train_losses, val_losses, tokens_seen = train_model_simple(
model, train_loader, val_loader, optimizer, device,
num_epochs=num_epochs, eval_freq=5, eval_iter=5,
start_context="Every effort moves you", tokenizer=tokenizer
)

end_time = time.time()
execution_time_minutes = (end_time - start_time) / 60
print(f"Training completed in {execution_time_minutes:.2f} minutes.")



# Show graphics with the training process
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
import math
def plot_losses(epochs_seen, tokens_seen, train_losses, val_losses):
fig, ax1 = plt.subplots(figsize=(5, 3))
ax1.plot(epochs_seen, train_losses, label="Training loss")
ax1.plot(
epochs_seen, val_losses, linestyle="-.", label="Validation loss"
)
ax1.set_xlabel("Epochs")
ax1.set_ylabel("Loss")
ax1.legend(loc="upper right")
ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
ax2 = ax1.twiny()
ax2.plot(tokens_seen, train_losses, alpha=0)
ax2.set_xlabel("Tokens seen")
fig.tight_layout()
plt.show()

# Compute perplexity from the loss values
train_ppls = [math.exp(loss) for loss in train_losses]
val_ppls = [math.exp(loss) for loss in val_losses]
# Plot perplexity over tokens seen
plt.figure()
plt.plot(tokens_seen, train_ppls, label='Training Perplexity')
plt.plot(tokens_seen, val_ppls, label='Validation Perplexity')
plt.xlabel('Tokens Seen')
plt.ylabel('Perplexity')
plt.title('Perplexity over Training')
plt.legend()
plt.show()

epochs_tensor = torch.linspace(0, num_epochs, len(train_losses))
plot_losses(epochs_tensor, tokens_seen, train_losses, val_losses)


torch.save({
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
},
"/tmp/model_and_optimizer.pth"
)

Давайте розглянемо пояснення крок за кроком

Функції для перетворення тексту <--> id

Це деякі прості функції, які можна використовувати для перетворення текстів з словника в id та навпаки. Це необхідно на початку обробки тексту та в кінці прогнозів:

# Functions to transform from tokens to ids and from to ids to tokens
def text_to_token_ids(text, tokenizer):
encoded = tokenizer.encode(text, allowed_special={'<|endoftext|>'})
encoded_tensor = torch.tensor(encoded).unsqueeze(0) # add batch dimension
return encoded_tensor

def token_ids_to_text(token_ids, tokenizer):
flat = token_ids.squeeze(0) # remove batch dimension
return tokenizer.decode(flat.tolist())

Генерація текстових функцій

У попередньому розділі була функція, яка просто отримувала найбільш ймовірний токен після отримання логітів. Однак це означатиме, що для кожного введення завжди буде генеруватися один і той же вихід, що робить його дуже детермінованим.

Наступна функція generate_text застосовуватиме концепції top-k, temperature та multinomial.

• top-k означає, що ми почнемо зменшувати до -inf всі ймовірності всіх токенів, за винятком топ k токенів. Отже, якщо k=3, перед прийняттям рішення лише 3 найімовірніші токени матимуть ймовірність, відмінну від -inf.
• temperature означає, що кожна ймовірність буде поділена на значення температури. Значення 0.1 покращить найвищу ймовірність у порівнянні з найнижчою, тоді як температура 5, наприклад, зробить її більш рівномірною. Це допомагає покращити варіацію у відповідях, яку ми хотіли б, щоб LLM мала.
• Після застосування температури знову застосовується функція softmax, щоб усі залишкові токени мали загальну ймовірність 1.
• Нарешті, замість вибору токена з найбільшою ймовірністю, функція multinomial застосовується для прогнозування наступного токена відповідно до фінальних ймовірностей. Отже, якщо токен 1 мав 70% ймовірності, токен 2 - 20%, а токен 3 - 10%, 70% часу буде обрано токен 1, 20% часу - токен 2, а 10% часу - токен 3.

# Generate text function
def generate_text(model, idx, max_new_tokens, context_size, temperature=0.0, top_k=None, eos_id=None):

# For-loop is the same as before: Get logits, and only focus on last time step
for _ in range(max_new_tokens):
idx_cond = idx[:, -context_size:]
with torch.no_grad():
logits = model(idx_cond)
logits = logits[:, -1, :]

# New: Filter logits with top_k sampling
if top_k is not None:
# Keep only top_k values
top_logits, _ = torch.topk(logits, top_k)
min_val = top_logits[:, -1]
logits = torch.where(logits < min_val, torch.tensor(float("-inf")).to(logits.device), logits)

# New: Apply temperature scaling
if temperature > 0.0:
logits = logits / temperature

# Apply softmax to get probabilities
probs = torch.softmax(logits, dim=-1)  # (batch_size, context_len)

# Sample from the distribution
idx_next = torch.multinomial(probs, num_samples=1)  # (batch_size, 1)

# Otherwise same as before: get idx of the vocab entry with the highest logits value
else:
idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # (batch_size, 1)

if idx_next == eos_id:  # Stop generating early if end-of-sequence token is encountered and eos_id is specified
break

# Same as before: append sampled index to the running sequence
idx = torch.cat((idx, idx_next), dim=1)  # (batch_size, num_tokens+1)

return idx

Loss functions

Функція calc_loss_batch обчислює крос-ентропію прогнозу одного пакета.
Функція calc_loss_loader отримує крос-ентропію всіх пакетів і обчислює середню крос-ентропію.

# Define loss functions
def calc_loss_batch(input_batch, target_batch, model, device):
input_batch, target_batch = input_batch.to(device), target_batch.to(device)
logits = model(input_batch)
loss = torch.nn.functional.cross_entropy(logits.flatten(0, 1), target_batch.flatten())
return loss

def calc_loss_loader(data_loader, model, device, num_batches=None):
total_loss = 0.
if len(data_loader) == 0:
return float("nan")
elif num_batches is None:
num_batches = len(data_loader)
else:
# Reduce the number of batches to match the total number of batches in the data loader
# if num_batches exceeds the number of batches in the data loader
num_batches = min(num_batches, len(data_loader))
for i, (input_batch, target_batch) in enumerate(data_loader):
if i < num_batches:
loss = calc_loss_batch(input_batch, target_batch, model, device)
total_loss += loss.item()
else:
break
return total_loss / num_batches

Завантаження даних

Функції create_dataloader_v1 та create_dataloader_v1 вже обговорювалися в попередньому розділі.

З цього моменту зверніть увагу, що визначено, що 90% тексту буде використано для навчання, тоді як 10% буде використано для валідації, і обидва набори зберігаються в 2 різних завантажувачах даних.
Зверніть увагу, що іноді частина набору даних також залишається для тестового набору, щоб краще оцінити продуктивність моделі.

Обидва завантажувачі даних використовують один і той же розмір партії, максимальну довжину, крок і кількість робітників (0 у цьому випадку).
Основні відмінності полягають у даних, що використовуються кожним, а валідатор не відкидає останній і не перемішує дані, оскільки це не потрібно для цілей валідації.

Також той факт, що крок такий же великий, як довжина контексту, означає, що не буде перекриття між контекстами, що використовуються для навчання даних (зменшує перенавчання, але також і навчальний набір даних).

Більше того, зверніть увагу, що розмір партії в цьому випадку становить 2, щоб розділити дані на 2 партії, основна мета цього - дозволити паралельну обробку та зменшити споживання на партію.

train_ratio = 0.90
split_idx = int(train_ratio * len(text_data))
train_data = text_data[:split_idx]
val_data = text_data[split_idx:]

torch.manual_seed(123)

train_loader = create_dataloader_v1(
train_data,
batch_size=2,
max_length=GPT_CONFIG_124M["context_length"],
stride=GPT_CONFIG_124M["context_length"],
drop_last=True,
shuffle=True,
num_workers=0
)

val_loader = create_dataloader_v1(
val_data,
batch_size=2,
max_length=GPT_CONFIG_124M["context_length"],
stride=GPT_CONFIG_124M["context_length"],
drop_last=False,
shuffle=False,
num_workers=0
)

Перевірка цілісності

Мета полягає в тому, щоб перевірити, чи є достатня кількість токенів для навчання, чи відповідають форми очікуваним, і отримати деяку інформацію про кількість токенів, використаних для навчання та валідації:

# Sanity checks
if total_tokens * (train_ratio) < GPT_CONFIG_124M["context_length"]:
print("Not enough tokens for the training loader. "
"Try to lower the `GPT_CONFIG_124M['context_length']` or "
"increase the `training_ratio`")

if total_tokens * (1-train_ratio) < GPT_CONFIG_124M["context_length"]:
print("Not enough tokens for the validation loader. "
"Try to lower the `GPT_CONFIG_124M['context_length']` or "
"decrease the `training_ratio`")

print("Train loader:")
for x, y in train_loader:
print(x.shape, y.shape)

print("\nValidation loader:")
for x, y in val_loader:
print(x.shape, y.shape)

train_tokens = 0
for input_batch, target_batch in train_loader:
train_tokens += input_batch.numel()

val_tokens = 0
for input_batch, target_batch in val_loader:
val_tokens += input_batch.numel()

print("Training tokens:", train_tokens)
print("Validation tokens:", val_tokens)
print("All tokens:", train_tokens + val_tokens)

Вибір пристрою для навчання та попередніх розрахунків

Наступний код просто вибирає пристрій для використання та обчислює втрати навчання та втрати валідації (без попереднього навчання) як відправну точку.

# Indicate the device to use

if torch.cuda.is_available():
device = torch.device("cuda")
elif torch.backends.mps.is_available():
device = torch.device("mps")
else:
device = torch.device("cpu")

print(f"Using {device} device.")

model.to(device) # no assignment model = model.to(device) necessary for nn.Module classes

# Pre-calculate losses without starting yet
torch.manual_seed(123) # For reproducibility due to the shuffling in the data loader

with torch.no_grad(): # Disable gradient tracking for efficiency because we are not training, yet
train_loss = calc_loss_loader(train_loader, model, device)
val_loss = calc_loss_loader(val_loader, model, device)

print("Training loss:", train_loss)
print("Validation loss:", val_loss)

Функції навчання

Функція generate_and_print_sample просто отримує контекст і генерує деякі токени, щоб зрозуміти, наскільки добре модель на даний момент. Це викликається функцією train_model_simple на кожному кроці.

Функція evaluate_model викликається так часто, як вказано в функції навчання, і використовується для вимірювання втрат під час навчання та втрат валідації на даний момент у навчанні моделі.

Тоді велика функція train_model_simple є тією, яка фактично навчає модель. Вона очікує:

• Завантажувач навчальних даних (з даними, вже розділеними та підготовленими для навчання)
• Завантажувач валідатора
• Оптимізатор, який буде використовуватися під час навчання: Це функція, яка використовуватиме градієнти та оновлюватиме параметри, щоб зменшити втрати. У цьому випадку, як ви побачите, використовується AdamW, але є багато інших.
• optimizer.zero_grad() викликається для скидання градієнтів на кожному раунді, щоб не накопичувати їх.
• Параметр lr є швидкістю навчання, яка визначає розмір кроків, які робляться під час процесу оптимізації при оновленні параметрів моделі. Менша швидкість навчання означає, що оптимізатор виконує менші оновлення ваг, що може призвести до більш точної конвергенції, але може уповільнити навчання. Більша швидкість навчання може прискорити навчання, але ризикує перепригнути мінімум функції втрат (перестрибнути точку, де функція втрат мінімізується).
• Зниження ваг модифікує крок Обчислення втрат, додаючи додатковий термін, який штрафує великі ваги. Це заохочує оптимізатор знаходити рішення з меншими вагами, балансуючи між хорошим підходом до даних і збереженням моделі простою, запобігаючи перенавчанню в моделях машинного навчання, заважаючи моделі надавати занадто велике значення будь-якій окремій ознаці.
• Традиційні оптимізатори, такі як SGD з L2 регуляризацією, поєднують зниження ваг з градієнтом функції втрат. Однак AdamW (варіант оптимізатора Adam) розділяє зниження ваг від оновлення градієнта, що призводить до більш ефективної регуляризації.
• Пристрій, який використовувати для навчання
• Кількість епох: Кількість разів, щоб пройти через навчальні дані
• Частота оцінки: Частота виклику evaluate_model
• Ітерація оцінки: Кількість пакетів, які використовуються при оцінці поточного стану моделі під час виклику generate_and_print_sample
• Початковий контекст: Яке початкове речення використовувати при виклику generate_and_print_sample
• Токенізатор

# Functions to train the data
def train_model_simple(model, train_loader, val_loader, optimizer, device, num_epochs,
eval_freq, eval_iter, start_context, tokenizer):
# Initialize lists to track losses and tokens seen
train_losses, val_losses, track_tokens_seen = [], [], []
tokens_seen, global_step = 0, -1

# Main training loop
for epoch in range(num_epochs):
model.train()  # Set model to training mode

for input_batch, target_batch in train_loader:
optimizer.zero_grad() # Reset loss gradients from previous batch iteration
loss = calc_loss_batch(input_batch, target_batch, model, device)
loss.backward() # Calculate loss gradients
optimizer.step() # Update model weights using loss gradients
tokens_seen += input_batch.numel()
global_step += 1

# Optional evaluation step
if global_step % eval_freq == 0:
train_loss, val_loss = evaluate_model(
model, train_loader, val_loader, device, eval_iter)
train_losses.append(train_loss)
val_losses.append(val_loss)
track_tokens_seen.append(tokens_seen)
print(f"Ep {epoch+1} (Step {global_step:06d}): "
f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")

# Print a sample text after each epoch
generate_and_print_sample(
model, tokenizer, device, start_context
)

return train_losses, val_losses, track_tokens_seen


def evaluate_model(model, train_loader, val_loader, device, eval_iter):
model.eval() # Set in eval mode to avoid dropout
with torch.no_grad():
train_loss = calc_loss_loader(train_loader, model, device, num_batches=eval_iter)
val_loss = calc_loss_loader(val_loader, model, device, num_batches=eval_iter)
model.train() # Back to training model applying all the configurations
return train_loss, val_loss


def generate_and_print_sample(model, tokenizer, device, start_context):
model.eval() # Set in eval mode to avoid dropout
context_size = model.pos_emb.weight.shape[0]
encoded = text_to_token_ids(start_context, tokenizer).to(device)
with torch.no_grad():
token_ids = generate_text(
model=model, idx=encoded,
max_new_tokens=50, context_size=context_size
)
decoded_text = token_ids_to_text(token_ids, tokenizer)
print(decoded_text.replace("\n", " "))  # Compact print format
model.train() # Back to training model applying all the configurations

Почати навчання

import time
start_time = time.time()

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.0004, weight_decay=0.1)

num_epochs = 10
train_losses, val_losses, tokens_seen = train_model_simple(
model, train_loader, val_loader, optimizer, device,
num_epochs=num_epochs, eval_freq=5, eval_iter=5,
start_context="Every effort moves you", tokenizer=tokenizer
)

end_time = time.time()
execution_time_minutes = (end_time - start_time) / 60
print(f"Training completed in {execution_time_minutes:.2f} minutes.")

Print training evolution

За допомогою наступної функції можна вивести еволюцію моделі під час її навчання.

import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
import math
def plot_losses(epochs_seen, tokens_seen, train_losses, val_losses):
fig, ax1 = plt.subplots(figsize=(5, 3))
ax1.plot(epochs_seen, train_losses, label="Training loss")
ax1.plot(
epochs_seen, val_losses, linestyle="-.", label="Validation loss"
)
ax1.set_xlabel("Epochs")
ax1.set_ylabel("Loss")
ax1.legend(loc="upper right")
ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
ax2 = ax1.twiny()
ax2.plot(tokens_seen, train_losses, alpha=0)
ax2.set_xlabel("Tokens seen")
fig.tight_layout()
plt.show()

# Compute perplexity from the loss values
train_ppls = [math.exp(loss) for loss in train_losses]
val_ppls = [math.exp(loss) for loss in val_losses]
# Plot perplexity over tokens seen
plt.figure()
plt.plot(tokens_seen, train_ppls, label='Training Perplexity')
plt.plot(tokens_seen, val_ppls, label='Validation Perplexity')
plt.xlabel('Tokens Seen')
plt.ylabel('Perplexity')
plt.title('Perplexity over Training')
plt.legend()
plt.show()

epochs_tensor = torch.linspace(0, num_epochs, len(train_losses))
plot_losses(epochs_tensor, tokens_seen, train_losses, val_losses)

Зберегти модель

Можливо зберегти модель + оптимізатор, якщо ви хочете продовжити навчання пізніше:

# Save the model and the optimizer for later training
torch.save({
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
},
"/tmp/model_and_optimizer.pth"
)
# Note that this model with the optimizer occupied close to 2GB

# Restore model and optimizer for training
checkpoint = torch.load("/tmp/model_and_optimizer.pth", map_location=device)

model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.1)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.train(); # Put in training mode

Або просто модель, якщо ви плануєте лише її використовувати:

# Save the model
torch.save(model.state_dict(), "model.pth")

# Load it
model = GPTModel(GPT_CONFIG_124M)

model.load_state_dict(torch.load("model.pth", map_location=device))

model.eval() # Put in eval mode

Завантаження ваг GPT2

Є 2 швидкі скрипти для локального завантаження ваг GPT2. Для обох ви можете клонувати репозиторій https://github.com/rasbt/LLMs-from-scratch локально, а потім:

• Скрипт https://github.com/rasbt/LLMs-from-scratch/blob/main/ch05/01_main-chapter-code/gpt_generate.py завантажить всі ваги та перетворить формати з OpenAI на ті, що очікуються нашим LLM. Скрипт також підготовлений з необхідною конфігурацією та з підказкою: "Кожне зусилля наближає вас"
• Скрипт https://github.com/rasbt/LLMs-from-scratch/blob/main/ch05/02_alternative_weight_loading/weight-loading-hf-transformers.ipynb дозволяє вам завантажити будь-які ваги GPT2 локально (просто змініть змінну CHOOSE_MODEL) і передбачити текст з деяких підказок.

Посилання

• https://www.manning.com/books/build-a-large-language-model-from-scratch