手搓484行代码是如何做到用普通CPU训练，耗时还不到20分钟就完成一个3.37M中文GPT模型？回答的效果还不错！一文讲清楚，30分钟轻松复刻

# 模型核心代码（src/model.py）
class GPT(nn.Module):
    def __init__(self, vocab_size, n_layer, n_head, n_embd, seq_len, dropout):
        super().__init__()
        self.seq_len = seq_len
        self.tok_emb = nn.Embedding(vocab_size, n_embd)  # 词嵌入
        self.blocks = nn.ModuleList([Block(n_embd, n_head, dropout) for _ in range(n_layer)])  # Transformer层
        self.norm = RMSNorm(n_embd)  # 简化版归一化
        self.head = nn.Linear(n_embd, vocab_size, bias=False)  # 输出层（共享词嵌入权重）

    def forward(self, idx):
        x = self.tok_emb(idx)  # 词嵌入
        mask = torch.tril(torch.ones(T, T, device=idx.device)).unsqueeze(0).unsqueeze(0)  # 下三角掩码
        for blk in self.blocks:
            x = blk(x, mask)  # 过Transformer层
        x = self.norm(x)
        return self.head(x)  # 输出logits

{"prompt": "什么是人工智能？", "completion": "人工智能是研究如何让机器模拟人类智能的技术。"}

# 训练循环核心（src/train.py）
accum = 0  # 累积计数器
for xb, yb in train_loader:
    xb, yb = xb.to(device), yb.to(device)
    logits = model(xb) #
    loss = loss_fn(logits.view(-1, logits.size(-1)), yb.view(-1)) #计算损失loss
    loss.backward()  # 累积梯度
    accum += 1
    if accum == bs // mb:  # 累积到等效batch_size
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # 梯度裁剪
        opt.step()  # 更新参数
        opt.zero_grad()
        sched.step()  # 调整学习率
        accum = 0

def forward(self, idx):
    B, T = idx.shape

    # 1. 词嵌入 (Token Embedding)
    x = self.tok_emb(idx) 
    # 对应图的最左侧输入。将输入的token ID序列转换为词嵌入向量。
    # x 的形状变为 (B, T, n_embd)
    # 2. 生成因果掩码 (Causal Mask)
    m = torch.tril(torch.ones(T, T, device=idx.device)).unsqueeze(0).unsqueeze(0)
    # 这是一个下三角矩阵，用于在自注意力计算中“遮蔽”未来的token，确保模型在预测第i个token时，只能看到前i-1个token。
    # 这个掩码 m 会被传递给每个 Block。
    # 3. 堆叠的 Decoder 层
    for blk in self.blocks:
        x = blk(x, m)
    # 循环遍历所有的Block。每一次循环，都执行一次完整的Decoder层计算（即图中的所有步骤）。
    # x 是上一层的输出，也是下一层的输入。
    # 6. 最终归一化和输出头
    x = self.norm(x)
    logits = self.head(x)
    # 所有Block计算完毕后，进行一次最终的归一化。
    # 然后通过一个线性层 `head` 将特征向量映射回词汇表大小，得到最终的预测logits。
    # 对应图的最右侧输出。
    return logits

def forward(self, x, mask):
    # 4. 第一个子层：多头注意力与残差连接
    x = x + self.attn(self.norm1(x), mask)
    # 这一行完美对应了图的上半部分：
    # a. `self.norm1(x)`: Layer Normalization (使用RMSNorm)。对应图中第一个Add & Norm左侧的Norm。
    # b. `self.attn(..., mask)`: Masked Multi-Head Attention。对应图中的紫色模块。
    # c. `x + ...`: Residual Connection (残差连接)。对应图中的第一个Add操作。
    # 5. 第二个子层：前馈网络与残差连接
    x = x + self.mlp(self.norm2(x))
    # 这一行完美对应了图的下半部分：
    # a. `self.norm2(x)`: Layer Normalization (使用RMSNorm)。对应图中第二个Add & Norm左侧的Norm。
    # b. `self.mlp(...)`: Feed-Forward Network (前馈网络)。对应图中的蓝色模块。
    # c. `x + ...`: Residual Connection (残差连接)。对应图中的第二个Add操作。
    return x

def forward(self, x, mask):
    # ... (形状定义)
    # 1. 生成 Q, K, V
    qkv = self.qkv(x)
    q, k, v = qkv.split(C, dim=-1)
    # 通过一个线性层一次性生成查询(Q), 键(K), 值(V)。
    # 2. 拆分多头
    q = q.view(B, T, h, self.head_dim).transpose(1, 2)
    k = k.view(B, T, h, self.head_dim).transpose(1, 2)
    v = v.view(B, T, h, self.head_dim).transpose(1, 2)
    # 将 Q, K, V 张量拆分成多个头，并调整维度以进行并行计算。
    # 3. 应用旋转位置编码 (RoPE)
    q, k = rope(q, k, T, self.head_dim, x.device)
    # 这是你的代码中一个重要的细节。它没有使用单独的位置嵌入层，而是采用了旋转位置编码。
    # 它将位置信息直接编码到 Q 和 K 的向量中，这是现代LLaMA等模型常用的方法。
    # 这一步对应了图中输入到注意力层之前隐含的“位置编码”步骤。
    # 4. 计算注意力分数并应用掩码
    att = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)
    att = att.masked_fill(mask == 0, float("-inf"))
    # 计算缩放点积注意力分数，并应用之前生成的因果掩码，将未来token的注意力分数设为负无穷。
    # 5. Softmax 和 Dropout
    att = F.softmax(att, dim=-1)
    att = self.drop(att)
    # 将注意力分数归一化为概率分布。
    # 6. 加权求和与输出投影
    y = att @ v
    y = y.transpose(1, 2).contiguous().view(B, T, C)
    y = self.proj(y)
    y = self.drop(y)
    # 将注意力权重应用到V上，然后将所有头的输出拼接起来，并通过一个最终的线性层（Projection）得到注意力子层的输出。

'''train方法整合了数据加载、模型构建、梯度累积、学习率调度（预热 + 余弦退火）、梯度裁剪、定期验证与保存，以及最终的模型量化等一系列关键技术，遵循了现代深度学习训练的最佳实践'''
def train(device_arg: str | None = None):
    # 1. 加载配置
    cfg = load_config("config.yaml")

    # 2. 设置随机种子以保证结果可复现
    set_seed(cfg["training"]["seed"])

    # 3. 设置PyTorch的线程数
    torch.set_num_threads(num_threads())
    # 4. 构建数据集和Tokenizer
    tok, train_ds, val_ds = build_datasets(cfg)

    # 5. 从配置中获取模型参数
    seq_len = cfg["model"]["seq_len"]

    # 6. 初始化GPT模型
    model = GPT(
        vocab_size=tok.vocab_size,
        n_layer=cfg["model"]["n_layer"],
        n_head=cfg["model"]["n_head"],
        n_embd=cfg["model"]["n_embd"],
        seq_len=seq_len,
        dropout=cfg["model"]["dropout"],
    )

    # 7. 确定并设置训练设备 (CPU/GPU)
    device = get_device(device_arg)
    model.to(device)
    # 8. 设置批处理大小和微批处理大小（用于梯度累积）
    bs = cfg["training"]["batch_size"]
    mb = cfg["training"]["micro_batch"]

    # 9. 创建DataLoader，用于高效加载数据
    train_loader = DataLoader(
        train_ds,
        batch_size=mb,
        shuffle=True,
        num_workers=0,
        collate_fn=lambda b: collate(b, seq_len, tok.pad_id),
    )
    # ... val_loader 类似 ...

    # 10. 定义优化器，用来更新梯度 
    opt = torch.optim.AdamW(
        model.parameters(),
        lr=cfg["training"]["lr"],
        weight_decay=cfg["training"]["weight_decay"],
    )

    # 11. 定义学习率调度器 (Warmup + Cosine Annealing)
    total_steps = cfg["training"]["max_steps"]
    warmup = cfg["training"]["warmup_steps"]
    def lr_lambda(step):
        if step < warmup:
            return step / max(1, warmup)  # 线性预热
        t = (step - warmup) / max(1, total_steps - warmup)
        return 0.5 * (1 + math.cos(math.pi * t)) # 余弦退火
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)

    # 12. 定义损失函数
    loss_fn = nn.CrossEntropyLoss(ignore_index=-100)
    # 13. 创建保存目录
    save_dir = cfg["training"]["save_dir"]
    ensure_dir(save_dir)

    step = 0
    accum = 0
    model.train()  # 设置模型为训练模式
    start_time = time.time()

    # 14. 循环直到达到总步数
    while step < total_steps:
        # 15. 遍历训练数据
        for xb, yb in train_loader:
            xb = xb.to(device)
            yb = yb.to(device)

            # 前向传播，logits张量是模型对输入 xb 进行预测后产生的原始、未归一化的分数
            logits = model(xb) # 触发一次完整的前向传播
            loss = loss_fn(logits.view(-1, logits.size(-1)), yb.view(-1)) # logits用计算损失

            # 反向传播，计算梯度
            loss.backward()

            accum += 1  #梯度累积 

            # 16. 当梯度累积到一定次数后，更新参数
            if accum == bs // mb:
                # 梯度裁剪，防止梯度爆炸
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                opt.step()      # 梯度更新，更新权重
                opt.zero_grad() # 清零梯度，为下一次累积做准备
                sched.step()    # 更新学习率,驱动整个学习率策略（从预热到衰减）的 “引擎”
                step += 1
                accum = 0

                # 17. 定期打印训练信息
                if step % 10 == 0:
                    print(f"step {step} loss {loss.item():.4f} lr {sched.get_last_lr()[0]:.6f}")

                # 18. 定期进行验证和保存
                if step % cfg["training"]["eval_interval"] == 0:
                    eval_loss = evaluate(model, val_loader, loss_fn, device)
                    # ... 打印并保存模型 ...
                    torch.save({"model": model.state_dict(), "cfg": cfg}, os.path.join(save_dir, "last.pt"))

                if step >= total_steps:
                    break

pip install -r requirements.txt
# 依赖很轻量：torch>=2.2、PyYAML、tokenizers，没有其他花里胡哨的库

python -m src.build_tokenizer
# 会在tokenizer/目录生成tokenizer.json（2048词表）

python -m src.train --device cpu

step 10 loss 5.2314 lr 0.000060
step 20 loss 4.8921 lr 0.000120
...
eval loss 3.2105 elapsed 600.2s  # 每20步验证一次
...

GPT_teacher-3.37M-cn % python -m src.infer --prompt "什么是注意力机制？" --show_label 
回答:注意力通过分配权重让模型重点参考关键位置，提升对序列关系的理解。

GPT_teacher-3.37M-cn % python -m src.infer --prompt "解释蒸馏水与纯水区别？" --show_label
回答:蒸馏水是通过蒸馏获得的水，去除了大部分杂质；纯水是指杂质含量极低的水，制备方式可以是蒸馏、反渗透等。

% python -m src.infer --prompt "RoPE的作用是" --show_label
回答:RoP把相对位置信息注入注意力，保留位序关系且高效。

效果不算惊艳，但作为 “从 0 到 1” 的成果，已经能体现 GPT 的核心逻辑 —— 用自注意力捕捉上下文关系。