打造属于自己的 GPT — 用 nanoGPT 从零开始训练语言模型

# Character-level 分词器（最简单）
text = "hello world"
chars = sorted(list(set(text)))
# chars = [' ', 'd', 'e', 'h', 'l', 'o', 'r', 'w']

stoi = {ch: i for i, ch in enumerate(chars)}  # char → int
itos = {i: ch for i, ch in enumerate(chars)}  # int → char

encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

print(encode("hello"))  # [3, 2, 4, 4, 5]
print(decode([3, 2, 4, 4, 5]))  # "hello"

import tiktoken
enc = tiktoken.get_encoding("gpt2")
tokens = enc.encode("让我们自己搭建一个 GPT！")
print(tokens)  # [171, 120, 230, 168, 245, ...]
print(len(tokens))  # ~15 tokens

import torch
import torch.nn as nn
import torch.nn.functional as F

class SelfAttention(nn.Module):
    def __init__(self, embed_dim, head_dim):
        super().__init__()
        self.query = nn.Linear(embed_dim, head_dim, bias=False)
        self.key = nn.Linear(embed_dim, head_dim, bias=False)
        self.value = nn.Linear(embed_dim, head_dim, bias=False)

    def forward(self, x):
        B, T, C = x.shape
        q = self.query(x)  # (B, T, head_dim)
        k = self.key(x)    # (B, T, head_dim)
        v = self.value(x)  # (B, T, head_dim)

        # Attention scores
        weights = q @ k.transpose(-2, -1)  # (B, T, T)
        weights = weights * (C ** -0.5)     # Scale

        # Causal mask — 未来的 token 不可见！
        mask = torch.tril(torch.ones(T, T))
        weights = weights.masked_fill(mask == 0, float('-inf'))

        weights = F.softmax(weights, dim=-1)
        out = weights @ v  # (B, T, head_dim)
        return out

class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        head_dim = embed_dim // num_heads
        self.heads = nn.ModuleList([
            SelfAttention(embed_dim, head_dim)
            for _ in range(num_heads)
        ])
        self.proj = nn.Linear(embed_dim, embed_dim)
        self.ffn = nn.Sequential(
            nn.Linear(embed_dim, 4 * embed_dim),
            nn.GELU(),
            nn.Linear(4 * embed_dim, embed_dim),
        )
        self.ln1 = nn.LayerNorm(embed_dim)
        self.ln2 = nn.LayerNorm(embed_dim)

    def forward(self, x):
        # Multi-Head Attention + Residual
        attn_out = torch.cat([h(self.ln1(x)) for h in self.heads], dim=-1)
        x = x + self.proj(attn_out)
        # Feed-Forward + Residual
        x = x + self.ffn(self.ln2(x))
        return x

class MicroGPT(nn.Module):
    def __init__(self, vocab_size, embed_dim=384, num_heads=6,
                 num_layers=6, block_size=256):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, embed_dim)
        self.pos_emb = nn.Embedding(block_size, embed_dim)
        self.blocks = nn.Sequential(*[
            TransformerBlock(embed_dim, num_heads)
            for _ in range(num_layers)
        ])
        self.ln_f = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_emb(idx)          # (B, T, embed_dim)
        pos_emb = self.pos_emb(torch.arange(T)) # (T, embed_dim)
        x = tok_emb + pos_emb

        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.head(x)  # (B, T, vocab_size)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1)
            )
        return logits, loss

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            logits, _ = self(idx[:, -256:])  # block_size 限制
            probs = F.softmax(logits[:, -1, :], dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, next_token], dim=1)
        return idx

# 在 spark01（GB10 128GB）上运行
cd ~/nanoGPT
python3 data/shakespeare_char/prepare.py
# → train: 1,003,854 tokens / val: 111,540 tokens

python3 train.py config/train_shakespeare_char.py \
  --device=cuda \
  --max_iters=5000 \
  --eval_interval=500 \
  --log_interval=100

ROMEO:
What say you to this? Let me not stay a whit;
And yet I feel the thing I have forgot
To take upon the honour of my word.