LLMs-from-Scratch 完全指南：从零构建 GPT 大语言模型

# 大多数开发者的 LLM 使用方式
from openai import OpenAI
client = OpenAI(api_key="...")
response = client.chat.completions.create(
    model="gpt-4o",
    messages=[{"role": "user", "content": "写一个快速排序"}]
)

# 克隆仓库
git clone https://github.com/rasbt/LLMs-from-scratch.git
cd LLMs-from-scratch

# 创建虚拟环境（推荐）
python -m venv venv
source venv/bin/activate  # Linux/macOS
# venv\Scripts\activate   # Windows

# 安装依赖
pip install -r requirements.txt

# test_setup.py
import torch
from ch04.ch04_main import MultiHeadAttention

# 检查 PyTorch 是否可用
print(f"PyTorch 版本: {torch.__version__}")
print(f"CUDA 可用: {torch.cuda.is_available()}")

# 快速测试注意力模块
mha = MultiHeadAttention(d_in=128, d_out=128, num_heads=4)
x = torch.randn(2, 10, 128)  # (batch, seq_len, embed_dim)
output = mha(x)
print(f"输出形状: {output.shape}")  # 应该是 (2, 10, 128)

# 直观例子
from tiktoken import encoding_for_model

enc = encoding_for_model("gpt-3.5-turbo")
text = "Building LLMs from scratch is rewarding."
tokens = enc.encode(text)
print(tokens)
# 输出: [73448, 2672, 310, 2428, 437, 6740, 13, 23441, 12047, 29879, 11337, 264, 4157, 12240, 264, 9820, 9391, 28671, 264, 4145, 220, 428, 1022, 372, 9006, 29892, 5614, 139, 231, 8698, 304, 10215, 13]

# 简化版自注意力（单头，无掩码）
def simple_attention(X):
    """
    X: 输入序列, shape (seq_len, embed_dim)
    返回: 注意力输出, shape (seq_len, embed_dim)
    """
    # Step 1: 生成 Q, K, V
    W_q = torch.randn(embed_dim, embed_dim)
    W_k = torch.randn(embed_dim, embed_dim)
    W_v = torch.randn(embed_dim, embed_dim)
    
    Q = X @ W_q
    K = X @ W_k
    V = X @ W_v
    
    # Step 2: 计算注意力分数
    scores = Q @ K.T  # shape: (seq_len, seq_len)
    
    # Step 3: 缩放（为什么除以 √d_k？让梯度更稳定）
    d_k = K.shape[-1]
    scores = scores / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))
    
    # Step 4: softmax 得到注意力权重
    attn_weights = torch.softmax(scores, dim=-1)
    
    # Step 5: 加权求和
    output = attn_weights @ V
    return output

# 准备训练数据
from ch06.ch06_main import create_dataloader

with open("data/the-verdict.txt", "r", encoding="utf-8") as f:
    raw_text = f.read()

# 词元化
tokenizer = tiktoken.get_encoding("gpt2")
token_ids = tokenizer.encode(raw_text)

# 创建数据加载器
train_loader = create_dataloader(
    token_ids,
    batch_size=4,
    max_length=256,    # 每个样本最多 256 个词元
    stride=128,         # 滑动窗口步长（有重叠）
    shuffle=True
)

# GPT 模型配置（小型版本，适合本地训练）
GPT_CONFIG_SMALL = {
    "vocab_size": 50257,      # GPT-2 词表大小
    "d_model": 768,           # 嵌入维度
    "n_heads": 12,            # 注意力头数
    "n_layers": 12,           # Transformer 块数
    "max_seq_len": 1024,      # 最大序列长度
    "dropout": 0.1,          # dropout 比例
    "qkv_bias": False         # 注意力的 QKV 是否加偏置
}

# 简化版训练循环
model = GPTModel(GPT_CONFIG_SMALL)
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.1)

for epoch in range(num_epochs):
    for batch in train_loader:
        input_ids, target_ids = batch
        
        # 前向传播
        logits = model(input_ids)        # shape: (batch, seq_len, vocab_size)
        loss = cross_entropy(logits, target_ids)
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

def generate(model, idx, max_new_tokens, temperature=1.0):
    """
    idx: 初始词元序列, shape (batch, seq_len)
    """
    model.eval()
    with torch.no_grad():
        for _ in range(max_new_tokens):
            # 前向传播得到下一个词元的 logits
            logits = model(idx)                      # (batch, seq_len, vocab_size)
            logits = logits[:, -1, :]              # 只看最后一个位置 (batch, vocab_size)
            
            # 温度采样
            if temperature > 0:
                logits = logits / temperature
                probs = torch.softmax(logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            else:
                # greedy 解码
                next_token = torch.argmax(logits, dim=-1, keepdim=True)
            
            # 把新词元追加到序列
            idx = torch.cat([idx, next_token], dim=1)
    
    return idx

# 使用
prompt = "Once upon a time"
token_ids = tokenizer.encode(prompt)
token_ids = torch.tensor(token_ids).unsqueeze(0)  # 加 batch 维度
output_ids = generate(model, token_ids, max_new_tokens=100, temperature=0.8)
print(tokenizer.decode(output_ids.squeeze().tolist()))