Training an LLM from Scratch, Locally — A Practical Walkthrough

1
class CausalSelfAttention(nn.Module):
2
    def __init__(self, config):
3
        super().__init__()
4
        assert config.n_embd % config.n_head == 0
5
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
6
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
7
        self.n_head = config.n_head
8
        self.n_embd = config.n_embd
9
        self.register_buffer("bias", torch.tril(
10
            torch.ones(config.block_size, config.block_size)
11
        ).view(1, 1, config.block_size, config.block_size))
12

13
    def forward(self, x):
14
        B, T, C = x.size()
15
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
16
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
17
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
18
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
19

20
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
21
        att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
22
        att = F.softmax(att, dim=-1)
23
        y = att @ v
24
        y = y.transpose(1, 2).contiguous().view(B, T, C)
25
        return self.c_proj(y)

1
class MLP(nn.Module):
2
    def __init__(self, config):
3
        super().__init__()
4
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
5
        self.gelu = nn.GELU()
6
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
7

8
    def forward(self, x):
9
        x = self.c_fc(x)
10
        x = self.gelu(x)
11
        x = self.c_proj(x)
12
        return x

1
class Block(nn.Module):
2
    def __init__(self, config):
3
        super().__init__()
4
        self.ln_1 = nn.LayerNorm(config.n_embd)
5
        self.attn = CausalSelfAttention(config)
6
        self.ln_2 = nn.LayerNorm(config.n_embd)
7
        self.mlp = MLP(config)
8

9
    def forward(self, x):
10
        x = x + self.attn(self.ln_1(x))
11
        x = x + self.mlp(self.ln_2(x))
12
        return x

1
class GPT(nn.Module):
2
    def __init__(self, config):
3
        super().__init__()
4
        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
5
        self.pos_emb = nn.Embedding(config.block_size, config.n_embd)
6
        self.blocks = nn.ModuleDict(dict(
7
            h=[Block(config) for _ in range(config.n_layer)]
8
        ))
9
        self.ln_f = nn.LayerNorm(config.n_embd)
10
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
11

12
    def forward(self, idx, targets=None):
13
        B, T = idx.size()
14
        tok_emb = self.tok_emb(idx)
15
        pos_emb = self.pos_emb(torch.arange(T, device=idx.device))
16
        x = tok_emb + pos_emb
17

18
        for block in self.blocks.h:
19
            x = block(x)
20

21
        x = self.ln_f(x)
22
        logits = self.lm_head(x)
23

24
        if targets is not None:
25
            loss = F.cross_entropy(
26
                logits.view(-1, logits.size(-1)),
27
                targets.view(-1)
28
            )
29
            return logits, loss
30
        return logits

1
def get_batch(split):
2
    data = train_data if split == 'train' else val_data
3
    ix = torch.randint(len(data) - block_size, (batch_size,))
4
    x = torch.stack([data[i:i+block_size] for i in ix])
5
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
6
    x, y = x.to(device), y.to(device)
7
    return x, y

1
max_steps = 5000
2
warmup_steps = 100
3
max_lr = 3e-4
4

5
def get_lr(step):
6
    if step < warmup_steps:
7
        return max_lr * step / warmup_steps
8
    decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps)
9
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
10
    return max_lr * coeff

1
model = GPT(config).to(device)
2
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
3

4
for step in range(max_steps):
5
    xb, yb = get_batch('train')
6
    logits, loss = model(xb, yb)
7
    optimizer.zero_grad(set_to_none=True)
8
    loss.backward()
9
    optimizer.step()
10

11
    if step % 200 == 0:
12
        with torch.no_grad():
13
            val_xb, val_yb = get_batch('val')
14
            _, val_loss = model(val_xb, val_yb)
15
        print(f"step {step}: train loss {loss.item():.4f}, val loss {val_loss.item():.4f}")
16

17
    if step % 1000 == 0:
18
        torch.save(model.state_dict(), f'checkpoint_{step}.pt')

1
def generate(model, idx, max_new_tokens, temperature=0.7, top_k=None):
2
    for _ in range(max_new_tokens):
3
        idx_cond = idx[:, -block_size:]
4
        logits, _ = model(idx_cond)
5
        logits = logits[:, -1, :] / temperature
6

7
        if top_k is not None:
8
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
9
            logits[logits < v[:, [-1]]] = float('-inf')
10

11
        probs = F.softmax(logits, dim=-1)
12
        idx_next = torch.multinomial(probs, num_samples=1)
13
        idx = torch.cat((idx, idx_next), dim=1)
14
    return idx

1
# Install dependencies
2
!pip install torch numpy tiktoken
3

4
# Download the Shakespeare dataset
5
!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt -O data/input.txt

1
model.py    — Transformer architecture (all the nn.Module classes)
2
train.py    — Data loading, training loop, evaluation
3
generate.py — Inference with temperature + top-k sampling