In the previous chapter, we gave the GPT model the ability of dropout, and the three major engineering optimization techniques for pre-training are complete.
Now let us consider the last important problem of model pre-training: how to save the training results.
Training a GPT model takes hours or even days. Obviously, we cannot retrain it every time we use it. We need to save the trained model so that we can load and use it directly next time.
What to save¶
All the knowledge of the model is in its trainable parameters: the weights and biases of every layer. We only need to save these values, not the code, the computational graph, or the gradients.
Save format¶
NumPy provides the np.savez function, which can save multiple arrays into a single file (an archive file in the .npz format).
All the trainable parameters of the GPT model are stored in the model’s parameter list model.parameters, and we only need to save each of them to the file in turn.
When to save¶
During training, we can evaluate the performance of the model on the test set after each epoch. If the loss of the current epoch on the test set reaches a new low, we save the model. This way, we make sure that the best training result is kept.
How to load¶
Next time we use the model, we can read the parameters from the .npz file and assign them in turn to the model parameters model.parameters.
import os
import numpy as np
from src.core import *np.random.seed(42)class CharDataset(Dataset):
def __init__(self, filename, batch_size=1, context_size=64, stride=None, split=0.9):
self.filename = filename
self.context_size = context_size
self.stride = stride if stride is not None else context_size // 2
self.split = split
super().__init__(batch_size)
def load(self):
with open(self.filename, encoding="utf-8") as f:
text = f.read()
self.vocab = sorted(set(text))
self.vocab_size = len(self.vocab)
self.stoi = {ch: i for i, ch in enumerate(self.vocab)}
self.itos = {i: ch for i, ch in enumerate(self.vocab)}
self.tokens = self.encode(text)
split = int(len(self.tokens) * self.split)
self.train_data = self._pack(self.tokens[:split])
self.test_data = self._pack(self.tokens[split:])
def _pack(self, tokens):
x, y = [], []
for i in range(0, len(tokens) - self.context_size - 1, self.stride):
x.append(tokens[i: i + self.context_size])
y.append(tokens[i + 1: i + self.context_size + 1])
return x, y
def encode(self, symbols):
return [self.stoi[s] for s in symbols]
def decode(self, tokens):
return "".join(self.itos[t] for t in tokens)Lower triangular matrix layer¶
class Tril(Layer):
def __init__(self, value=-1e9):
super().__init__()
self.value = value
def forward(self, x: Tensor):
keep = np.tril(np.ones(x.shape[-2:]))
p = Tensor(np.where(keep, x.data, self.value))
def gradient_fn():
x.grad += p.grad * keep
return p.attach(gradient_fn, {x})Normalization layer¶
class LayerNorm(Layer):
def __init__(self, normalized_size, eps=1e-5):
super().__init__()
self.eps = eps
self.weight = Tensor(np.ones(normalized_size, dtype=DTYPE))
self.bias = Tensor(np.zeros(normalized_size, dtype=DTYPE))
def forward(self, x: Tensor):
mean = np.mean(x.data, axis=-1, keepdims=True)
var = np.var(x.data, axis=-1, keepdims=True, ddof=0)
norm = (x.data - mean) / np.sqrt(var + self.eps)
p = Tensor(self.weight.data * norm + self.bias.data)
def gradient_fn():
axis = tuple(range(p.grad.ndim - 1)) if p.grad.ndim > 1 else None
self.weight.grad += np.sum(p.grad * norm, axis=axis)
self.bias.grad += np.sum(p.grad, axis=axis)
grad = p.grad * self.weight.data
grad_mean = np.mean(grad, axis=-1, keepdims=True)
norm_mean = np.mean(grad * norm, axis=-1, keepdims=True)
x.grad += (grad - grad_mean - norm * norm_mean) / np.sqrt(var + self.eps)
return p.attach(gradient_fn, {x})
@property
def parameters(self):
return [self.weight, self.bias]Dropout layer¶
class Dropout(Layer):
def __init__(self, prob=0.1):
super().__init__()
self.prob = prob
def forward(self, x: Tensor):
if not self.training or self.prob == 0:
return x
keep_prob = 1.0 - self.prob
mask = (np.random.rand(*x.shape) < keep_prob).astype(DTYPE) / keep_prob
p = Tensor(x.data * mask)
def gradient_fn():
x.grad += p.grad * mask
return p.attach(gradient_fn, {x})class GELU(Layer):
def __init__(self):
super().__init__()
self.c = np.sqrt(2.0 / np.pi)
def forward(self, x: Tensor):
tanh = np.tanh(self.c * (x.data + 0.044715 * x.data ** 3))
a = Tensor(0.5 * x.data * (1.0 + tanh))
def gradient_fn():
grad = 0.5 * (1.0 + tanh) + 0.5 * x.data * (1.0 - tanh ** 2) * self.c * (1.0 + 3.0 * 0.044715 * x.data ** 2)
x.grad += a.grad * grad
return a.attach(gradient_fn, parents=[x])class GPTLoss(Loss):
def loss(self, p: Tensor, y: Tensor):
exp = np.exp(p.data - np.max(p.data, axis=-1, keepdims=True))
softmax = exp / np.sum(exp, axis=-1, keepdims=True)
flat_softmax = softmax.reshape(-1, softmax.shape[-1])
flat_y = y.data.reshape(-1).astype(np.int64)
n = len(flat_y)
rows = np.arange(n)
log = np.log(np.clip(flat_softmax[rows, flat_y], 1e-10, 1))
ce = Tensor(0 - np.sum(log) / n)
def gradient_fn():
flat_grad = flat_softmax.copy()
flat_grad[rows, flat_y] -= 1
p.grad += ce.grad * flat_grad.reshape(softmax.shape) / n
return ce.attach(gradient_fn, parents=[p])class GPTEmbedding(Composite):
def __init__(self, vocab_size, context_size, embedding_size):
self.embedding = Embedding(vocab_size, embedding_size)
self.positional_embedding = Embedding(context_size, embedding_size)
self.dropout = Dropout()
super().__init__([self.embedding,
self.positional_embedding,
self.dropout])
def forward(self, x: Tensor):
token = self.embedding(x)
position = self.positional_embedding(Tensor(range(x.shape[1])))
return self.dropout(token + position)GPT attention layer¶
class GPTAttention(Composite):
def __init__(self, embedding_size, heads=1):
self.embedding_size = embedding_size
self.heads = heads
self.normalize = LayerNorm(embedding_size)
self.query = Linear(embedding_size, embedding_size)
self.key = Linear(embedding_size, embedding_size)
self.value = Linear(embedding_size, embedding_size)
self.mask = Tril()
self.softmax = Softmax()
self.output = Linear(embedding_size, embedding_size)
self.dropout = Dropout()
super().__init__([self.normalize,
self.query,
self.key,
self.value,
self.mask,
self.softmax,
self.output,
self.dropout])
def forward(self, x: Tensor):
norm = self.normalize(x)
head_embedding_size = self.embedding_size // self.heads
multi_head_shape = (-1, x.shape[1], self.heads, head_embedding_size)
shape = (-1, x.shape[1], self.embedding_size)
query = self.query(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))
key = self.key(norm).reshape(multi_head_shape).transpose((0, 2, 3, 1))
value = self.value(norm).reshape(multi_head_shape).transpose((0, 2, 1, 3))
scale = Tensor(np.array(1.0 / np.sqrt(head_embedding_size)))
scores = query @ key * scale
weights = self.softmax(self.mask(scores))
return x + self.dropout(self.output((weights @ value).transpose((0, 2, 1, 3)).reshape(shape)))GPT feed-forward network layer¶
class GPTFeedForward(Composite):
def __init__(self, embedding_size):
self.normalize = LayerNorm(embedding_size)
self.input = Linear(embedding_size, embedding_size * 4)
self.gelu = GELU()
self.output = Linear(embedding_size * 4, embedding_size)
self.dropout = Dropout()
super().__init__([self.normalize,
self.input,
self.gelu,
self.output,
self.dropout])
def forward(self, x: Tensor):
norm = self.normalize(x)
h = self.gelu(self.input(norm))
return x + self.dropout(self.output(h))GPT Transformer block¶
class GPTTransformer(Composite):
def __init__(self, embedding_size, heads):
self.attention = GPTAttention(embedding_size, heads)
self.feed_forward = GPTFeedForward(embedding_size)
super().__init__([self.attention,
self.feed_forward])
def forward(self, x: Tensor):
x = self.attention(x)
return self.feed_forward(x)GPT output layer¶
class GPTOutput(Composite):
def __init__(self, embedding_size, vocab_size):
self.normalize = LayerNorm(embedding_size)
self.output = Linear(embedding_size, vocab_size)
super().__init__([self.normalize,
self.output])
def forward(self, x: Tensor):
norm = self.normalize(x)
return self.output(norm)GPT model¶
class GPT(Composite):
def __init__(self, vocab_size, context_size, embedding_size, heads, blocks):
self.embedding = GPTEmbedding(vocab_size, context_size, embedding_size)
self.transformers = [GPTTransformer(embedding_size, heads) for _ in range(blocks)]
self.output = GPTOutput(embedding_size, vocab_size)
super().__init__([self.embedding] + self.transformers + [self.output])
def forward(self, x: Tensor):
x = self.embedding(x)
for layer in self.transformers:
x = layer(x)
return self.output(x)GPT trainer¶
We add the functions to save and load model parameters to the GPT trainer:
save(the save function): number all the parameters of the model in order and save them to an.npzfile;load(the load function): load the data from the.npzfile and assign it in order to the parameters of the model;train(the training function):Shuffle: before each epoch, we added data shuffling. It randomizes the order of the samples and adds randomness to the training;
Log: track the changes of the loss and the learning rate during training;
Save the best model parameters: after each epoch, evaluate the training result on the test set. If the result is better, replace the saved parameters.
class GPTTrainer(Trainer):
def train(self, dataset, epochs, scheduler=None, filename=None):
dataset.train()
self.model.train()
steps = 0
# Best loss
best_losses = float('inf')
for epoch in range(epochs):
order = list(range(len(dataset)))
# Shuffle
np.random.shuffle(order)
total_loss = 0.0
for step, i in enumerate(order):
if scheduler is not None:
self.optimizer.lr = scheduler.step(steps)
feature, label = dataset[i]
prediction = self.model(feature)
loss = self.loss_fn(prediction, label)
self.optimizer.zero_grad()
loss.backward()
total_loss += float(loss.data)
self.optimizer.clip_grad_norm()
self.optimizer.step()
steps += 1
# Training log
if (step + 1) % 1000 == 0:
lr = f" lr {self.optimizer.lr:.6f}" if scheduler is not None else ""
print(f"epoch {epoch + 1} step {step + 1}/{len(dataset)} loss {(total_loss / 1000):.4f}{lr}")
total_loss = 0.0
_, total_loss = self.test(dataset)
# Test log
print(f"epoch {epoch + 1} done, test loss {total_loss:.4f}")
# Save the best model parameters
if total_loss < best_losses and filename is not None:
best_losses = total_loss
self.save(filename)
print(f"epoch {epoch + 1} saved model to {filename}")
print("GPT-training completed")
def test(self, dataset):
dataset.eval()
self.model.eval()
predictions = []
total_loss = 0.0
with Tensor.no_grad():
for i in range(len(dataset)):
feature, label = dataset[i]
prediction = self.model(feature)
loss = self.loss_fn(prediction, label)
predictions.append(prediction)
total_loss += float(loss.data)
loss = total_loss / len(dataset)
dataset.train()
self.model.train()
return predictions, loss
def generate(self, dataset, prompt, steps=256):
self.model.eval()
tokens = dataset.encode(prompt)
with Tensor.no_grad():
for _ in range(steps):
feature = Tensor([tokens[-dataset.context_size:]])
logits = self.model(feature)
last_logits = logits.data[0, -1]
exp = np.exp(last_logits - np.max(last_logits))
probs = exp / np.sum(exp)
token = np.random.choice(len(probs), p=probs)
tokens.append(token)
return dataset.decode(tokens)
# Save function
def save(self, filename):
params = {f"param_{i}": p.data for i, p in enumerate(self.model.parameters)}
np.savez(filename, **params)
# Load function
def load(self, filename):
if os.path.isfile(filename):
data = np.load(filename, allow_pickle=False)
for i, p in enumerate(self.model.parameters):
p.data = data[f"param_{i}"]
p.grad = np.zeros_like(p.data)Hyperparameters¶
In this chapter, we scale up the hyperparameters of the model and try to train a relatively more effective model. The final model has about 10.8M parameters.
💡 Tip: training takes about 1-2 hours on a Mac mini M4 with 16 GB of memory.
Learning rate¶
LEARNING_RATE = 0.001Batch size¶
BATCH_SIZE = 16Number of epochs¶
EPOCHS = 10Context size¶
CONTEXT_SIZE = 256Embedding size¶
EMBEDDING_SIZE = 384Number of heads¶
HEADS = 6Number of Transformer blocks¶
BLOCKS = 6Modeling¶
DATA_FILE = "../../../../data/tinyshakespeare.txt"
dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
loss_fn = GPTLoss()
optimizer = AdamWOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = GPTTrainer(model, loss_fn, optimizer)Training¶
During training we pass in MODEL_FILE, and the best model parameters are saved automatically after each epoch.
MODEL_FILE = "../../../../data/tinyshakespeare-gpt.npz"
scheduler = WarmupCosineScheduler(LEARNING_RATE, EPOCHS * len(dataset), 100, LEARNING_RATE / 10)
# Automatically save the best model parameters during training
trainer.train(dataset, EPOCHS, scheduler, MODEL_FILE)epoch 1 done, test loss 2.3235
epoch 1 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 2 done, test loss 2.0444
epoch 2 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 3 done, test loss 1.8580
epoch 3 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 4 done, test loss 1.7599
epoch 4 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 5 done, test loss 1.6474
epoch 5 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 6 done, test loss 1.5968
epoch 6 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 7 done, test loss 1.5635
epoch 7 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 8 done, test loss 1.5383
epoch 8 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 9 done, test loss 1.5339
epoch 9 saved model to ../../../../data/tinyshakespeare-gpt.npz
epoch 10 done, test loss 1.5331
epoch 10 saved model to ../../../../data/tinyshakespeare-gpt.npz
GPT-training completed
Evaluation¶
prediction, loss = trainer.test(dataset)
print(f'prediction:\t{len(prediction)} steps, each {prediction[0].data.shape}\nloss:\t{loss}')prediction: 54 steps, each (16, 256, 65)
loss: 1.5330633212018896
Inference¶
After training is completed, the best model parameters have been saved to the MODEL_FILE file. Next time we use the model, we can load it directly without retraining.
# Load the model parameters
trainer.load(MODEL_FILE)
print(trainer.generate(dataset, prompt="ROMEO:"))ROMEO:
As have done!
LEONTES:
Ho! when be letimest gentleman learness?
PAULINA:
Preventage straight the understand; such nor
Your victory fallse to dare.
DUKE VINCENTIO:
There, by the crown.
GLOUCESTER:
Indeed: thou art uncle and Verona wilt
Of thy sun-she s
In this chapter, we added the functions to save and load model parameters to the GPT trainer: the parameters of the best model are saved automatically during training, and can be loaded and used directly next time.
With this, the pre-training part of the GPT model is complete. In the next stage, we will enter the inference stage: how to use the pre-trained model to generate text.
Exercises¶
Try changing the number of epochs
EPOCHSto 20 and train for longer. Observe how the size of the saved model parameter file and the quality of the generated text change.