一个微型大语言模型DIY训练部署（以Tiny9为例），打通 Transformer、MoE、预训练、微调、量化与部署的完整工程链路

LLM大语言模型发展到今天，已经形成很多成熟的产业链。但绝大部分开发者是有没有算力条件去DIY一个完整的小模型来理解整个LLM的生命周期和训练过程的。现在以Tiny9为例帮助理解整个模型从数据、预训练到最终运行的简单过程示意，帮助理解LLM的运行机制。

---

基础设施与环境准备

在开始之前，我们需要搭建一个标准的 GPU 训练环境。

以Tiny9模型为例，因为参数足够小。

硬件假设：4张 NVIDIA A100 (80GB) 或 8张 RTX 4090 (24GB)。

软件依赖：

# 创建虚拟环境
conda create -n tiny9 python=3.10
conda activate tiny9

# 安装核心依赖 (建议使用 PyTorch 2.1+ 以支持更好的 FlashAttention)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install transformers==4.38.0 datasets==2.17.0 accelerate==0.27.0
pip install trl==0.7.10 deepspeed==0.13.1 peft==0.8.0
pip install vllm==0.3.3 autoawq==0.2.0

---

一、 阶段一：数据工程（Data Engineering）

1. 训练专属 Tokenizer

Tiny9 需要自己的分词器。我们使用 tokenizers 库基于 BPE 算法训练。

**实施步骤与代码 (train_tokenizer.py)**：

from tokenizers import Tokenizer, models, pre_tokenizers, trainers
from datasets import load_dataset

# 1. 加载清洗后的纯文本数据 (例如维基百科或 CommonCrawl 子集)
dataset = load_dataset("text", data_files={"train": "data/corpus.txt"})

# 2. 初始化 BPE Tokenizer
tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)

# 3. 配置训练器 (词表大小设为 32000)
trainer = trainers.BpeTrainer(
    vocab_size=32000,
    special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"],
    min_frequency=2,
    limit_alphabet=6000
)

# 4. 训练并保存
tokenizer.train_from_iterator(dataset["train"]["text"], trainer=trainer)
tokenizer.save("tiny9_tokenizer.json")
print("Tokenizer 训练完成！")

2. 数据打包 (Sequence Packing)

痛点：如果直接把短文本喂给模型，大量的 <pad> 会浪费算力。
方案：将多条文本拼接，用 <eos> 隔开，并重新生成 Labels，确保模型只在真实 token 上计算 Loss。

graph LR
    A[Text 1] -->|加 eos| B[Text 1 ]
    C[Text 2] -->|加 eos| D[Text 2 ]
    B --> E[拼接]
    D --> E
    E --> F[截断/填充至 2048]
    F --> G[生成 Labels: 右移一位]

**核心逻辑代码 (data_packing.py)**：

def packing_examples(examples, max_seq_length=2048, eos_token_id=3):
    """将多条短文本打包成固定长度的序列"""
    packed_input_ids, packed_attention_mask, packed_labels = [], [], []
    
    current_ids, current_labels = [], []
    
    for text in examples["text"]:
        # 编码单条文本
        encoded = tokenizer(text, add_special_tokens=False)
        input_ids = encoded["input_ids"] + [eos_token_id]
        
        # 如果加上当前文本会超长，则截断并保存当前包
        if len(current_ids) + len(input_ids) > max_seq_length:
            # 截断
            remain_length = max_seq_length - len(current_ids)
            current_ids.extend(input_ids[:remain_length])
            # Labels: 预测下一个 token，所以 labels 是 input_ids 右移一位，第一个设为 -100 (忽略 loss)
            current_labels = current_ids[1:] + [-100] 
            
            packed_input_ids.append(current_ids)
            packed_labels.append(current_labels)
            packed_attention_mask.append([1] * max_seq_length)
            
            # 重置，把剩下的 input_ids 放入下一个包
            current_ids = input_ids[remain_length:]
        else:
            current_ids.extend(input_ids)
            
    return {
        "input_ids": packed_input_ids,
        "attention_mask": packed_attention_mask,
        "labels": packed_labels
    }

---

二、 阶段二：模型架构与预训练 (Pre-training)

1. 核心代码：实现 Tiny9 的 MoE 层

这是理解 MoE 的关键。我们将标准的 FFN 替换为 MoE。

import torch
import torch.nn as nn
import torch.nn.functional as F

class Expert(nn.Module):
    """单个专家网络 (简单的 MLP)"""
    def __init__(self, dim, hidden_dim):
        super().__init__()
        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
        self.act = nn.SiLU()

    def forward(self, x):
        return self.w2(self.act(self.w1(x)))

class MoELayer(nn.Module):
    """混合专家层 (MoE)"""
    def __init__(self, dim, hidden_dim, num_experts=8, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        
        # 初始化 N 个专家
        self.experts = nn.ModuleList([Expert(dim, hidden_dim) for _ in range(num_experts)])
        
        # 路由器 (Gate)
        self.gate = nn.Linear(dim, num_experts, bias=False)
        
        # 负载均衡辅助损失权重
        self.aux_loss_coef = 0.01 

    def forward(self, x):
        batch_size, seq_len, dim = x.shape
        x_flat = x.view(-1, dim) # (batch*seq, dim)
        
        # 1. 路由计算
        router_logits = self.gate(x_flat) # (batch*seq, num_experts)
        routing_probs = F.softmax(router_logits, dim=-1)
        
        # 2. Top-K 选择
        top_k_probs, top_k_indices = torch.topk(routing_probs, self.top_k, dim=-1)
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True) # 归一化
        
        # 3. 专家计算与加权求和
        output = torch.zeros_like(x_flat)
        for i in range(self.top_k):
            expert_idx = top_k_indices[:, i]
            prob = top_k_probs[:, i].unsqueeze(-1)
            
            # 将 token 分配给对应的专家计算 (这里用循环简化，实际工程会用 grouped GEMM 优化)
            for exp_id in range(self.num_experts):
                mask = (expert_idx == exp_id)
                if mask.any():
                    token_for_expert = x_flat[mask]
                    expert_output = self.experts[exp_id](token_for_expert)
                    output[mask] += prob[mask] * expert_output
                    
        # 4. 计算负载均衡 Loss (Auxiliary Loss)
        # 鼓励每个 expert 被选中的概率均匀，且路由概率均匀
        # 简化版实现：计算 expert 被选中的频率与路由概率的乘积
        # (实际 DeepSpeed/Megatron 中有更高效的实现)
        aux_loss = self._compute_aux_loss(routing_probs, top_k_indices)
        
        return output.view(batch_size, seq_len, dim), aux_loss

    def _compute_aux_loss(self, routing_probs, top_k_indices):
        # 计算每个 expert 的路由概率均值
        mean_routing_prob = routing_probs.mean(dim=0) 
        # 计算每个 expert 被选中的频率 (one-hot 后求均值)
        expert_mask = F.one_hot(top_k_indices, self.num_experts).float()
        mean_expert_freq = expert_mask.mean(dim=0).sum(dim=0) / self.top_k
        
        # 辅助损失 = 路由概率均值 * 频率均值 * num_experts
        return (mean_routing_prob * mean_expert_freq).sum() * self.num_experts * self.aux_loss_coef

2. 预训练实施步骤与 DeepSpeed 配置

预训练需要分布式。我们使用 DeepSpeed ZeRO-3 来切分模型状态。

ds_config_zero3.json:

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu", "pin_memory": true },
    "offload_param": { "device": "cpu", "pin_memory": true },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto"
  },
  "fp16": { "enabled": "auto", "loss_scale": 0, "initial_scale_power": 16 },
  "bf16": { "enabled": "auto" },
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": 1.0,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto"
}

训练启动命令 (train_pretrain.sh):

# 使用 accelerate 启动 DeepSpeed 训练
accelerate launch --num_processes 4 --multi_gpu --mixed_precision bf16 \
    --main_process_port 29500 \
    run_pretrain.py \
    --model_name_or_path ./tiny9_config \
    --dataset_name ./packed_data \
    --output_dir ./checkpoints/tiny9-pt \
    --deepspeed ds_config_zero3.json \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --learning_rate 3e-4 \
    --max_steps 100000 \
    --save_steps 2000 \
    --logging_steps 10

---

三、 阶段三：对齐与微调 (SFT & DPO)

预训练后的 Tiny9 只会“续写”，我们需要用 trl (Transformer Reinforcement Learning) 库让它“听话”。

1. 监督微调 (SFT)

数据格式：{"instruction": "你好", "output": "你好！我是Tiny9。"}

SFT 核心代码 (run_sft.py):

from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, SFTConfig
from datasets import load_dataset

# 1. 加载预训练好的 Tiny9 和 Tokenizer
model = AutoModelForCausalLM.from_pretrained("./checkpoints/tiny9-pt", torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained("./tiny9_tokenizer")
tokenizer.pad_token = tokenizer.eos_token

# 2. 格式化数据 (将 instruction 和 output 拼接成 chat 格式)
def formatting_prompts_func(example):
    text = f"<bos>### 指令:\n{example['instruction']}\n\n### 回答:\n{example['output']}<eos>"
    return {"text": text}

dataset = load_dataset("json", data_files="data/sft_data.json")["train"]
dataset = dataset.map(formatting_prompts_func)

# 3. 配置并启动 SFTTrainer
training_args = SFTConfig(
    output_dir="./checkpoints/tiny9-sft",
    per_device_train_batch_size=4,
    learning_rate=2e-5,
    num_train_epochs=3,
    logging_steps=10,
    save_strategy="epoch",
    bf16=True,
    report_to="tensorboard"
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
    max_seq_length=2048,
    # 关键：只对 output 部分计算 loss，忽略 instruction 部分
    # 在 formatting_prompts_func 中可以通过设置 labels 的特定位置为 -100 实现
)

trainer.train()

2. 直接偏好优化 (DPO)

SFT 后，模型可能产生“虽然格式正确，但人类不喜欢”的回答。DPO 通过对比 chosen (好回答) 和 rejected (坏回答) 来优化。

数据格式：{"prompt": "如何评价苹果？", "chosen": "苹果很好吃...", "rejected": "苹果是手机..."}

DPO 核心代码 (run_dpo.py):

from trl import DPOTrainer, DPOConfig

# 加载 SFT 后的模型作为 reference model 和 policy model
# (实际工程中，policy model 是正在训练的，ref model 是冻结的)
model = AutoModelForCausalLM.from_pretrained("./checkpoints/tiny9-sft")
ref_model = AutoModelForCausalLM.from_pretrained("./checkpoints/tiny9-sft")

training_args = DPOConfig(
    output_dir="./checkpoints/tiny9-dpo",
    beta=0.1, # 控制偏离 reference model 的程度
    learning_rate=5e-6,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    max_length=1024,
    max_prompt_length=512,
    bf16=True
)

trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    args=training_args,
    train_dataset=load_dataset("json", data_files="data/dpo_data.json")["train"],
    tokenizer=tokenizer,
)

trainer.train()

---

四、 阶段四：量化与推理部署 (Quantization & Inference)

训练好的 Tiny9 (BF16) 大约需要 2GB 显存。为了在边缘设备或高并发下部署，我们进行 INT4 量化。

1. 模型量化 (AWQ)

AWQ (Activation-aware Weight Quantization) 是目前效果最好的 PTQ（训练后量化）算法之一。

量化实施代码 (quantize_awq.py):

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = "./checkpoints/tiny9-dpo"
quant_path = "./checkpoints/tiny9-dpo-awq-int4"
quant_config = { "zero_point": True, "q_group_size": 128, "version": "GEMM" }

# 1. 加载模型和 tokenizer
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 2. 执行量化 (内部会自动进行校准 calibration)
print("开始 AWQ INT4 量化...")
model.quantize(tokenizer, quant_config=quant_config)

# 3. 保存量化后的模型
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)
print(f"量化完成，保存至 {quant_path}")

2. 推理部署 (vLLM)

使用 vLLM 部署高性能推理服务，支持 Continuous Batching 和 PagedAttention。

部署启动命令:

# 启动 OpenAI 兼容的 API 服务
python -m vllm.entrypoints.openai.api_server \
    --model ./checkpoints/tiny9-dpo-awq-int4 \
    --tokenizer ./checkpoints/tiny9-dpo-awq-int4 \
    --dtype float16 \
    --max-model-len 2048 \
    --gpu-memory-utilization 0.9 \
    --port 8000 \
    --host 0.0.0.0

客户端调用测试:

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="empty")

response = client.chat.completions.create(
    model="tiny9-dpo-awq-int4",
    messages=[{"role": "user", "content": "请用一句话解释什么是 MoE 架构？"}],
    temperature=0.7,
    max_tokens=100
)
print(response.choices[0].message.content)

---

五、 完整实施 Checklist (避坑指南)

在真正跑通这套流程时，请务必对照以下 Checklist：

📊 数据阶段

• 去重检查：是否使用了 MinHash 或 SimHash 去除重复文档？（重复数据会导致模型在特定 token 上 Loss 异常低，引发梯度爆炸）。
• Packing 检查：打包后的 labels 是否正确右移？跨文本拼接处是否用 <eos> 隔离？（防止模型把两篇不相关的文章当成连贯上下文）。

🏗️ 预训练阶段

• MoE 负载均衡：在 TensorBoard 中监控 aux_loss。如果 aux_loss 持续为 0 或极大，说明 Router 坍塌（所有 token 都去了同一个 Expert）或梯度消失。
• EP 通信检查：如果使用了 Expert Parallelism (EP)，检查 GPU 间的 All-to-All 通信时间占比。如果占比 > 20%，说明网络带宽是瓶颈，需调整 EP 和 TP 的比例。

🎯 对齐阶段 (SFT/DPO)

• Loss Masking：SFT 时，是否只对 output 计算 Loss？（如果对 instruction 也计算 Loss，模型会学到“如何提问”而不是“如何回答”）。
• DPO 数据质量：DPO 对数据质量极度敏感。chosen 和 rejected 的差异必须明显，且 rejected 不能是毫无逻辑的乱码（会导致模型学坏）。

🚀 部署阶段

• 量化掉点评估：INT4 量化后，务必在 MMLU、GSM8K 等榜单上跑一下评测。如果复杂推理任务掉点超过 5%，建议退回 INT8 或 FP16。
• KV Cache 显存监控：在 vLLM 中，监控 gpu_cache_usage_perc。如果经常达到 100%，说明并发请求过多，需要增加 max-model-len 限制或增加 GPU 节点。

---

延伸阅读

• PyTorch：Transformer 模块官方文档，理解基础 Transformer 组件与接口。

• Hugging Face Transformers：量化官方文档，了解 INT8、4-bit 与 QLoRA 的工程实现。

模型技术和算法革新经常在变，以 Tiny9 的模型名称与参数配置用于教学示例，文中的代码是逐步构建研发系统的基础实现，并非已完成训练、全面验证和性能优化的商业级模型发布包。