Training from scratch creates a foundation language model rather than adapting an existing one. It starts with an initialized model and a large text corpus, then learns parameters by repeatedly predicting the next token and adjusting the model when its prediction is wrong. It is difficult and resource-intensive; if an available base model fits the problem, finetuning is usually the easier starting point. section 5.2 section 5.2.1 section 1.3
The practical sequence is:
- Choose an architecture to suit the task, dataset size, and compute budget. Transformer variants are typical; the blueprint specifies layers and their connections. section 5.2.1
- Acquire, clean, and curate diverse text, then tokenize it into numerical inputs. Data quality and variety strongly affect the result. section 5.2.1 section 5.5
- Initialize the model, batch the data, calculate predictions and loss, and update parameters through backpropagation. Monitor performance on unseen data: lower training loss alone does not rule out overfitting. section 5.5 section 5.2.1
Budget before you train
A rough sizing rule calls for VRAM greater than four times the model’s parameter count in billions, with room for training batches. A 1B-parameter model may need at least 5–6 GB; a 70B-parameter model with a 32K-token context is estimated around 300 GB. Batch size and context length change the requirement. Finetuning can need substantially less memory and data. section 5.2
Listing 5.1 An example of training from scratch
import os
import torch
from accelerate import Accelerator
import bitsandbytes as bnb
class GPT(torch.nn.Module): #1
def __init__(self):
super().__init__()
self.token_embedding = torch.nn.Embedding(vocab_size, n_embed)
self.positional_embedding = torch.nn.Embedding(block_size, n_embed)
self.blocks = torch.nn.Sequential(
*[Block(n_embed, n_head=n_head) for _ in range(n_layer)]
)
self.ln_f = torch.nn.LayerNorm(n_embed)
self.lm_head = torch.nn.Linear(n_embed, vocab_size)
self.apply(self._init_weights)
def forward(self, idx, targets=None):
B, T = idx.shape
tok_emb = self.token_embedding(idx)
pos_emb = self.positional_embedding(torch.arange(T, device=device))
x = tok_emb + pos_emb
x = self.blocks(x)
x = self.ln_f(x)
logits = self.lm_head(x)
if targets is None:
loss = None
else:
B, T, C = logits.shape
logits = logits.view(B * T, C)
targets = targets.view(B * T)
loss = torch.nn.functional.cross_entropy(logits, targets)
return logits, loss
def _init_weights(self, module):
if isinstance(module, torch.nn.Linear):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module, torch.nn.Embedding):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
def generate(self, idx, max_new_tokens):
for _ in range(max_new_tokens):
idx_cond = idx[:, -block_size:]
logits, loss = self(idx_cond)
logits = logits[:, -1, :]
probs = torch.nn.functional.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
idx = torch.cat((idx, idx_next), dim=1)
return idx
class Block(torch.nn.Module): #2
def __init__(self, n_embed, n_head):
super().__init__()
head_size = n_embed // n_head
self.self_attention = MultiHeadAttention(n_head, head_size)
self.feed_forward = FeedFoward(n_embed)
self.ln1 = torch.nn.LayerNorm(n_embed)
self.ln2 = torch.nn.LayerNorm(n_embed)
def forward(self, x):
x = x + self.self_attention(self.ln1(x))
x = x + self.feed_forward(self.ln2(x))
return x
class MultiHeadAttention(torch.nn.Module):
def __init__(self, num_heads, head_size):
super().__init__()
self.heads = torch.nn.ModuleList(
[Head(head_size) for _ in range(num_heads)]
)
self.projection = torch.nn.Linear(head_size * num_heads, n_embed)
self.dropout = torch.nn.Dropout(dropout)
def forward(self, x):
out = torch.cat([h(x) for h in self.heads], dim=-1)
out = self.dropout(self.projection(out))
return out
class Head(torch.nn.Module):
def __init__(self, head_size):
super().__init__()
self.key = torch.nn.Linear(n_embed, head_size, bias=False)
self.query = torch.nn.Linear(n_embed, head_size, bias=False)
self.value = torch.nn.Linear(n_embed, head_size, bias=False)
self.register_buffer(
"tril", torch.tril(torch.ones(block_size, block_size))
)
self.dropout = torch.nn.Dropout(dropout)
def forward(self, x):
_, T, _ = x.shape
k = self.key(x)
q = self.query(x)
attention = q @ k.transpose(-2, -1) * k.shape[-1] ** 0.5
attention = attention.masked_fill(
self.tril[:T, :T] == 0, float("-inf")
)
attention = torch.nn.functional.softmax(attention, dim=-1)
attention = self.dropout(attention)
v = self.value(x)
out = attention @ v
return out
class FeedFoward(torch.nn.Module):
def __init__(self, n_embed):
super().__init__()
self.net = torch.nn.Sequential(
torch.nn.Linear(n_embed, 4 * n_embed),
torch.nn.ReLU(),
torch.nn.Linear(4 * n_embed, n_embed),
torch.nn.Dropout(dropout),
)
def forward(self, x):
return self.net(x)
def encode(string): #3
return [utt2int[c] for c in string]
def decode(line):
return "".join([int2utt[i] for i in line])
def get_batch(split):
data = train_data if split == "train" else val_data
idx = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([data[i : i + block_size] for i in idx])
y = torch.stack([data[i + 1 : i + block_size + 1] for i in idx])
x, y = x.to(device), y.to(device)
return x, y
@torch.no_grad()
def estimate_loss():
out = {}
model.eval()
for split in ["train", "val"]:
losses = torch.zeros(eval_iters)
for k in range(eval_iters):
X, Y = get_batch(split)
logits, loss = model(X, Y)
losses[k] = loss.item()
out[split] = losses.mean()
model.train()
return out
if __name__ == "__main__": #4
batch_size = 64 # Number of utterances at once #5
block_size = 256 # Maximum context window size
max_iters = 5000
eval_interval = 500
learning_rate = 3e-4
eval_iters = 200
n_embed = 384
n_head = 6
n_layer = 6
dropout = 0.2
accelerator = Accelerator()
device = accelerator.device
doing_quantization = False # Change to True if imported bitsandbytes
with open("./data/crimeandpunishment.txt", "r", encoding="utf-8") as f: #6
text = f.read()
chars = sorted(list(set(text))) #7
vocab_size = len(chars)
utt2int = {ch: i for i, ch in enumerate(chars)}
int2utt = {i: ch for i, ch in enumerate(chars)}
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]
model = GPT().to(device) #8
print("Instantiated Model")
print(
sum(param.numel() for param in model.parameters()) / 1e6,
"Model parameters",
)
optimizer = (
torch.optim.AdamW(model.parameters(), lr=learning_rate)
if not doing_quantization
else bnb.optim.Adam(model.parameters(), lr=learning_rate)
)
print("Instantiated Optimizer")
model, optimizer, train_data = accelerator.prepare(
model, optimizer, train_data
)
print("Prepared model, optimizer, and data")
#
for iter in range(max_iters): #9
print(f"Running Epoch {iter}")
if iter % eval_interval == 0 or iter == max_iters - 1:
losses = estimate_loss()
print(
f"| step {iter}: train loss {losses['train']:.4f} "
"| validation loss {losses['val']:.4f} |"
)
xb, yb = get_batch("train")
logits, loss = model(xb, yb)
optimizer.zero_grad(set_to_none=True)
accelerator.backward(loss)
optimizer.step()
model_dir = "./models/scratchGPT/" #10
if not os.path.exists(model_dir):
os.makedirs(model_dir)
model_path = model_dir + "model.pt" #11
torch.save(
model.state_dict(),
model_path,
)
loaded = GPT().load_state_dict(model_path) #12
context = torch.zeros((1, 1), dtype=torch.long, device=device) #13
print(decode(loaded.generate(context, max_new_tokens=500)[0].tolist()))
A PyTorch GPT training example shows the model components and initialization behind the workflow. section 5.2.1
