从0训练LLM原理解析

1. 统一背景：大模型到底在训练什么
   1.1 大模型看到的不是文字，而是 token id
   自然语言会先被 tokenizer 转换成整数序列：

"你好，世界"

-> tokenizer

-> 101, 2048, 330, 512, 2

模型只处理这些整数 id。每个 id 会先查 embedding 表，变成向量：

input_ids: B, T

embedding: B, T, H

其中：

B = batch size

T = sequence length

H = hidden size

V = vocabulary size

1.2 Decoder-only Causal LM 的基本结构

大多数自回归语言模型可以抽象成：

input_ids

-> token embedding

-> 多层 Transformer Decoder Block

-> final norm

-> lm_head

-> logits

张量形状：

input_ids: B, T

hidden: B, T, H

logits: B, T, V

每个位置的 logits 是对词表中所有 token 的打分。logitsb, t 表示第 b 个样本、第 t 个位置预测下一个 token 的分布。

一层常见 Decoder Block：

x = x + SelfAttention(RMSNorm(x))

x = x + FFN(RMSNorm(x))

常见组件：

RMSNorm:

比 LayerNorm 更简单，只按均方根缩放。

RoPE:

旋转位置编码，把位置信息注入 Q/K。

GQA:

query heads 多，key/value heads 少，减少 KV cache 和计算开销。

SwiGLU FFN:

down_proj(silu(gate_proj(x)) * up_proj(x))

MoE:

用多个 expert 替换普通 FFN，每个 token 只路由到 top-k 个 expert。

1.3 MoE 的 router auxiliary loss 怎么算

如果模型使用 MoE，FFN 不再是一个固定网络，而是多个 expert：

token hidden state

-> router / gate

-> 得到每个 expert 的概率

-> 选择 top-k expert

-> 只让 token 经过被选中的 expert

假设：

B = batch size

T = sequence length

N = B * T 当前 batch 里一共有多少个 token

H = hidden size

E = num_experts expert 数量

K = top-k 每个 token 选几个 expert

router 输入输出形状通常是：

x_flat: N, H 把 B, T, H 展平成 token 列表

router_logits: N, E 每个 token 对每个 expert 的打分

scores: N, E softmax 后的 expert 概率

topk_idx: N, K 每个 token 选中的 expert 编号

topk_weight: N, K 选中 expert 的概率权重

为什么需要 auxiliary loss？因为如果不约束 router，它可能把大部分 token 都送到同一个 expert：

expert 0: 90% token

expert 1: 5% token

expert 2: 3% token

expert 3: 2% token

这样会有两个问题：

1. 负载不均衡:

   少数 expert 很忙，其它 expert 几乎没训练到。

2. 容量浪费:

   名义上有多个 expert，但实际像只用了一个大 FFN。

   router auxiliary loss 的目标就是鼓励：

每个 expert 都被相对均匀地使用。

一种常见的简化计算方式是：

scores: N, E

scoresn, e 表示第 n 个 token 被 router 分给 expert e 的概率。

scores = softmax(router_logits, dim=-1)

topk_idx: N, K

表示每个 token 实际选择了哪些 expert。

topk_weight, topk_idx = topk(scores, k=K, dim=-1)

load: E

统计每个 expert 实际被选中的比例。

如果 E=4，load=0.25,0.25,0.25,0.25 就很均衡。

load = one_hot(topk_idx, num_classes=E).float().mean(dim=(0, 1))

importance: E

统计 router 给每个 expert 的平均概率。

它不只看最终 top-k 选择，也看 router 的软概率分布。

importance = scores.mean(dim=0)

aux_loss 越大，说明"实际选择"和"概率倾向"越集中到少数 expert。

乘 E 是为了让 loss 尺度在不同 expert 数量下更稳定。

router_aux_loss = (load * importance).sum() * E * aux_loss_coef

有些实现里 K=1，即每个 token 只选一个 expert。这时：

topk_idx: N, 1

load: E

也有实现会把 top-k 维度先展平再平均，核心思想一样：统计每个 expert 被选中的比例。

用一个小例子看。假设 E=4，当前 batch 有 8 个 token，top-1 路由结果是：

token 0 -> expert 0

token 1 -> expert 0

token 2 -> expert 0

token 3 -> expert 0

token 4 -> expert 1

token 5 -> expert 1

token 6 -> expert 2

token 7 -> expert 3

那么实际负载比例：

load = 4/8, 2/8, 1/8, 1/8

= 0.50, 0.25, 0.125, 0.125

如果 router 的平均概率也偏向 expert 0：

importance = 0.60, 0.20, 0.10, 0.10

那么：

(load * importance).sum()

= 0.500.60 + 0.25 0.20 + 0.1250.10 + 0.125 0.10

= 0.375

再乘 expert 数量和系数：

router_aux_loss = 0.375 * 4 * aux_loss_coef

如果完全均衡：

load = 0.25, 0.25, 0.25, 0.25

importance = 0.25, 0.25, 0.25, 0.25

(load * importance).sum()

= 4 * 0.25 * 0.25

= 0.25

可以看到，不均衡时这个值更大。训练时把它加到主 loss：

total_loss = language_modeling_loss + router_aux_loss

梯度会推动 router 不要总偏向少数 expert。

需要注意：router auxiliary loss 不是让每个 token 随机平均分配，而是在不破坏语言模型主任务的前提下，给 router 一个"别太偏科"的约束。主 loss 仍然是 next-token prediction，aux loss 只是辅助平衡。

1.4 Causal LM 的核心 loss

自回归语言模型做的是 next-token prediction：

给定前面的 token，预测下一个 token。

假设：

input_ids = A, B, C, D

训练目标是：

看到 A -> 预测 B

看到 A,B -> 预测 C

看到 A,B,C -> 预测 D

在张量上体现为错位：

shift_logits = logits:, :-1, : # B, T-1, V

shift_labels = labels:, 1: # B, T-1

loss = cross_entropy(shift_logits, shift_labels)

如果某些 token 不想参与训练，就把 label 设成 -100：

labels = token_id, token_id, -100, -100

ignore_index=-100 会让这些位置不计入 loss。

1.5 三种 mask 不要混

训练代码里常见三种 mask：

attention_mask

给 attention 用。

1 表示真实 token，0 表示 padding。

作用是避免模型注意到 pad token。

labels

给 supervised CE loss 用。

token id 表示参与训练，-100 表示忽略。

loss_mask / response_mask / completion_mask

给 DPO/RL 等自定义 loss 用。

1 表示这个 token 的 logprob 参与 loss，0 表示不参与。

一个 SFT 样本可以这样理解：

input_ids: system, user, assistant, pad

attention_mask: 1, 1, 1, 0

labels: -100, -100, assistant, -100

一个 Agent RL 样本可以这样理解：

input_ids: prompt, model_tool_call, tool_response, model_final_answer, pad

attention_mask: 1, 1, 1, 1, 0

response_mask: 0, 1, 0, 1, 0

工具返回结果需要被模型看到，所以 attention_mask=1；但它不是模型生成的，所以 response_mask=0。

1.6 怎么读懂文档里的维度

很多人读大模型训练代码卡住，不是因为公式本身难，而是因为不知道 B, T, V 这些维度在现实里对应什么。先用一个很小的例子建立直觉。

假设：

B = 2 一个 batch 有 2 条样本

T = 6 每条样本固定成 6 个 token

H = 4 每个 token 的隐藏向量长度是 4

V = 10 词表里一共有 10 个 token

那么 input_ids: B, T = 2, 6 可以想成一张表：

第 0 条样本: bos, 我, 喜欢, 猫, eos, pad

第 1 条样本: bos, 天气, 很, 好, eos, pad

模型第一步会把每个整数 token 查表变成向量：

input_ids: 2, 6

embedding: 2, 6, 4

也就是说，每个 token 不再只是一个整数，而是一个长度为 4 的向量：

"我" -> 0.12, -0.31, 0.08, 0.44

经过 Transformer 后，模型输出 logits: B, T, V = 2, 6, 10。这表示：

对 batch 里的每条样本，

对序列里的每个位置，

都给词表里的 10 个 token 各打一个分。

可以把 logits0, 2, : 理解为：

第 0 条样本，第 2 个位置，模型对"下一个 token 是词表中每个 token"的打分。

如果词表是：

0: pad

1: bos

2: eos

3: 我

4: 喜欢

5: 猫

6: 狗

7: 天气

8: 很

9: 好

那么 logits0, 2, : 可能是：

-3.1, -5.0, -2.7, -1.2, -0.4, 3.8, 0.6, -2.2, -1.9, -0.8

分数最高的是 token 5，也就是"猫"。这表示模型认为这个位置后面最可能接"猫"。

1.7 gather、mask、mean 到底在做什么

训练时我们不关心模型对所有词表 token 的分数，只关心"真实答案 token"的分数。

如果当前位置真实下一个 token 是"猫"，token id 是 5，那么：

log_probsb, t, : 是模型对整个词表的 log 概率

labelsb, t = 5 表示真实答案是 token 5

gather 后得到 log_probsb, t, 5

所以 gather 可以理解为：

从模型给整个词表的概率表里，挑出真实 token 对应的那一格。

mask 可以理解为：

哪些位置算分，哪些位置不算分。

例如：

token_logp: -0.2, -0.5, -1.0, -0.3, -0.1

mask: 1, 1, 0, 1, 0

乘起来：

token_logp * mask = -0.2, -0.5, 0, -0.3, 0

最后求平均或求和：

只让 mask=1 的位置影响 loss。

所以你看到这样的代码：

loss = (token_loss * mask).sum() / mask.sum()

可以直接翻译成：

只统计需要训练的位置，并对这些位置取平均。

1.8 为什么很多地方都是 B, T-1

语言模型训练永远有一个错位：

input_ids: bos, 我, 喜欢, 猫, eos

真正训练关系是：

看到 bos -> 预测 我

看到 bos,我 -> 预测 喜欢

看到 bos,我,喜欢 -> 预测 猫

看到 ...猫 -> 预测 eos

所以长度为 T 的序列，只能产生 T-1 个预测目标。最后一个 token 后面没有答案，第一个 token 没有前文预测它。

这就是为什么你经常看到：

logits:, :-1, : # 去掉最后一个位置

labels:, 1: # 去掉第一个位置

维度从 B, T 变成 B, T-1 是正常的，不是少了数据。

2. Tokenizer 训练
   2.1 Tokenizer 解决什么问题
   Tokenizer 是语言模型的"词典"和"切分规则"。它负责：

文本 -> token ids

token ids -> 文本

如果 tokenizer 改了，模型的 embedding 表和输出层也必须跟着变。一个已经训练好的模型不能随便换 tokenizer。

2.2 常见数据格式

Tokenizer 通常只需要纯文本流。可以来自预训练文本，也可以从对话数据里抽取所有 content。

对话格式：

{

"conversations": [

{"role": "user", "content": "你好"},

{"role": "assistant", "content": "你好，有什么可以帮你？"}

]

}

抽取成：

你好

你好，有什么可以帮你？

2.3 ByteLevel BPE 的直观原理

ByteLevel BPE 大致流程：

1. 把文本拆成 byte 级基础符号。
2. 统计相邻符号对出现频率。
3. 把最高频的相邻对合并成新 token。
4. 重复合并，直到达到 vocab_size。
5. 保存 tokenizer.json 和 tokenizer_config.json。
   特殊 token 必须提前加入，例如：

<|endoftext|>

<|im_start|>

<|im_end|>

<tool_call>

</tool_call>

<tool_response>

</tool_response>

这些 token 不只是"文本片段"，它们定义了训练和推理协议。

2.4 伪代码

def iter_texts(jsonl_path):

这个生成器负责把各种 jsonl 数据统一变成"纯文本流"。

tokenizer 训练不关心 role、label、reward，只关心文本里出现过哪些字符和片段。

for line in open(jsonl_path):

sample = json.loads(line)

    # 预训练数据常见格式：{"text": "..."}
    # 这种格式可以直接把 text 字段喂给 tokenizer 训练器。
    if "text" in sample:
        yield sample["text"]

    # 对话数据常见格式：{"conversations": [...]}
    # tokenizer 训练时不需要保留 user/assistant 边界，
    # 只需要抽取每条消息的 content，让 tokenizer 见过这些文本分布。
    elif "conversations" in sample:
        contents = []
        for msg in sample["conversations"]:
            if msg.get("content"):
                contents.append(msg["content"])

        # 用换行拼接，避免不同消息内容直接粘在一起。
        yield "\n".join(contents)

tokenizer = ByteLevelBPE()

tokenizer.train(

iterator 是惰性文本流，不必一次把全部数据加载进内存。

iterator=iter_texts(data_path),

# vocab_size 控制词表大小。
# 小词表会让模型参数更少，但文本会被切得更碎；
# 大词表压缩率更好，但 embedding/lm_head 参数更多。
vocab_size=6400,

# 特殊 token 必须提前固定下来。
# 后续 chat template、工具调用、思考标签都依赖这些 token 的稳定存在。
special_tokens=[
    "<|endoftext|>", "<|im_start|>", "<|im_end|>",
    "<think>", "</think>",
    "<tool_call>", "</tool_call>",
    "<tool_response>", "</tool_response>",
],

)

保存后，训练和推理都必须使用同一份 tokenizer。

tokenizer.save(output_dir)

3. 预训练：Pretraining

3.1 预训练在学什么

预训练是最基础的语言建模阶段。它不关心"用户/助手"角色，也不关心回答是否礼貌，核心目标只有一个：

预测下一个 token。

它让模型学到：

词语搭配

语法结构

事实和常识的统计相关性

代码/数学/多语言模式

长文本上下文关系

3.2 数据格式

最简单的预训练数据是一行一个文本：

{"text": "清晨的阳光透过窗帘洒进房间，桌上的书页被风轻轻翻动。"}

{"text": "Transformer 通过自注意力机制建模上下文关系，是现代语言模型的重要结构。"}

3.3 数据处理

处理流程：

1. 把原始文本转成 token id。

add_special_tokens=False 表示先不要自动加 bos/eos，

因为这里希望手动控制序列边界。

tokens = tokenizer(text, add_special_tokens=False)

2. 给 bos/eos 预留两个位置。

如果 max_length=1024，正文最多只能占 1022 个 token。

tokens = tokens:max_length - 2

3. 加上序列开始和结束标记。

bos/eos 能帮助模型学习"文本从哪里开始、在哪里结束"。

tokens = bos_id + tokens + eos_id

4. padding 到固定长度，便于组成 batch。

Transformer 通常要求一个 batch 内的样本长度一致。

input_ids = pad(tokens, max_length, pad_id)

5. labels 默认复制 input_ids。

Causal LM 内部会做 shift：用 input_idst 预测 labelst+1。

labels = input_ids.copy()

6. padding 位置不是真实文本，不应该产生 loss。

-100 是 PyTorch cross_entropy 常用的 ignore_index。

labelsinput_ids == pad_id = -100

张量形状：

单条样本:

input_ids: T

labels: T

batch 后:

input_ids: B, T

labels: B, T

logits: B, T, V

例子：

text = "你好"

tokens = bos, 你, 好, eos, pad, pad

labels = bos, 你, 好, eos, -100, -100

真正训练时：

logits0 预测 labels1

logits1 预测 labels2

logits2 预测 labels3

3.4 损失函数

模型前向输出每个位置对词表的打分。

logitsb, t, v 表示第 b 条样本第 t 个位置预测 token v 的分数。

logits = model(input_ids) # B, T, V

Causal LM 的第 t 个 logits 预测第 t+1 个 token。

所以最后一个 logits 没有下一个 token 可预测，要丢掉。

shift_logits = logits:, :-1, : # B, T-1, V

labels 的第 0 个 token 是 bos，没有前文用来预测它，要丢掉。

shift_labels = labels:, 1: # B, T-1

reshape 成二维后，cross_entropy 会把每个 token 位置当成一个分类样本。

ignore_index=-100 会跳过 padding 或不该训练的位置。

loss = cross_entropy(

shift_logits.reshape(-1, V),

shift_labels.reshape(-1),