实战：用 Colossal-LLaMA 训中文增强模型

Colossal-LLaMA 是 ColossalAI 团队 2023-2024 年开源的项目， 完整演示了如何在英文模型基础上做中文增强 + 二次预训练 + SFT。 本章按 applications/Colossal-LLaMA/ 目录把这个流程跑一遍。

8.1项目结构

applications/Colossal-LLaMA/
├── README.md
├── prepare_pretrain_dataset.py     ★ 准备预训练数据
├── prepare_sft_dataset.py           ★ 准备 SFT 数据
├── train.py                         ★ 主训练脚本（PT/CPT 共用）
├── inference_example.py
├── stream_chat_example.py
├── colossal_llama/
│   ├── dataset/
│   │   ├── conversation.py          中文 prompt 模板
│   │   └── loader.py                streaming dataset
│   ├── tokenizer/
│   │   ├── init_tokenizer.py        ★ 词表扩展
│   │   └── chat_template.py
│   └── utils/
│       ├── ckpt_io.py
│       └── flash_attention_patch.py
└── docs/

8.2step 1：扩中文词表

原版 LLaMA tokenizer 用 BPE，词表 32K。每个中文字符通常被切成 2–3 个 token，处理中文极其低效。Colossal-LLaMA 用 SentencePiece 在中文语料上训练一个 30K 的中文 tokenizer，再合并到 LLaMA 原表里。

看 colossal_llama/tokenizer/init_tokenizer.py：

# 1) 先在中文语料训一个 SentencePiece
python -m sentencepiece.SentencePieceTrainer \
    --input=chinese.txt \
    --model_prefix=chinese_sp \
    --vocab_size=30000 \
    --character_coverage=0.9995 \
    --model_type=bpe

# 2) 合并到 LLaMA 原 tokenizer
python init_tokenizer.py \
    --source_tokenizer_dir Llama-3-8B \
    --target_tokenizer_dir colossal_llama_tokenizer \
    --expand_tokens_file chinese_sp.model

扩完后词表从 32K → ~58K，每个汉字平均只占 1 token，中文输入序列长度减少 65%。

为什么扩词表很关键

$$\text{每 token 训练 cost} \propto \text{seqlen}^2 \;\Longrightarrow\; \text{seqlen 减 65\%, cost 减 } 1 - 0.35^2 = 87.75\%$$

（attention 是 O(n²)，中文场景这一项收益巨大。）

8.3step 2：embedding 扩展 + 初始化

扩了词表后，模型的 embed_tokens 和 lm_head 尺寸必须改。看 train.py 里的处理：

old_size = model.config.vocab_size           # 128256 (Llama-3)
new_size = len(tokenizer)                     # ~154256
model.resize_token_embeddings(new_size)

# 用旧 embedding 的平均值初始化新加的 token
with torch.no_grad():
    mean = model.get_input_embeddings().weight[:old_size].mean(dim=0)
    model.get_input_embeddings().weight[old_size:] = mean
    mean_out = model.get_output_embeddings().weight[:old_size].mean(dim=0)
    model.get_output_embeddings().weight[old_size:] = mean_out

注意：不要用随机初始化，那样新 token 在第一步训练时 loss 会非常高，模型可能直接发散。

8.4step 3：准备 streaming 预训练数据

看 prepare_pretrain_dataset.py，主要做三件事：

1. 读 jsonl / parquet 原始语料；
2. 用扩好的 tokenizer 切；
3. 按 max_length packing 成 sample，落盘到 indexed_dataset 格式（与 Megatron 兼容）。

python prepare_pretrain_dataset.py \
    --data_input_dirs raw_zh/ raw_code/ raw_en/ \
    --tokenizer_dir colossal_llama_tokenizer \
    --data_output_dirs ./train_data \
    --max_length 4096 \
    --num_spliced_dataset_bins 64

跑完 ./train_data/ 下会有 64 个 bin 文件，训练时 streaming 读取。

8.5step 4：启动二次预训练（CPT）

看 train.py 关键参数：

colossalai run --nproc_per_node 8 train.py \
    --pretrained meta-llama/Meta-Llama-3-8B \
    --tokenizer_dir colossal_llama_tokenizer \
    --dataset ./train_data \
    --plugin hybrid_parallel \
    --tp 2 --pp 1 \
    --zero 1 \
    --batch_size 1 \
    --num_epochs 1 \
    --max_length 4096 \
    --lr 2e-5 \
    --mixed_precision bf16 \
    --warmup_steps 100 \
    --save_interval 1000 \
    --save_dir ./outputs/cpt \
    --use_flash_attn \
    --pad_token unk

训练循环简化版：

colossalai.launch_from_torch()
plugin = HybridParallelPlugin(tp_size=2, pp_size=1, zero_stage=1, precision="bf16",
                              enable_flash_attention=True, enable_fused_normalization=True)
booster = Booster(plugin=plugin)

model = AutoModelForCausalLM.from_pretrained("Meta-Llama-3-8B", torch_dtype=torch.bfloat16)
model.resize_token_embeddings(len(tokenizer))
optimizer = HybridAdam(model.parameters(), lr=2e-5)
scheduler = ...

model, optimizer, _, dl, scheduler = booster.boost(model, optimizer, None, dl, scheduler)

for epoch in range(num_epochs):
    for batch in dl:
        if pp_size > 1:
            booster.execute_pipeline(iter([batch]), model,
                                     criterion=lambda o, b: o.loss,
                                     optimizer=optimizer, return_loss=True)
        else:
            loss = model(**batch).loss
            booster.backward(loss, optimizer)
        optimizer.step()
        optimizer.zero_grad()
        scheduler.step()

8.6step 5：SFT

CPT 训完后做 SFT。Colossal-LLaMA 的 SFT 数据格式（chat 模板）：

{"messages": [
    {"from": "human", "content": "请介绍 ColossalAI"},
    {"from": "assistant", "content": "ColossalAI 是..."}
]}

chat template 在 colossal_llama/tokenizer/chat_template.py 定义，关键点：对 user 部分的 token 做 label_mask = -100（不算 loss），只学 assistant 输出。

python prepare_sft_dataset.py \
    --data_input_dirs raw_sft.jsonl \
    --tokenizer_dir colossal_llama_tokenizer \
    --data_output_dirs ./sft_data \
    --conversation_template colossal_llama \
    --max_length 4096

colossalai run --nproc_per_node 8 train.py \
    --pretrained ./outputs/cpt/epoch-1-step-final \
    --dataset ./sft_data \
    --plugin hybrid_parallel \
    --tp 2 --zero 1 \
    --lr 2e-5 --num_epochs 3 \
    --save_dir ./outputs/sft

8.7资源估算

8B 模型、bf16、TP=2、ZeRO-1、seqlen=4096、micro_bs=1：

8.8常见问题

8.9这章你需要带走的

• 中文增强 = 扩词表 → embedding resize → CPT → SFT 四步走；
• 词表扩展用 SentencePiece + 与 LLaMA 原表合并，注意 embedding 用均值初始化；
• CPT 用 HybridParallelPlugin(tp=2, zero=1)，bf16，开 flash + fused norm；
• SFT 阶段必须给 user token 打 label_mask = -100；
• 8×A100-40GB 跑 Llama-3-8B CPT 约 150 samples/s（4K seq）；
• 训完用 booster.save_model(sharded_strategy='hf') 直接得到 HF 兼容 ckpt。