1
2
3
4

# 原始 token 序列: [101, 7592, 2088, 102, 2023]
# 拆分为输入和目标：
X = torch.tensor(input_ids[:-1], dtype=torch.long)  # [101, 7592, 2088, 102]
Y = torch.tensor(input_ids[1:], dtype=torch.long)    # [7592, 2088, 102, 2023]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

# 训练时：并行计算所有位置（使用真实历史）
X = [101, 7592, 2088, 102]     # 真实的前 N-1 个 token
Y = [7592, 2088, 102, 2023]    # 真实的后 N-1 个 token
logits = model(X)               # 一次前向传播，得到所有位置的预测
loss = CrossEntropyLoss(logits, Y)  # 一次性计算损失

# 推理时：逐个生成（使用模型自己的预测）
tokens = [101]
for i in range(max_len):
    next_token = model(tokens).argmax()  # 模型自己预测下一个
    tokens.append(next_token)

1

{"text": "<|im_start|>我知道底边长度和高，分别是5cm和8cm，请告诉我计算出来的面积是多少。三角形的面积为20平方厘米。<|im_end|> <|im_start|>为什么理智告诉我们不要跟陌生人接触？因为陌生人可能会对我们的个人安全造成威胁。<|im_end|>"}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

from tokenizers import models, pre_tokenizers, trainers, Tokenizer

# 1. 初始化 BPE tokenizer
tokenizer = Tokenizer(models.BPE())
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)

# 2. 定义特殊 token
special_tokens = ["<|endoftext|>", "<|im_start|>", "<|im_end|>"]

# 3. 配置训练器
trainer = trainers.BpeTrainer(
    vocab_size=6400,                                    # 词表大小
    special_tokens=special_tokens,                      # 特殊 token
    show_progress=True,
    initial_alphabet=pre_tokenizers.ByteLevel.alphabet() # 初始字母表
)

# 4. 从语料中训练
texts = read_texts_from_jsonl(data_path)
tokenizer.train_from_iterator(texts, trainer=trainer)

# 5. 设置解码器
tokenizer.decoder = decoders.ByteLevel()

1
2
3
4
5
6

<|im_start|>system
You are a helpful assistant<|im_end|>
<|im_start|>user
你来自哪里？<|im_end|>
<|im_start|>assistant
我来自地球<|im_end|>

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

input_ids
    |
    v
[Embedding 层] --> token 嵌入向量 (batch, seq_len, 512)
    |
    v
[NanoQwenBlock x 8] --> 8 层 Transformer Block
    |   每层包含：
    |   - RMSNorm
    |   - GQA 注意力（含 RoPE）
    |   - 残差连接
    |   - RMSNorm
    |   - FeedForward（SwiGLU）/ MOE
    |   - 残差连接
    |
    v
[RMSNorm] --> 最终归一化
    |
    v
[LM Head] --> logits (batch, seq_len, 6400)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

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

    def _norm(self, x):
        # rsqrt = 1 / sqrt(x)，即计算 RMS 的倒数再乘以 x
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        return self.weight * self._norm(x.float()).type_as(x)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
    # 计算每对维度的频率：theta_i = 1 / (base^(2i/d))
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    # 生成位置序列 [0, 1, 2, ..., end-1]
    t = torch.arange(end, device=freqs.device)
    # 外积：每个位置 x 每个频率 -> 角度矩阵
    freqs = torch.outer(t, freqs).float()
    # 计算 cos 和 sin，并拼接（对应向量前半和后半）
    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)
    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)
    return freqs_cos, freqs_sin

def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    def rotate_half(x):
        # 将后半部分取反并移到前面：[x1,x2,x3,x4] -> [-x3,-x4,x1,x2]
        return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)

    # 旋转公式：q' = q * cos + rotate_half(q) * sin
    q_embed = (q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))
    k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))
    return q_embed, k_embed

1
2
3

num_attention_heads = 8    # Query 头数
num_key_value_heads = 2    # KV 头数
n_rep = 8 / 2 = 4          # 每组 KV 头被 4 个 Query 头共享

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    bs, slen, num_key_value_heads, head_dim = x.shape
    if n_rep == 1:
        return x  # MHA 的情况，无需复制
    return (
        x[:, :, :, None, :]
        .expand(bs, slen, num_key_value_heads, n_rep, head_dim)
        .reshape(bs, slen, num_key_value_heads * n_rep, head_dim)
    )
# 输入形状: (batch, seq_len, 2, 64)  -- 2个KV头
# 输出形状: (batch, seq_len, 8, 64)  -- 扩展为8个头，与Q匹配

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

class FeedForward(nn.Module):
    def __init__(self, config: NanoQwenConfig):
        super().__init__()
        if config.intermediate_size is None:
            # 自动计算中间层维度：hidden_size * 8/3，并对齐到 64 的倍数
            intermediate_size = int(config.hidden_size * 8 / 3)
            config.intermediate_size = 64 * ((intermediate_size + 64 - 1) // 64)

        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]  # SiLU 激活

    def forward(self, x):
        # SwiGLU: down_proj( silu(gate_proj(x)) * up_proj(x) )
        return self.dropout(self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

class MoEGate(nn.Module):
    def __init__(self, config):
        self.weight = nn.Parameter(torch.empty((n_routed_experts, hidden_size)))

    def forward(self, hidden_states):
        # 计算每个 token 对每个专家的得分
        logits = F.linear(hidden_states, self.weight)
        scores = logits.softmax(dim=-1)
        # 选择 Top-K 个专家
        topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1)
        # 归一化权重
        topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
        return topk_idx, topk_weight, aux_loss

1
2
3
4
5
6
7
8

class MOEFeedForward(nn.Module):
    def __init__(self, config):
        # 多个路由专家
        self.experts = nn.ModuleList([FeedForward(config) for _ in range(n_routed_experts)])
        # 门控网络
        self.gate = MoEGate(config)
        # 共享专家（所有 token 都经过）
        self.shared_experts = nn.ModuleList([FeedForward(config) for _ in range(n_shared_experts)])

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

class NanoQwenConfig(PretrainedConfig):
    model_type = "nanoqwen"

    def __init__(self,
        hidden_size=512,           # 隐藏维度
        num_hidden_layers=8,       # Transformer层数
        num_attention_heads=8,     # Query头数
        num_key_value_heads=2,     # KV头数（GQA）
        vocab_size=6400,           # 词表大小
        hidden_act='silu',         # 激活函数
        rms_norm_eps=1e-05,        # RMSNorm epsilon
        rope_theta=1000000.0,      # RoPE base
        flash_attn=True,           # Flash Attention
        use_moe=False,             # 是否启用 MOE
        ...
    ):

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

class NanoQwenBlock(nn.Module):
    def forward(self, hidden_states, position_embeddings, ...):
        # Pre-Norm + Attention + 残差
        residual = hidden_states
        hidden_states, present_key_value = self.self_attn(
            self.input_layernorm(hidden_states),  # Pre-Norm
            position_embeddings, ...
        )
        hidden_states += residual  # 残差连接

        # Pre-Norm + FFN/MOE + 残差
        hidden_states = hidden_states + self.mlp(
            self.post_attention_layernorm(hidden_states)  # Pre-Norm
        )
        return hidden_states, present_key_value

1
2
3
4
5
6

class NanoQwenForCausalLM(PreTrainedModel, GenerationMixin):
    def __init__(self, config):
        self.model = NanoQwenModel(config)           # Transformer 主体
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)  # 输出头
        # 权重绑定：Embedding 和 LM Head 共享权重
        self.model.embed_tokens.weight = self.lm_head.weight

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

# 1. 创建自动混合精度上下文
ctx = torch.cuda.amp.autocast()

# 2. 创建梯度缩放器（将小梯度放大，防止 FP16 下溢）
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))

# 3. 训练循环
with ctx:
    res = model(X)                    # 前向传播在混合精度下进行
    loss = loss_fct(res.logits, Y)    # 损失计算

scaler.scale(loss).backward()         # 缩放后的损失反向传播
scaler.unscale_(optimizer)            # 恢复真实梯度
torch.nn.utils.clip_grad_norm_(...)   # 梯度裁剪
scaler.step(optimizer)                # 更新参数
scaler.update()                       # 更新缩放因子

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

# 配置：accumulation_steps = 8
# 效果：实际 batch_size = 32 * 8 = 256

for step, (X, Y, loss_mask) in enumerate(train_loader):
    with ctx:
        res = model(X)
        loss = compute_loss(res, Y, loss_mask)
        loss = loss / args.accumulation_steps  # 损失除以累积步数

    scaler.scale(loss).backward()  # 梯度累积（不清零）

    if (step + 1) % args.accumulation_steps == 0:
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)  # 累积完毕后才清零

1
2
3

def get_lr(current_step, total_steps, lr):
    # lr_min = lr/10, lr_max = lr
    return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

# 启动命令
# torchrun --nproc_per_node 2 train_pretrain.py

# 初始化进程组
dist.init_process_group(backend="nccl")
ddp_local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(f"cuda:{ddp_local_rank}")

# 包装模型
model = DistributedDataParallel(model, device_ids=[ddp_local_rank])

# 使用分布式采样器
train_sampler = DistributedSampler(train_ds)
train_loader = DataLoader(train_ds, sampler=train_sampler, ...)

1
2
3
4
5
6

# 在 Attention.forward 中：
if past_key_value is not None:
    # 将新的 K, V 与缓存拼接
    xk = torch.cat([past_key_value[0], xk], dim=1)
    xv = torch.cat([past_key_value[1], xv], dim=1)
past_kv = (xk, xv) if use_cache else None

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

class PretrainDataset(Dataset):
    def __init__(self, data_path, tokenizer, max_length=512):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.samples = self.load_data(data_path)

    def load_data(self, path):
        samples = []
        with open(path, 'r', encoding='utf-8') as f:
            for line in f:
                data = json.loads(line.strip())
                samples.append(data)
        return samples

    def __getitem__(self, index):
        sample = self.samples[index]

        # Tokenize：将文本编码为 token IDs
        encoding = self.tokenizer(
            str(sample['text']),
            max_length=self.max_length,
            padding='max_length',      # 填充到 max_length
            truncation=True,           # 超长截断
            return_tensors='pt'
        )
        input_ids = encoding.input_ids.squeeze()

        # 构建 loss_mask：非 padding 位置为 True
        loss_mask = (input_ids != self.tokenizer.pad_token_id)

        # 构建自回归训练对
        X = torch.tensor(input_ids[:-1], dtype=torch.long)       # 前 N-1 个 token
        Y = torch.tensor(input_ids[1:], dtype=torch.long)        # 后 N-1 个 token
        loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long) # 对应 Y 的 mask

        return X, Y, loss_mask

1
2
3
4
5
6

# input_ids: [101, 7592, 2088, 102, 0, 0, 0]  （0 是 pad_token_id）
# loss_mask: [  1,    1,    1,   1, 0, 0, 0]

# 损失计算
loss = loss_fct(logits.view(-1, vocab_size), Y.view(-1)).view(Y.size())
loss = (loss * loss_mask).sum() / loss_mask.sum()  # 只计算非 padding 位置的平均损失

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

# 加载模型和分词器
model = NanoQwenForCausalLM(config)
model.load_state_dict(torch.load(checkpoint_path))
model.eval()

# 构建输入
prompt = "人工智能技术正在"
new_prompt = tokenizer.bos_token + prompt
inputs = tokenizer(new_prompt, return_tensors="pt").to(device)

# 生成文本
generated_ids = model.generate(
    inputs["input_ids"],
    max_new_tokens=512,
    do_sample=True,              # 采样生成
    temperature=0.85,            # 温度参数
    top_p=0.85,                  # 核采样
    pad_token_id=tokenizer.pad_token_id,
    eos_token_id=tokenizer.eos_token_id,
    streamer=TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
)

1
2

temperature = 0.85  # 略低于 1.0，偏向确定性
top_p = 0.85        # 保留累积概率 85% 的 token

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

from fastapi import FastAPI
from vllm import LLM, SamplingParams

app = FastAPI(title="NanoQwen Pretrain API")

# 加载模型
model = LLM(
    model="./out/",
    trust_remote_code=True,
    tensor_parallel_size=1,
    gpu_memory_utilization=0.8,
    max_model_len=2048,
    dtype="bfloat16"
)

# 非流式接口
@app.post("/chat")
async def chat(request: ChatRequest):
    prompt = f"<|im_start|>user\n{request.message}<|im_end|>\n<|im_start|>assistant\n"
    sampling_params = SamplingParams(
        temperature=request.temperature,
        top_p=request.top_p,
        max_tokens=request.max_tokens
    )
    outputs = model.generate([prompt], sampling_params)
    return ChatResponse(response=outputs[0].outputs[0].text)

# 流式接口
@app.post("/chat/stream")
async def chat_stream(request: ChatRequest):
    return StreamingResponse(generate_stream(...), media_type="text/plain")

1
2

python api_server.py
# 服务运行在 http://0.0.0.0:8000

1
2
3
4
5
6

# 1. 创建 Conda 环境
conda create -n pretrain python=3.10
conda activate pretrain

# 2. 安装依赖
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

# 3. 训练分词器（可选，已提供预训练好的分词器）
cd scripts
python train_tokenizer.py

# 4. 启动预训练
cd trainer
python train_pretrain.py

# 5. 分布式训练（多卡）
torchrun --nproc_per_node 2 train_pretrain.py --ddp

# 6. 模型评估
python eval_model.py

# 7. 部署 API
python api_server.py

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

大模型预训练
├── 1. 技术全栈概览
│   ├── 预训练 -> 微调(SFT) -> 对齐(RLHF) -> 部署
│   └── 预训练 = 知识的上限
│
├── 2. 预训练原理
│   ├── 自回归语言建模
│   │   ├── 条件概率 P(x_t | x_1,...,x_{t-1})
│   │   └── X = ids[:-1], Y = ids[1:]
│   └── Teacher Forcing
│       ├── 训练用真实历史（并行）
│       └── 推理用模型预测（自回归）
│
├── 3. 数据构建
│   ├── 数据来源（多个开源数据集整合）
│   ├── 数据清洗（长度过滤、格式整理）
│   ├── PII 检测（规则+NER+大模型）
│   └── 数据去重（哈希/MinHash/N-gram）
│
├── 4. 分词器
│   ├── BPE 算法
│   ├── 特殊 Token（endoftext/im_start/im_end）
│   ├── 词表大小选择（6400 vs 150K+）
│   └── chat_template 对话模板
│
├── 5. 模型架构
│   ├── RMSNorm（无均值归一化）
│   ├── RoPE（旋转位置编码）
│   ├── GQA（分组查询注意力）
│   ├── SwiGLU（门控激活FFN）
│   ├── MOE（混合专家）
│   │   ├── 门控网络（路由选择）
│   │   ├── 路由专家 + 共享专家
│   │   └── 辅助损失（负载均衡）
│   └── 权重绑定（Embedding = LM Head）
│
├── 6. 训练技术
│   ├── 混合精度（AMP + GradScaler）
│   ├── 梯度累积
│   ├── 余弦退火学习率
│   ├── 梯度裁剪
│   ├── DDP 分布式训练
│   └── KV-Cache（推理优化）
│
├── 7. 数据集处理
│   ├── JSONL 格式
│   ├── Tokenize + Padding + Truncation
│   └── Loss Mask（忽略 padding）
│
├── 8. 评估与推理
│   ├── 文本接龙测试
│   └── 采样策略（Temperature + Top-p）
│
└── 9. API 部署
    ├── FastAPI（Web框架）
    ├── vLLM（推理引擎）
    └── 流式/非流式接口