AspadaX/tiny-llm

GitHub: AspadaX/tiny-llm

一个用纯 Rust 编写的极简 MiniCPM5-1B 推理引擎,通过单文件代码和 TUI 可视化帮助开发者从底层理解大语言模型的推理过程。

Stars: 10 | Forks: 0

# tiny-llm 一个极简的 LLM 推理引擎,专为 [MiniCPM5-1B](https://huggingface.co/openbmb/MiniCPM5-1B) 打造,完全使用 Rust 编写。

我构建这个项目是为了学习大语言模型的基础知识。所有的理解都被直接转化为代码——内容密集,但井然有序。希望它也能帮助你学习。 灵感来源于 [tiny-vllm](https://github.com/kuawo/tiny-llm)。 ## 功能 - **纯 CPU 推理** — 无需 GPU - **无 KV-cache** — 专注于核心模型机制 - **单文件实现** — `src/main.rs`(约 1800 行) - **TUI 可视化** — 实时查看模型是如何“思考”的 ### 已实现内容 | 组件 | 细节 | | -------------------- | ------------------------------------------------------- | | Embedding | 通过 embedding 表进行 Token ID 查找 | | RoPE | 旋转位置嵌入 | | Attention | 带有 GQA (Grouped Query Attention) 的 Multi-head attention | | MLP | SwiGLU 激活函数 | | Normalization | RMSNorm | | Residual connections | Attention 和 MLP 之后的标准跳跃连接 | ## 快速开始 ### 1. 下载模型 ``` # 从 HuggingFace 下载 MiniCPM5-1B # 将模型文件放在目录中,例如 ./models/minicpm5-1b/ # 该目录应包含: # - config.json # - model-00000-of-00001.safetensors # - tokenizer.json ``` ### 2. 构建并运行 ``` cargo run --release -- "" "" ``` 示例: ``` cargo run --release -- ./models/minicpm5-1b "What is artificial intelligence?" ``` 按 `q` 键可提前退出。 ## 代码解析 阅读此代码的最佳方式是从上到下阅读 `src/main.rs`: 1. **Tensor 运算** — 一个基于 candle tensors 的极简 `TinyTensor` 封装,包含矩阵乘法、reshape、transpose、softmax 等操作。我会尝试手动编写这些数学运算。 2. **模型加载** — `ModelConfigurations`、`LlamaModel` 和 `TransformerBlock` 结构体,它们直接映射到 safetensors 的键。 3. **`predict_next_token`** — 遍历每个 transformer 层的推理循环。每个操作都附带了 shape 注释和耗时记录。 4. **TUI** — Ratatui 组件实时渲染 attention 图、tensor shape 和 logits。 ## 许可证 MIT — 请参阅 [LICENSE](LICENSE)
标签:AI, DLL 劫持, Rust, 可视化界面, 大语言模型, 推理引擎, 教学项目, 网络流量审计, 自动化代码审查, 通知系统