参数	值
参数量	~0.2B
词嵌入维度 (d_model)	1024
Transformer 层数	16
注意力头数 (Q 头)	16
KV 头数 (GQA)	4
词表大小	32128
最大序列长度	1024
激活函数	SwiGLU
位置编码	RoPE
归一化	RMSNorm
混合精度训练	BF16

以上是默认配置，代码里所有超参数都集中在 config_zh.yaml，一改全改，随便魔改。

架构亮点

组件	实现方式
RoPE	复数乘法旋转，`precompute_rope_operators` 预计算全部位置频率矩阵，零重复计算
GQA	Q=16头，KV=4头，推理时 KV Cache 显存缩减 4×
Flash Attention	PyTorch 2.0 原生 `F.scaled_dot_product_attention`，无第三方依赖
SwiGLU	FFN 隐藏层维度对齐 `multiple_of=256`，适配硬件 GEMM 最优 Tile
RMSNorm	Pre-Norm 结构，省去均值中心化，训练更稳定
权重共享	Embedding 与 Unembedding 层权重绑定，节省约 30% 参数量
KV Cache	推理前预分配全零张量，原地写入，无显存碎片化
MoE（可选）	4专家 Top-2 路由，代码保留但当前版本 `use_moe=False`——在 0.2B 规模下 dense 效果更优

快速开始

只想跑推理？按顺序执行下面 5 步即可，全程终端操作，无需修改任何代码。

Step 1 — 克隆仓库

bash

git clone https://github.com/Tianyu-Zhou1964/PIE-Handmaking_LLM.git
cd PIE-Handmaking_LLM

Step 2 — 安装 Python 依赖

bash

# 按你的 CUDA 版本选（推理也可以用 CPU，跳过 --index-url 直接 pip install torch）
pip install torch --index-url https://download.pytorch.org/whl/cu121   # CUDA 12.1

pip install -r requirements.txt

Step 3 — 安装 Rust 并编译 BPE 引擎

分词器核心用 Rust + PyO3 实现，需要本地编译，不在 PyPI 上。

Linux / macOS：

bash

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
source $HOME/.cargo/env

Windows（PowerShell）：

powershell

winget install Rustlang.Rustup
# 或者去 https://rustup.rs 下载 rustup-init.exe 双击安装

安装完 Rust 后，编译 Python 可调用的 .so / .pyd：

bash

pip install maturin

cd Chinese/src/custom_bpe/
maturin develop --release
cd ../../..

编译成功后在该目录下执行 python -c "import custom_bpe; print('ok')" 验证。

Step 4 — 下载预训练权重

bash

# 方式一：huggingface-hub（推荐，自动断点续传）
pip install huggingface_hub
python -c "
from huggingface_hub import hf_hub_download
hf_hub_download(
    repo_id='Tianyu-Zhou/PIE1.0-0.2B-dense-base',
    filename='PIE-0.2B-dense.pth',
    local_dir='./Checkpoint'
)
"

# 方式二：git lfs（需要提前安装 git-lfs）
git lfs install
git clone https://huggingface.co/Tianyu-Zhou/PIE1.0-0.2B-dense-base ./hf_model
cp hf_model/PIE-0.2B-dense.pth Checkpoint/

# 国内用户：从 ModelScope 下载
pip install modelscope
python -c "
from modelscope.hub.snapshot_download import snapshot_download
snapshot_download('Zaoshangzhou/PIE1.0-0.2B-dense-base', cache_dir='./Checkpoint')
"

下载完成后确认文件存在：

bash

ls Checkpoint/
# 应该看到 PIE-0.2B-dense.pth

Step 5 — 启动推理

bash

cd Chinese/src/
python inference.py

支持 CUDA / Apple MPS / CPU，device="auto" 自动检测，无需手动指定。

采样参数（temperature、top_k、top_p、repetition_penalty 等）统一在 Chinese/config_zh.yaml 中调整。

配套教程

这个项目是抖音系列课程的代码实现：👉 前往作者主页

系列	内容	进度
「手撕大模型」	理论篇：Attention → RoPE → MoE，每集讲透一个机制	18 集，完结
「手搓大模型」	实战篇：Rust BPE 分词器 → 完整多卡训练 base 模型	5 大集，完结

代码注释里的 Ep1、Ep9、Ep15 直接对应视频集数，看代码时可以随时跳转理论讲解。

联系方式

📧 tyro1964@gmail.com
📱 抖音私信

开源协议

Apache-2.0

「你不需要很厉害才能开始，但你需要开始才能变得很厉害。」