用 Rust 从头实现 Transformer 模型

完整的大语言模型实现

纯 Rust 编写，无外部 ML 框架依赖。从零开始构建，仅使用 ndarray 进行矩阵运算。

本项目演示了如何用 Rust 从零开始构建基于 transformer 的语言模型，包括：

• 基于事实文本补全的预训练
• 对话 AI 的指令微调
• 交互式聊天模式用于测试
• 完整的反向传播和梯度裁剪
• 模块化架构，关注点清晰分离

这不是生产级 LLM，与大型模型相差甚远。

这只是一个演示这些模型内部工作原理的玩具项目。

🔍 关键文件探索

从这两个核心文件开始了解实现：

src/main.rs - 训练管道、数据准备和交互模式

src/llm.rs - 核心 LLM 实现，包含前向/反向传播和训练逻辑

该模型使用基于 transformer 的架构，包含以下组件：

Input Text → Tokenization → Embeddings → Transformer Blocks → Output Projection → Predictions

src/
├── main.rs              # 🎯 Training pipeline and interactive mode
├── llm.rs               # 🧠 Core LLM implementation and training logic
├── lib.rs               # 📚 Library exports and constants
├── transformer.rs       # 🔄 Transformer block (attention + feed-forward)
├── self_attention.rs    # 👀 Multi-head self-attention mechanism
├── feed_forward.rs      # ⚡ Position-wise feed-forward networks
├── embeddings.rs        # 📊 Token embedding layer
├── output_projection.rs # 🎰 Final linear layer for vocabulary predictions
├── vocab.rs            # 📝 Vocabulary management and tokenization
├── layer_norm.rs       # 🧮 Layer normalization
└── adam.rs             # 🏃 Adam optimizer implementation

tests/
├── llm_test.rs         # Tests for core LLM functionality
├── transformer_test.rs # Tests for transformer blocks
├── self_attention_test.rs # Tests for attention mechanisms
├── feed_forward_test.rs # Tests for feed-forward layers
├── embeddings_test.rs  # Tests for embedding layers
├── vocab_test.rs       # Tests for vocabulary handling
├── adam_test.rs        # Tests for optimizer
└── output_projection_test.rs # Tests for output layer

🧪 模型学习的内容

实现包括两个训练阶段：

预训练：从事实陈述中学习基本世界知识

• "The sun rises in the east and sets in the west"
• "Water flows downhill due to gravity"
• "Mountains are tall and rocky formations"

指令微调：学习对话模式

• "User: How do mountains form? Assistant: Mountains are formed through tectonic forces..."
• 处理问候、解释和后续问题

# Clone and run
git clone https://github.com/tekaratzas/RustGPT.git
cd RustGPT
cargo run

# The model will:
# 1. Build vocabulary from training data
# 2. Pre-train on factual statements (100 epochs)
# 3. Instruction-tune on conversational data (100 epochs)
# 4. Enter interactive mode for testing

训练后，以交互方式测试模型：

Enter prompt: How do mountains form?
Model output: Mountains are formed through tectonic forces or volcanism over long geological time periods

Enter prompt: What causes rain?
Model output: Rain is caused by water vapor in clouds condensing into droplets that become too heavy to remain airborne

🧮 技术实现

• 词汇表大小：动态构建（从训练数据生成）
• 嵌入维度：128（由 src/lib.rs 中的 EMBEDDING_DIM 定义）
• 隐藏维度：256（由 src/lib.rs 中的 HIDDEN_DIM 定义）
• 最大序列长度：80 tokens（由 src/lib.rs 中的 MAX_SEQ_LEN 定义）
• 架构：3 个 Transformer 块 + 嵌入层 + 输出投影层
• 优化器：Adam，带梯度裁剪
• 预训练学习率：0.0005（100 轮）
• 指令微调学习率：0.0001（100 轮）
• 损失函数：交叉熵损失
• 梯度裁剪：L2 范数上限为 5.0
• 自定义分词，支持标点符号处理
• 贪心解码用于文本生成
• 梯度裁剪提升训练稳定性
• 模块化层系统，接口清晰
• 全面的测试覆盖，所有组件均有测试

# Run all tests
cargo test

# Test specific components
cargo test --test llm_test
cargo test --test transformer_test
cargo test --test self_attention_test

# Build optimized version
cargo build --release

# Run with verbose output
cargo test -- --nocapture

🧠 学习资源

本实现演示了关键的 ML 概念：

• Transformer 架构（注意力、前馈、层归一化）
• 神经网络反向传播
• 语言模型训练（预训练 + 微调）
• 分词和词汇管理
• 基于梯度的 Adam 优化
• 完美理解现代 LLM 内部运作的方式！

依赖项：

• ndarray - 矩阵运算用的 N 维数组库
• rand + rand_distr - 初始化的随机数生成
• 没有 PyTorch、TensorFlow 或 Candle - 只有纯 Rust 和线性代数！

欢迎贡献！本项目非常适合学习和实验。

高优先级功能需求

• 🏪 模型持久化 - 将训练后的参数保存/加载到磁盘（当前全在内存中）
• ⚡ 性能优化 - SIMD、并行训练、内存效率
• 🎯 更好的采样 - 波束搜索、top-k/top-p、温度缩放
• 📊 评估指标 - 困惑度、基准测试、训练可视化

改进方向

• 高级架构（多头注意力、位置编码、RoPE）
• 训练改进（不同优化器、学习率调度、正则化）
• 数据处理（更大数据集、分词器改进、流式处理）
• 模型分析（注意力可视化、梯度分析、可解释性）

贡献指南

• 创建功能分支：git checkout -b feature/model-persistence
• 进行更改并添加测试
• 运行测试套件：cargo test
• 提交包含清晰说明的 pull request
• 遵循标准 Rust 约定（cargo fmt）
• 为新功能添加全面测试
• 根据需要更新文档和 README
• 保持"从零开始"的哲学 - 避免重型 ML 依赖

贡献创意

• 🚀 初级：模型保存/加载、更多训练数据、配置文件
• 🔥 中级：波束搜索、位置编码、训练检查点
• ⚡ 高级：多头注意力、层并行化、自定义优化

有问题？开启一个 issue 或开始讨论！

没有 PyTorch、TensorFlow 或 Candle - 只有纯 Rust 和线性代数！