这是一个专门用于中文语言处理的大型语言模型实现，使用纯 Rust 构建，不依赖任何外部的机器学习框架。完全基于 ndarray 实现矩阵运算，采用现代 Pre-LN Transformer 架构（GPT-2 标准）。
本项目展示了如何在 Rust 中从零开始构建专门处理中文的 Transformer 语言模型，包括：
- 现代 Pre-LN Transformer 架构 - GPT-2/3 标准，具有明确的残差连接
- 中文预训练：在中文事实知识文本上进行预训练
- 中文指令微调：针对中文对话场景进行微调
- 中文交互聊天模式：支持中文交互式对话
- 完整反向传播：包含梯度裁剪和 Adam 优化器
- 模块化架构：清晰的关注点分离
- 中文优化分词器：使用 jieba-rs 进行中文分词，全局单例优化（快 50-70%）
- 多头自注意力机制（8 头）更好地理解中文语法
- 上下文窗口管理：保持对话历史记录
- 高级解码方法（top-k/top-p 采样，束搜索，温度缩放）
- 正则化技术（Dropout，层归一化）提升稳定性
- 性能监控：详细的定时和性能分析
这不是一个生产级别的大语言模型，距离大型模型还很远。
这只是一个演示项目，展示了这些模型在底层的工作原理。
- ✅ 训练公开入口收缩：当前推荐公开入口仅保留
train(...)、train_monitored(...)、train_bucketed_sequential(...)、train_with_checkpointing(...) - ✅ 共享训练单步：训练主线围绕
PreparedTrainingStep、prepare_training_step(...)、backward_with_ctx(...)组织，避免重复复制单样本训练逻辑 - ✅ Loss 口径统一：公开训练 API 统一使用 token-weighted mean loss，不同入口输出的
Loss/PPL现在语义一致 - ✅ 批次语义校准：
train_bucketed_sequential(...)表示“按批次组织样本 + 批内逐样本顺序更新”，不是严格的 batch 梯度平均更新
- ✅ Batch/PAD mask 接线：
SelfAttention/TransformerBlock的forward_batch(...)会消费attention_mask（Key padding mask），为后续真正的 batch 向量化训练提供正确性门槛 - ✅ 消除
PAD_ID==0硬编码：padding/loss/grad 以vocab.pad_token_id()为单一事实源；BatchLoader支持注入pad_token_id - ✅ 训练信号护栏：shape mismatch / target 越界 / 全 PAD 时跳过 optimizer step，避免 Adam 动量“0 梯度漂移”，并提升可诊断性
- ✅ 新增/更新测试：
self_attention_forward_batch_mask_test、compute_gradients_step_test
- 🚀 检查点管理器 - 支持Best/Last/周期性保存策略，自动清理旧检查点
- ✅ 完整状态保存 - 模型参数 + Adam优化器状态（m, v, timestep），支持断点续训
- ✅ 早停集成 - 自动保存最佳检查点，触发早停时自动回滚到最佳状态
- ✅ Resume训练 - 从检查点恢复训练，支持中断续训，保持训练连续性
- ✅ CLI参数支持 -
--resume,--resume-from,--checkpoint-dir等完整参数支持 - ✅ 集成测试 - 验证保存/恢复后loss连续性（loss差异 < 0.1%），确保优化器状态正确恢复
- 🚀 阶段1训练优化 - 引入训练监控、学习率调度与缓存策略
- ✅ 数据预处理缓存 - 避免重复 tokenization
- ✅ 余弦退火学习率 - 后续主训练入口已收敛为 warmup + cosine（默认无重启）
- ✅ 早停机制 - 自动检测收敛
- ✅ 增强训练监控 - Loss, PPL, LR, Grad, Speed, ETA 监控
- ✅ 梯度累积 - 支持 micro-batch 累积，当前由
train_monitored(...)控制
- ✅ Pre-LN Transformer 架构 - 从 Post-LN 升级到 Pre-LN (GPT-2 标准) 以获得更好的训练稳定性
- ✅ 明确的残差连接 - 从子层移动到 TransformerBlock 以提高清晰度
- ✅ 移除语义增强器 - 通过删除未经验证的实验特性来简化模型
- ✅ 性能优化 - Jieba 单例优化 (快 50-70%)，注意力重塑优化 (快 20-30%)
- ✅ 编译器优化 - LTO，opt-level 3，codegen-units 1 用于发布版本构建
- ✅ 性能监控 - 添加全面的性能跟踪和分析
从以下核心文件开始了解实现：
src/main.rs- 训练流水线、数据准备和交互模式src/llm.rs- 核心 LLM 实现和训练逻辑src/transformer.rs- Pre-LN Transformer 块，具有明确的残差连接
模型使用 Pre-LN Transformer 架构（GPT-2 标准），包含以下组件：
输入文本 → Jieba 分词 → 词嵌入 + 位置编码
↓
[2个 Transformer 块]
每个块包含：
• 层归一化 → 多头注意力 (8 头) → Dropout → 残差连接
• 层归一化 → 前馈网络 → Dropout → 残差连接
↓
输出投影 → Softmax → 词预测
Pre-LN（子层前的层归一化）是现代 GPT-2、GPT-3 及后续版本使用的标准：
- ✅ 更稳定的训练 - 更好的梯度流动
- ✅ 更快的收敛 - 减少梯度消失/爆炸
- ✅ 更稳定 - 对学习率不那么敏感
架构对比：
Post-LN (旧版): Pre-LN (当前 - GPT-2 标准):
输入 输入
↓ ↓
注意力 层归一化
↓ ↓
层归一化 注意力
↓ ↓
Dropout Dropout
↓ ↓
(+输入) (+输入) ← 明确的残差
↓ ↓
前馈网络 层归一化
↓ ↓
层归一化 前馈网络
↓ ↓
Dropout Dropout
↓ ↓
输出 (+X) ← 明确的残差
↓
输出
src/
├── main.rs # 🎯 训练流水线和交互模式
├── llm.rs # 🧠 核心 LLM 实现和训练逻辑
├── lib.rs # 📚 库导出和常量
├── transformer.rs # 🔄 Pre-LN Transformer 块，具有明确的残差连接
├── self_attention.rs # 👀 多头自注意力机制 (8 头)
├── feed_forward.rs # ⚡ 逐位置前馈网络
├── embeddings.rs # 📊 词嵌入层，包含位置编码
├── output_projection.rs # 🎰 最终线性层，用于词汇预测
├── vocab.rs # 📝 词汇管理与优化的 jieba-rs 分词
├── layer_norm.rs # 🧮 层归一化 (可学习的 γ 和 β)
├── dropout.rs # 🚫 Dropout 正则化 (10% 率，反向 Dropout)
├── position_encoding.rs # 📍 正弦位置编码
├── adam.rs # 🎓 Adam 优化器 (β₁=0.9, β₂=0.999)
├── performance_monitor.rs # ⏱️ 性能分析和定时
└── dataset_loader.rs # 📁 训练数据加载
实现包括两个专门针对中文的训练阶段：
-
中文预训练：从中文事实陈述中学习中文世界知识
- "太阳从东方升起，在西方落下"
- "水由于重力而从高处流向低处"
- "山脉是高大而多岩石的地形"
- 增强了中文文化知识、成语和历史事实
-
中文指令微调：学习中文对话模式
- "用户：山脉是如何形成的？助手：山脉通过构造力或火山活动在长时间的地质时期内形成..."
- 处理中文问候、解释和后续问题
- 包含中文文化引用和成语
# 克隆并运行
git clone https://github.com/H-Chris233/RustGPT-Chinese.git
cd RustGPT-Chinese
cargo run
# 模型将：
# 1. 从中文训练数据构建词汇表（使用jieba-rs分词）
# 2. 在中文事实陈述上进行预训练（500 轮，带早停）
# 3. 在中文对话数据上进行指令微调（500 轮，带早停）
# 4. 自动保存检查点（最佳模型 + 最新模型）
# 5. 进入中文交互模式进行测试# 正常训练（自动保存检查点）
cargo run
# 从检查点恢复训练（自动查找最佳或最新检查点）
cargo run -- --resume
# 从指定检查点恢复训练
cargo run -- --resume --resume-from=checkpoints/checkpoint_best_epoch_50_loss_2.3456.bin
# 自定义resume参数
cargo run -- --resume --epochs=1000 --lr=0.0001 --patience=50 --checkpoint-dir=my_checkpoints
# 快速测试模式（仅预训练，无检查点）
cargo run -- --quick --pretrain-epochs=30 --lr=0.0001 --patience=10| 参数 | 说明 | 默认值 |
|---|---|---|
--resume |
启用resume训练模式 | - |
--resume-from=<path> |
指定检查点文件路径 | 自动查找 |
--checkpoint-dir=<dir> |
检查点保存/加载目录 | checkpoints |
--epochs=<n> |
训练的最大epoch数 | 500 |
--lr=<f> |
学习率 | 从检查点继承 |
--patience=<n> |
早停patience | 30 |
--quick |
快速测试模式（仅预训练） | - |
--pretrain-epochs=<n> |
预训练epoch数 | 30 |
--freeze-attn |
冻结注意力层参数更新 | - |
--no-interactive |
跳过交互模式 | - |
checkpoints/
├── checkpoint_best_epoch_42_loss_2.1234.bin # 最佳模型检查点
├── checkpoint_best_epoch_42_loss_2.1234.json # 元数据（可读）
├── checkpoint_last.bin # 最新模型检查点
├── checkpoint_last.json # 元数据（可读）
└── model_final.bin # 训练完成后的最终模型
检查点包含：
- ✅ 模型参数：所有层的权重和偏置
- ✅ Adam优化器状态：一阶矩m、二阶矩v、timestep
- ✅ 训练元数据：epoch、loss、学习率、时间戳、训练阶段
- ✅ 词汇表：完整的token到ID映射
自动保存策略：
- BestAndLast策略（默认）：同时保存最佳模型和最新模型
checkpoint_best_*.bin: 当loss改善时保存，保留最佳3个检查点checkpoint_last.bin: 每个epoch都更新，确保能从最新状态恢复
- Periodic策略：每N个epoch保存一次，适合长时间训练
早停集成：
- 训练时自动跟踪最佳loss和对应的epoch
- 触发早停时，自动从最佳检查点回滚模型参数
- 避免过拟合，确保使用最优模型参数
Resume训练保证：
- 保存完整的Adam优化器状态（m, v, timestep）
- 加载后loss完全一致（经过集成测试验证，差异 < 0.1%）
- 训练连续性得到保证，学习率调度正确恢复
- 支持跨训练阶段恢复（预训练 → 指令微调）
训练完成后，可以交互式地测试中文模型：
输入提示：山脉是如何形成的？
模型输出：山脉通过构造力或火山活动在长时间的地质时期内形成
输入提示：降雨的原因是什么？
模型输出：降雨是由云中的水蒸气凝结成水滴，当水滴变得太重而无法悬浮在空气中时形成的
| API | 输入 | 适用场景 | 说明 |
|---|---|---|---|
train(...) |
原始文本 | 最基础教学入口 | 展示最小训练闭环，适合入门阅读 |
train_monitored(...) |
原始文本 | 推荐主训练入口 | 带学习率调度、早停、监控与梯度累积 |
train_bucketed_sequential(...) |
原始文本 | 进阶训练入口 | 使用 bucketing / padding mask 组织批次，但参数更新仍按批内逐样本顺序执行 |
train_with_checkpointing(...) |
预 tokenized 数据 | 可恢复训练 / 长训练任务 | 组合检查点保存、恢复、早停与训练连续性保障 |
说明：当前项目把公开训练指标统一为 token-weighted mean loss，以避免不同训练入口之间出现不可比的
Loss/PPL。
- 词汇表大小: 动态（基于训练数据构建，集成jieba-rs）
- 嵌入维度: 256（见
EMBEDDING_DIM） - 隐藏维度: 512（见
HIDDEN_DIM） - 最大序列长度: 128（见
MAX_SEQ_LEN） - 架构: 2 Pre-LN Transformer 块 + 嵌入 + 输出投影（见
src/main.rs默认构建） - 总参数: 与
vocab_size线性相关（主要由 Embeddings / OutputProjection 主导）；运行时会打印llm.total_parameters()
- 优化器: Adam (β₁=0.9, β₂=0.999, ε=1e-8) 带梯度裁剪
- 学习率调度: 余弦退火 + Warmup（无重启，见
LLM::cosine_with_warmup_lr） - 预训练/微调默认 LR: 0.0001（
src/main.rs默认训练入口） - 损失函数:
log_softmax+ NLL（从 log_probs 计算交叉熵），pad_token_id不参与 loss/grad - 梯度裁剪: L2 范数限制在 1.0（训练主路径默认）
- 正则化: Dropout 层，10% 率（反向 Dropout）
- 现代 Pre-LN Transformer - GPT-2/3 标准架构，用于稳定训练
- 明确的残差连接 - 清晰且可维护的架构
- 优化的中文分词 - jieba-rs 与全局单例 (快 50-70%)
- 多头自注意力 - 8 头，带优化的重塑操作 (快 20-30%)
- 高级解码方法:
- 贪心解码 (argmax)
- Top-k 采样 (核采样)
- Top-p 采样 (累积概率)
- 束搜索与对数概率
- 温度缩放输出多样性
- 梯度裁剪 - L2 范数用于训练稳定性
- 模块化层系统 - 清晰接口与 Layer trait
- 全面的测试覆盖 - 所有组件的单元测试
- 上下文窗口管理 - 对话历史的滑动窗口
- 性能监控 - 详细的定时和性能分析工具
- 编译器优化 - LTO, opt-level 3, 单代码生成单元
| 优化 | 加速 | 状态 |
|---|---|---|
| Jieba 单例 (OnceLock) | 50-70% | ✅ 已实现 |
| 注意力重塑 (切片操作) | 20-30% | ✅ 已实现 |
| 编译器优化 (LTO) | 10-20% | ✅ 已实现 |
可选 BLAS 加速（--features blas） |
依赖环境 | ✅ 已实现 |
| 总计预期提升 | 60-80% | ✅ 已实现 |
# 运行所有测试
cargo test
# 测试特定组件
cargo test --test llm_test
cargo test --test transformer_test
cargo test --test self_attention_test
cargo test --test chinese_tests
cargo test --test vocab_test
# 构建优化版本
cargo build --release
# 运行详细输出
cargo test -- --nocapture
# 格式化代码
cargo fmt
# 运行 linter
cargo clippy为获得最大性能，请使用发布模式：
cargo build --release
./target/release/llm发布模式启用了：
- 链接时优化 (LTO) - 跨箱内联
- 最大优化级别 (opt-level 3)
- 单代码生成单元 - 更好的优化机会
- 预期加速: 比调试模式快 10-20%
此实现展示了中文大语言模型的关键机器学习概念：
- Transformer 架构（注意力、前馈、层归一化）
- 通过神经网络的反向传播
- 中文语言模型训练（预训练 + 微调）
- 中文分词和词汇管理，使用jieba-rs
- 基于梯度的 Adam 优化
- 上下文管理以维护对话历史
- 正则化技术以提高稳定性
是了解现代 LLM 在底层如何工作的完美选择！
ndarray- 用于矩阵运算的 N 维数组（可选 BLAS：cargo build --features blas）jieba-rs- 中文文本分词和处理lru- 分词缓存（减少重复 tokenization）rand- 随机数生成（初始化权重等）regex- 正则表达式匹配中文成语识别serde+serde_json- 序列化（元数据/导出）bincode- 二进制序列化（checkpoint）log+simple_logger- 日志
没有 PyTorch、TensorFlow 或 Candle - 只有纯 Rust 和线性代数！
欢迎贡献！这个项目非常适合学习和实验。
- 🧮 真正的 batch 向量化训练 - 基于
forward_batch/backward_batch做批量矩阵运算，并实现“累积/平均梯度后一次 optimizer step”（当前train_bucketed_sequential仍是 batch 内顺序 SGD） - 📊 评估指标 - 困惑度，基准测试，训练可视化
- 🎯 注意力可视化 - 可视化中文文本的注意力模式
- 📈 训练曲线 - 损失/准确率绘图
- 高级架构 (旋转位置嵌入 (RoPE)，Flash Attention)
- 训练改进 (严格 batch 更新语义，混合精度，更完善的日志/诊断)
- 中文数据处理 (更大的中文数据集，流式数据加载)
- 模型分析 (注意力可视化，梯度分析，可解释性)
- ✅ Pre-LN Transformer - 现代 GPT-2 标准架构
- ✅ 明确的残差连接 - 清晰且可维护
- ✅ 性能优化 - 比初版快 60-80%
- ✅ Padding mask 支持 -
forward_with_padding_mask(...)/forward_batch(..., attention_mask)（训练默认单序列无需 padding） - ✅ 梯度累积 - 可配置（默认禁用以提升稳定性）
- ✅ Warmup + 余弦退火 - 默认训练入口启用 warmup（无重启）
- Fork 仓库
- 创建功能分支：
git checkout -b feature/vectorized-batch - 进行更改并添加测试
- 运行测试套件：
cargo test - 格式化和检查：
cargo fmt && cargo clippy - 提交 PR 并提供清晰的描述
- 遵循标准 Rust 约定 (
cargo fmt) - 为新功能添加全面测试
- 根据需要更新文档和 README
- 保持"从零开始"的哲学 - 避免重量级 ML 依赖
- 专注于中文语言处理改进
- 为复杂算法添加解释注释
- 🚀 初学者: 更多中文训练数据，配置文件，补测试/修文档
- 🔥 中级: 真正的 batch 向量化训练，评估指标，训练可视化
- ⚡ 高级: Flash Attention，RoPE，混合精度训练
有问题？开一个 issue 或开始讨论！
此项目是开源的，可用于教育目的。
使用 🦀 Rust 和 ❤️ 构建，用于理解中文大语言模型
没有 PyTorch、TensorFlow 或 Candle - 只有纯 Rust 和线性代数！