NanoChat 完整教程

# nanochat/model.py (简化版)
import torch
import torch.nn as nn
from torch.nn import functional as F

class CausalSelfAttention(nn.Module):
    """因果自注意力层"""
    def __init__(self, config):
        super().__init__()
        # Q, K, V投影
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        # 输出投影
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        # Causal mask
        self.register_buffer("bias", torch.tril(
            torch.ones(config.block_size, config.block_size)
        ).view(1, 1, config.block_size, config.block_size))
        
    def forward(self, x):
        B, T, C = x.size()  # batch, seq_len, embedding_dim
        
        # 计算Q, K, V
        qkv = self.c_attn(x)
        q, k, v = qkv.split(self.n_embd, dim=2)
        
        # 多头注意力
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        
        # 注意力分数（Flash Attention优化）
        att = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        
        # 合并多头
        y = att.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y)

class Block(nn.Module):
    """Transformer Block"""
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = nn.Sequential(
            nn.Linear(config.n_embd, 4 * config.n_embd),
            nn.GELU(),
            nn.Linear(4 * config.n_embd, config.n_embd),
        )
        
    def forward(self, x):
        # Pre-norm架构
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

# nanochat/dataset.py (简化版)
class DataLoader:
    def __init__(self, split, batch_size, seq_len, device):
        self.batch_size = batch_size
        self.seq_len = seq_len
        self.device = device
        
        # 加载数据shard
        self.shards = self._load_shards(split)
        self.reset()
    
    def _load_shards(self, split):
        """从磁盘加载预处理的数据文件"""
        data_dir = Path("data") / split
        shards = sorted(data_dir.glob("*.npy"))
        return shards
    
    def reset(self):
        """重置迭代器"""
        self.current_shard = 0
        self.current_position = 0
        self.tokens = np.load(self.shards[self.current_shard])
    
    def next_batch(self):
        """获取下一个batch"""
        B, T = self.batch_size, self.seq_len
        buf = self.tokens[self.current_position : self.current_position + B*T + 1]
        
        # X: 输入序列，Y: 目标序列（右移1位）
        x = torch.tensor(buf[:-1], dtype=torch.long).view(B, T)
        y = torch.tensor(buf[1:], dtype=torch.long).view(B, T)
        
        # 移动到GPU
        x, y = x.to(self.device), y.to(self.device)
        
        # 更新位置
        self.current_position += B * T
        if self.current_position + B*T + 1 > len(self.tokens):
            self._load_next_shard()
        
        return x, y