lolspider / notes
Clipping · 8 分钟

NanoWhale 项目深度解析:从零学习 DeepSeek-V4 架构

NanoWhale 项目深度解析:从零学习 DeepSeek-V4 架构

一句话总结:NanoWhale 是 HuggingFace 发布的一个教学项目,它用约 1.1 亿参数(~110M)实现了一个"迷你版" DeepSeek-V4 大语言模型,从零开始训练,让你能用一张 H100 显卡就跑通完整的大模型训练流程。


目录


一、这个项目在做什么?

大模型训练的一般流程

训练一个大语言模型(LLM)通常分为两个阶段:

plaintext
阶段 1:预训练(Pre-training)          阶段 2:微调(SFT)
┌──────────────────────┐          ┌──────────────────────┐
│ 海量无标注文本(书籍、网页)│          │ 高质量对话数据        │
│ ↓                     │          │ ↓                     │
│ 模型学习"语言规律"       │          │ 模型学会"对话和指令"    │
│ → 会补全句子,但不会聊天  │   ───→   │ → 能回答问题、执行指令  │
└──────────────────────┘          └──────────────────────┘

NanoWhale 的做法:

  1. 用 DeepSeek-V4 的架构(缩小版),在 FineWeb-Edu 数据集上预训练(约 26 亿 token)
  2. 用 SmolTalk 对话数据集进行 SFT 微调
  3. 最终得到一个 110M 参数的微型聊天模型

二、核心概念速查表

术语 全称 通俗解释
LLM Large Language Model 大语言模型,能理解和生成文本的 AI
Transformer — 当前主流的神经网络架构,GPT、BERT 都基于此
参数(Parameters) — 模型内部的可学习数值,类似大脑中的"突触",110M = 1.1 亿个
Token — 文本的基本单位。"Hello world" → ["Hello", " world"](2 个 token)
Hidden Size — 模型内部数据的"通道宽度",越大模型越强但越慢
Layer — 模型的层数,每层做一次"注意力 + 前馈"计算
Head Attention Head 注意力的"视角"数量,多个头能同时关注不同特征
MLA Multi-Head Latent Attention DeepSeek 的注意力创新,用低秩压缩节省内存
MoE Mixture of Experts 混合专家:多个"专家"子网络,每次只激活几个
RoPE Rotary Position Embedding 旋转位置编码,告诉模型"这个词在第几位"
RMSNorm Root Mean Square Normalization 一种归一化方法,稳定训练过程
SwiGLU Swish-Gated Linear Unit 一种激活函数,比 ReLU 效果更好
SFT Supervised Fine-Tuning 监督微调,用对话数据教模型"聊天"
bf16 / fp32 Brain Float 16 / Float 32 数值精度。bf16 省内存,fp32 更精确
Gradient Checkpointing — 用时间换内存的训练技巧,减少显存占用
Perplexity — 困惑度,衡量模型预测能力,越低越好
Sinkhorn Sinkhorn-Knopp Algorithm 一种矩阵归一化算法,使每行每列的和为 1

三、项目文件结构

plaintext
nanowhale/
├── modeling_deepseek_v4.py      # 🧠 模型架构代码(最核心)
├── configuration_deepseek_v4.py  # ⚙️ 模型配置类
├── requirements.txt              # 📦 依赖库
├── configs/
│   └── main_100m.yaml            # 📋 110M 模型的训练配置
├── scripts/
│   ├── train_pretrain.py          # 🏋️ 预训练脚本
│   ├── train_sft.py               # 💬 SFT 微调脚本
│   ├── eval_smoke.py              # 🧪 冒烟测试脚本
│   ├── chat.py                    # 🗣️ 对话交互脚本
│   └── count_params.py            # 🔢 参数统计工具
└── tokenizer/                     # 📝 分词器(DeepSeek-V4 的 tokenizer)

四、配置文件详解 — main_100m.yaml

这是整个项目的"设计蓝图",决定了模型的所有超参数:

yaml
model:
  vocab_size: 129280        # 词表大小:模型认识多少个 token
  hidden_size: 320          # 隐藏层维度:数据在模型内部的"宽度"
  num_hidden_layers: 8      # 层数:模型有多"深"
  num_attention_heads: 8    # 注意力头数:同时有几个"视角"
  num_key_value_heads: 1    # KV 头数 = 1 → MQA(Multi-Query Attention)

  # MoE 相关
  moe_intermediate_size: 640    # 每个专家的中间层大小
  n_routed_experts: 4           # 路由专家数量
  n_shared_experts: 1           # 共享专家数量(所有 token 都经过)
  num_experts_per_tok: 2        # 每个 token 选 2 个专家

  # MLA 注意力相关
  q_lora_rank: 160              # Q 的低秩压缩维度
  head_dim: 96                  # 每个头的维度
  qk_rope_head_dim: 32          # 用 RoPE 编码的维度
  o_groups: 2                   # 输出分组数
  o_lora_rank: 80               # 输出的低秩压缩维度

  # Hyper-Connections
  hc_mult: 4                    # HC 复制份数

逐项解释

vocab_size: 129280(词表大小)

  • 模型能识别 129,280 个不同的 token
  • 这个数字用的是 DeepSeek-V4 官方的分词器
  • 词表越大,Embedding 层参数越多(129280 × 320 ≈ 41M 参数,占总量的 37%)

hidden_size: 320(隐藏维度)

  • 相当于数据在模型内部流动的"管道直径"
  • 真正的 DeepSeek-V4 用 4096 或更大,这里缩小到 320

num_hidden_layers: 8(层数)

  • 模型有 8 层 Transformer Block
  • 每层都做:注意力计算 → 前馈网络

num_attention_heads: 8, num_key_value_heads: 1

  • 8 个 Query 头,但只有 1 个 KV 头
  • 这叫 MQA(Multi-Query Attention):所有 Query 头共享同一组 Key-Value
  • 好处:节省 KV 缓存的内存

MoE 参数

  • 4 个路由专家 + 1 个共享专家
  • 每个 token 经过门控网络选择 2 个路由专家
  • 同时所有 token 都会经过那个共享专家

五、模型架构详解 — modeling_deepseek_v4.py

这是整个项目的核心,约 500 行 Python 代码实现了 DeepSeek-V4 的完整架构。

整体数据流

plaintext
输入 token → Embedding → 扩展为 hc_mult 份
    → [Layer 1: HC-Attention + HC-FFN]
    → [Layer 2: HC-Attention + HC-FFN]
    → ...
    → [Layer 8: HC-Attention + HC-FFN]
    → HC Head 合并为 1 份 → RMSNorm → LM Head → 输出 logits

5.1 RMSNorm(均方根归一化)

python
class DeepseekV4RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))  # 可学习的缩放参数

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        x = x.float()                          # 转为 float32 计算
        var = x.pow(2).mean(-1, keepdim=True)   # 计算均方值
        x = x * torch.rsqrt(var + self.eps)     # 归一化
        return (self.weight * x).to(dtype)      # 缩放并恢复精度

5.2 RoPE(旋转位置编码)

python
def precompute_freqs_cis(dim, seqlen, base=10000.0):
    """预计算旋转位置编码的 cos/sin 值"""
    freqs = 1.0 / (base ** (torch.arange(0, dim, 2) / dim))  # 不同频率
    t = torch.arange(seqlen)              # 位置 [0, 1, 2, ..., seqlen-1]
    freqs = torch.outer(t, freqs)         # 位置 × 频率
    cos = freqs.cos()
    sin = freqs.sin()
    return torch.stack([cos, sin], dim=0)  # [2, seqlen, dim//2]

def apply_rotary_emb(x, cos_sin):
    """对输入应用旋转位置编码"""
    cos, sin = cos_sin[0], cos_sin[1]
    d = x.shape[-1] // 2
    x1, x2 = x[..., :d], x[..., d:]       # 把向量切成两半
    y1 = x1 * cos + x2 * sin               # 旋转公式
    y2 = -x1 * sin + x2 * cos
    return torch.cat([y1, y2], dim=-1)

5.3 Multi-Head Latent Attention(多头潜在注意力)

这是 DeepSeek 最重要的架构创新之一,目标是用更少的内存做注意力计算。

python
class DeepseekV4Attention(nn.Module):
    def __init__(self, config, layer_idx):
        super().__init__()
        self.num_heads = 8       # 8 个注意力头
        self.head_dim = 96       # 每头 96 维
        self.qk_rope_head_dim = 32  # 其中 32 维用 RoPE
        self.nope_head_dim = 64   # 另外 64 维不用位置编码

        # ====== Q 投影:低秩压缩 ======
        # 传统:hidden(320) → Q(8×96=768)  ← 参数多
        # MLA:  hidden(320) → 压缩(160) → Q(8×96=768)  ← 参数少!
        self.wq_a = nn.Linear(320, 160, bias=False)   # 压缩:320→160
        self.q_norm = RMSNorm(160)                     # 归一化
        self.wq_b = nn.Linear(160, 768, bias=False)    # 展开:160→768

        # ====== KV 投影:直接投影(只有 1 个 KV 头)=====
        self.wkv = nn.Linear(320, 96, bias=False)      # 320→96
        self.kv_norm = RMSNorm(96)

        # ====== O 投影:分组低秩 ======
        # 输出分 2 组,每组先压缩到 80 维,再合并回 320 维
        self.wo_a = nn.Linear(384, 160, bias=False)    # 分组压缩
        self.wo_b = nn.Linear(160, 320, bias=False)    # 恢复维度

MLA 的计算流程图:

plaintext
输入 hidden_states [B, S, 320]
│
├─→ Q 计算(低秩路径)
│   wq_a: 320 → 160(压缩)
│   q_norm: 归一化
│   wq_b: 160 → 768(展开为 8 头 × 96 维)
│   → Q [B, 8, S, 96]
│
├─→ KV 计算(只有 1 个头!MQA)
│   wkv: 320 → 96
│   kv_norm: 归一化
│   → KV [B, 1, S, 96](所有 8 个 Q 头共享这 1 个 KV)
│
├─→ 只在最后 32 维上应用 RoPE
│   Q = [Q_nope(64维), Q_rope(32维)]
│   KV = [KV_nope(64维), KV_rope(32维)]
│
├─→ 注意力计算:scaled_dot_product_attention
│   attn_output [B, 8, S, 96]
│
└─→ O 投影(分组低秩)
    分 2 组 → wo_a 压缩 → wo_b 展开 → output [B, S, 320]

5.4 Hyper-Connections(超连接)

这是 DeepSeek-V4 引入的另一个重要创新,替代了传统的残差连接(Residual Connection)。

python
def hc_split_sinkhorn(mixes, hc_scale, hc_base, hc_mult=4, sinkhorn_iters=20, eps=1e-6):
    """HC 的核心:拆分 + Sinkhorn 归一化"""
    # mixes 是线性层输出,拆为三部分
    pre_raw  = mixes[..., :4]              # 前连接权重(4 维)
    post_raw = mixes[..., 4:8]             # 后连接权重(4 维)
    comb_raw = mixes[..., 8:]              # 组合矩阵(4×4=16 维)

    # sigmoid 归一化到 [0, 1]
    pre = sigmoid(pre_raw * scale + base) + eps

    # softmax + Sinkhorn 迭代(使矩阵每行每列和为 1)
    comb = softmax(comb_raw)
    for _ in range(sinkhorn_iters - 1):
        comb = comb / comb.sum(dim=-1, keepdim=True)  # 行归一化
        comb = comb / comb.sum(dim=-2, keepdim=True)  # 列归一化

    return pre, post, comb

HC 在 Transformer Block 中的使用:

python
class DeepseekV4Block(nn.Module):
    def forward(self x, ...):  # x 形状: [B, S, hc_mult=4, D]
        # ---- 注意力子层 ----
        residual = x                        # 保存 4 份副本
        y, post, comb = self.hc_pre(x)      # 4 份 → 1 份(加权合并)
        y = self.attn_norm(y)               # RMSNorm
        y = self.attn(y)                    # 注意力计算
        x = self.hc_post(y, residual, post, comb)  # 1份 → 4份(用组合矩阵更新)

        # ---- 前馈子层 ----
        residual = x                        # 保存 4 份副本
        y, post, comb = self.hc_pre(x)      # 4 份 → 1 份
        y = self.ffn_norm(y)                # RMSNorm
        y = self.ffn(y)                     # MoE 前馈
        x = self.hc_post(y, residual, post, comb)  # 1份 → 4份

        return x

5.5 MoE(混合专家系统)

MoE 是让模型"变大但不变慢"的关键技术。

python
class DeepseekV4Expert(nn.Module):
    """单个专家:SwiGLU 前馈网络"""
    def __init__(self, hidden_size=320, intermediate_size=640):
        self.w1 = nn.Linear(320, 640, bias=False)  # 门控投影
        self.w2 = nn.Linear(640, 320, bias=False)   # 下投影
        self.w3 = nn.Linear(320, 640, bias=False)   # 上投影

    def forward(self, x):
        gate = self.w1(x)      # 计算"门控值"
        up = self.w3(x)        # 计算"上变换"
        x = silu(gate) * up    # 门控 × 上变换(SwiGLU)
        return self.w2(x)      # 压缩回原始维度
python
class DeepseekV4Gate(nn.Module):
    """门控网络:决定每个 token 路由到哪些专家"""

    def forward(self, x):
        scores = linear(x, self.weight)         # 计算每个 token 对每个专家的"匹配分"

        # sqrtsoftplus 评分:比 softmax 更平滑
        scores = softplus(scores).sqrt()

        # 加上可学习偏置
        scores = scores + self.bias

        # 选 Top-2 分数最高的专家
        indices = scores.topk(2)                # 每个 token 选 2 个专家
        weights = scores[indices]               # 取出对应分数
        weights = weights / weights.sum()       # 归一化为权重
        weights = weights * 1.5                 # 乘以缩放因子
        return weights, indices
python
class DeepseekV4MoE(nn.Module):
    """完整 MoE 层"""

    def __init__(self, config):
        self.gate = DeepseekV4Gate(config)           # 门控网络
        self.experts = [Expert() for _ in range(4)]   # 4 个路由专家
        self.shared_expert = Expert()                  # 1 个共享专家

    def forward(self, x):
        # Step 1: 门控选择
        weights, indices = self.gate(x)    # [B*S, 2], [B*S, 2]

        # Step 2: 路由到选中的专家
        y = zeros_like(x)
        for i in range(4):                 # 遍历 4 个专家
            mask = (indices == i)          # 找到选中专家 i 的 token
            if mask.any():
                expert_out = self.experts[i](x[mask])
                y[mask] += weights[mask] * expert_out  # 加权累加

        # Step 3: 加上共享专家的输出(所有 token 都经过)
        y = y + self.shared_expert(x)

        return y

5.6 SwiGLU 激活函数

plaintext
传统 FFN:   output = W2 × ReLU(W1 × x)
SwiGLU:     output = W2 × (SiLU(W1 × x) ⊙ W3 × x)
                                  ↑          ↑
                               门控信号     值信号

5.7 完整 Transformer Block

plaintext
输入 token embedding [B, S, 320]
    │
    ├── 扩展为 4 份 [B, S, 4, 320]  ← Hyper-Connections
    │
    │   ┌─────── 注意力子层 ───────┐
    ├──→│ HC-Pre: 4份 → 1份        │
    │   │ RMSNorm                  │
    │   │ MLA Attention             │
    │   │ HC-Post: 1份 → 4份       │
    │   └──────────────────────────┘
    │
    │   ┌─────── 前馈子层 ─────────┐
    ├──→│ HC-Pre: 4份 → 1份        │
    │   │ RMSNorm                  │
    │   │ MoE (4路由 + 1共享)       │
    │   │ HC-Post: 1份 → 4份       │
    │   └──────────────────────────┘
    │
    └──→ 输出 [B, S, 4, 320] → 送入下一层

5.8 完整模型组装

python
class DeepseekV4Model(nn.Module):
    def __init__(self, config):
        self.embed_tokens = nn.Embedding(129280, 320)   # 词嵌入
        self.layers = [DeepseekV4Block() for _ in range(8)]  # 8 层
        self.norm = RMSNorm(320)                         # 最终归一化
        self.hc_head = ...                               # HC 合并头

    def forward(self, input_ids):
        # 1. 词嵌入
        x = self.embed_tokens(input_ids)        # [B, S] → [B, S, 320]

        # 2. 扩展为 hc_mult=4 份
        x = x.unsqueeze(2).expand(-1, -1, 4, -1)  # [B, S, 4, 320]

        # 3. 逐层计算
        for layer in self.layers:
            x = layer(x)

        # 4. 合并 4 份 → 1 份
        x = self.hc_head(x)     # [B, S, 4, 320] → [B, S, 320]

        # 5. 最终归一化
        x = self.norm(x)        # [B, S, 320]
        return x


class DeepseekV4ForCausalLM(nn.Module):
    """因果语言模型(可以预测下一个 token)"""
    def __init__(self, config):
        self.model = DeepseekV4Model(config)
        self.lm_head = nn.Linear(320, 129280, bias=False)  # 输出层

    def forward(self, input_ids, labels=None):
        hidden = self.model(input_ids)              # [B, S, 320]
        logits = self.lm_head(hidden)               # [B, S, 129280] 每个 token 的概率

        if labels is not None:
            # 计算损失:预测下一个 token
            shift_logits = logits[..., :-1, :]      # 去掉最后一个
            shift_labels = labels[..., 1:]           # 去掉第一个
            loss = cross_entropy(shift_logits, shift_labels)
            return loss, logits

        return logits

六、训练流程详解

6.1 预训练 — train_pretrain.py

python
# 核心流程
model = DeepseekV4ForCausalLM(config)          # 1. 随机初始化模型
tokenizer = PreTrainedTokenizerFast(path)       # 2. 加载分词器
dataset = load_dataset("fineweb-edu")           # 3. 加载训练数据

trainer = SFTTrainer(                           # 4. 创建训练器
    model=model,
    train_dataset=dataset,
    args=SFTConfig(
        max_seq_length=2048,                    # 最大序列长度
        packing=True,                            # 把多个短文本拼到一起,不浪费空间
        per_device_train_batch_size=8,           # 每 GPU 每步 8 个样本
        gradient_accumulation_steps=4,           # 每 4 步更新一次参数
        # → 有效 batch size = 8 × 4 = 32
        learning_rate=6e-4,                      # 学习率
        lr_scheduler_type="cosine",              # 余弦退火调度
        warmup_ratio=0.03,                       # 前 3% 步数预热
        bf16=True,                               # 用 bf16 精度训练
        gradient_checkpointing=True,             # 梯度检查点(省内存)
    ),
)

trainer.train()                                  # 5. 开始训练!

6.2 SFT 微调 — train_sft.py

python
model = DeepseekV4ForCausalLM.from_pretrained("checkpoints/pretrain/final")  # 加载预训练模型
dataset = load_dataset("HuggingFaceTB/smol-smoltalk")                        # 加载对话数据

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    args=SFTConfig(
        learning_rate=2e-5,       # SFT 用更小的学习率(预训练用 6e-4)
        max_steps=3000,           # SFT 步数比预训练少得多
    ),
)
trainer.train()

七、评估与推理

7.1 冒烟测试 — eval_smoke.py

一个自动化的"健康检查"脚本,验证模型能否正常工作:

python
# 测试 1:加载 tokenizer → 检查词表
# 测试 2:加载 config → 检查架构参数
# 测试 3:加载模型 → 检查参数量
# 测试 4:文本生成 → 5 个 prompt 测试补全
# 测试 5:对话模板 → 测试 chat_template
# 测试 6:困惑度 → 计算 perplexity(越低越好)

7.2 对话交互 — chat.py

python
# 简单的交互式对话循环
model = DeepseekV4ForCausalLM.from_pretrained(model_path)
tokenizer = PreTrainedTokenizerFast.from_pretrained(model_path)

while True:
    user_input = input("User: ")
    messages.append({"role": "user", "content": user_input})

    # 用 chat_template 格式化
    prompt = tokenizer.apply_chat_template(messages)

    # 生成回复
    output = model.generate(input_ids, temperature=0.7, top_p=0.9)
    response = tokenizer.decode(output)
    print(f"Assistant: {response}")

八、训练结果与已知问题

训练结果

阶段 数据量 最终 Loss Token 准确率 硬件
预训练 ~2.6B token ~5.3 33.8% 1× H100 80GB
SFT SmolTalk eval loss 2.607 — —

Perplexity:SFT 后从 13.62 降到 12.90(越低越好)

已知问题


九、术语总表

术语 说明 在本项目中的值
vocab_size 词表大小 129,280
hidden_size 隐藏层维度 320
num_hidden_layers Transformer 层数 8
num_attention_heads 注意力头数 8
num_key_value_heads KV 头数(MQA) 1
head_dim 每头维度 96 = 32(RoPE) + 64(NoPE)
q_lora_rank Q 低秩压缩维度 160
o_groups 输出分组数 2
o_lora_rank 输出低秩维度 80
hc_mult Hyper-Connections 份数 4
n_routed_experts 路由专家数 4
n_shared_experts 共享专家数 1
num_experts_per_tok 每 token 选专家数 2
moe_intermediate_size 专家中间层大小 640
scoring_func 门控评分函数 sqrtsoftplus
routed_scaling_factor 路由缩放因子 1.5
rope_theta RoPE 频率基数 10,000
rms_norm_eps RMSNorm 稳定值 1e-6
initializer_range 权重初始化标准差 0.02
max_position_embeddings 最大上下文长度 2,048
learning_rate (pretrain) 预训练学习率 6e-4
learning_rate (sft) SFT 学习率 2e-5
batch_size 每 GPU 批量 8
gradient_accumulation 梯度累积步数 4
warmup_ratio 预热比例 3%
total_params 总参数量 ~110M

参考资料