llama2.🔥

你是否曾想过用纯Mojo来推理一个小型Llama 2模型？如果没有？那么现在你可以了！

支持版本: Mojo 0.26.2.0

随着Mojo的发布，我受到启发，将我的Python版Llama2.py移植到Mojo中。结果如何呢？这个版本利用了Mojo的SIMD和向量化原语，使性能相比Python版本提升了近250倍。更令人印象深刻的是，在进行了一些原生优化后，Mojo版本在多线程推理方面的表现比原始的llama2.c高出30%，并且在CPU上对小规模Llama模型的推理速度也比llama.cpp快了20%。这充分展示了通过Mojo的高级特性实现硬件级优化的巨大潜力。

支持的模型

目前，以下模型已成功通过llama2.mojo运行：

模型
stories 260K, 15M, 42M, 110M
Tinyllama-1.1B-Chat-v0.2

在Apple M1 Max上的全面基准测试

mojo与6种编程语言的对比

基准测试（更新）

Mac M1 Max（6线程）

模型	llama2.c（OMP/并行化）	llama2.mojo（并行化）	llama.cpp（CPU，6线程）	llama2.py
stories15M.bin	730 tok/s	1025 tok/s	890 tok/s	38 tok/s（pypi）
stories42M.bin	270 tok/s	490 tok/s	420 tok/s	-
stories110M.bin	102 tok/s	195 tok/s	187 tok/s	-
TinyLlama-1.1B	-	23 tok/s	-	-

Ubuntu 20.04, Intel(R) Core(TM) i7-8700 CPU @ 3.20GHz, 6核，12线程

模型	llama2.c（OMP/并行化）	llama2.mojo（并行化）	llama2.mojo（朴素矩阵乘法）	llama2.py
stories15M.bin	435 tok/s	440 tok/s	67.26 tok/s	1.3 tok/s
stories110M.bin	64 tok/s	63 tok/s	9.20 tok/s	-
TinyLlama-1.1B	7.25 tok/s	7.25 tok/s	-	-

先决条件

请确保你已经安装并配置好Mojo环境。

感受🔥的魔力

首先，导航到你存放项目的文件夹，并将此仓库克隆到该文件夹中：

git clone https://github.com/tairov/llama2.mojo.git

然后，进入仓库目录：

cd llama2.mojo

接下来，我们下载模型：

wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin

最后，直接运行Mojo代码：

mojo llama2.mojo stories15M.bin -s 100 -n 256 -t 0.5 -i "从前有个..."

命令行选项：

• -s <int> - 随机种子（默认为当前时间的毫秒数）
• -n <int> - 运行步数（默认为256，0表示最大序列长度）
• -t <float> - 温度值，范围为[0,1.0]（默认为0.9）
• -i <string> - 输入提示
• -z <string> - 分词器路径（默认为tokenizer.bin）
• -j <int> - 并行工作线程数（默认为可用的物理核心数）
• -pc <int> - 是否打印配置信息（0或1）

示例输出

并行工作线程数：4 SIMD宽度：16
总读取字节数：60816028 估计检查点大小：57 MB
层数：6 | 词汇表大小：32000
从前有个小女孩叫莉莉。她喜欢和她的玩具以及朋友们一起玩。有一天，莉莉的妈妈让她帮忙做晚饭。莉莉很高兴地照做了，拿出了盘子和叉子。
在做饭的过程中，莉莉的妈妈又让她帮忙收拾厨房。莉莉同样很开心，开始收拾盘子。突然，她看到柜台上有一把闪亮的刀。她拿起刀给妈妈看。
“莉莉，小心那把刀哦，它很锋利，可能会伤到你。”妈妈提醒道。
莉莉小心翼翼地把刀放回原处。她完成了做饭的任务，心里感到非常自豪。妈妈拥抱了她，说：“莉莉，你真是个好帮手啊！谢谢你帮我！”
每秒生成token数：148.95267649340573

引用 llama2.🔥

如果你在学术研究中使用或讨论了llama2.mojo，请引用该项目以帮助提高其知名度：

@misc{llama2.mojo,
  author = {Aydyn Tairov}, 
  title = {用纯Mojo单文件实现Llama2推理},
  year = {2023},
  month = {09},
  howpublished = {\url{https://github.com/tairov/llama2.mojo}},
  note = {Llama2 Mojo，MIT许可证}
}

我们诚挚地请求你在发表的论文中包含指向GitHub仓库的链接。这样可以让感兴趣的读者轻松找到项目的最新更新和扩展内容。

llama2.mojo旨在鼓励关于Transformer架构高效实现、Llama模型以及Mojo编程语言应用的学术研究。引用该项目有助于推动围绕这些主题的知识社区发展。感谢您对llama2.mojo的支持！

在科研与工业界

我们的项目已被学术界使用。如果您发表了使用本项目的论文或文章，请提交一个 Pull Request，我们将在此处展示。

• https://arxiv.org/abs/2410.17736 - MojoBench：针对 Mojo 的语言建模与基准测试
• https://arxiv.org/abs/2505.04080 - MojoFrame：Mojo 语言中的 DataFrame 库
• https://ieeexplore.ieee.org/abstract/document/10883176/metrics - Mojo：一种高性能编程语言的全面综述
• https://arxiv.org/abs/2502.01651 - LLaMA 2 推理的微调：面向最佳效率的语言实现比较研究

使用 Tinyllama-1.1B-Chat-v0.2

TinyLlama 是一个在 3 万亿 tokens 上训练的 11 亿参数的 Llama 模型。其紧凑的规模使其能够满足对计算和内存占用要求较低的多种应用场景。这也是我们选择它作为首个支持模型的原因。

首先，导航到您存放项目的文件夹，并将此仓库克隆到该文件夹中：

git clone https://github.com/tairov/llama2.mojo.git

然后，进入仓库目录：

cd llama2.mojo

接下来，下载模型和分词器：

wget https://huggingface.co/kirp/TinyLlama-1.1B-Chat-v0.2-bin/resolve/main/tok_tl-chat.bin
wget https://huggingface.co/kirp/TinyLlama-1.1B-Chat-v0.2-bin/resolve/main/tl-chat.bin

最后，直接运行 Mojo 程序：

mojo llama2.mojo tl-chat.bin \
    -z tok_tl-chat.bin \
    -n 256 -t 0 -s 100 -i "«|im_start|>user\n给我一个生成斐波那契数列的 Python 函数<|im_end|>\n«|im_start|>assistant\n"

示例输出

硬件线程数：6
SIMD 向量宽度：16
检查点大小：4,400,767,004 [4196 MB]
层数：22
词汇表大小：32003
«|im_start|>user
给我一个生成斐波那契数列的 Python 函数<|im_end|>
«|im_start|>assistant
当然，这里有一个生成斐波那契数列的 Python 函数：

def fibonacci(n):
    if n <= 0:
        return 0
    elif n == 1:
        return 1
    else:
        return fibonacci(n-1) + fibonacci(n-2)

这个函数接受一个整数 n 作为参数，返回下一个斐波那契数。它使用递归方法来计算斐波那契数，从 0 开始逐步累加。函数会返回当前递归层级的结果，可能是 0 或者某个斐波那契数。

许可证

MIT

llama2.mojo 快速上手指南

llama2.mojo 是一个使用纯 Mojo 语言实现的高效 Llama 2 模型推理工具。它利用 Mojo 的 SIMD 和向量化原语，在单文件内实现了比 Python 版本快近 250 倍的性能，并在多线程推理中超越了原始的 llama2.c 和 llama.cpp。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：支持 macOS (如 Apple Silicon M1/M2) 或 Linux (如 Ubuntu)。
• Mojo 环境：必须安装并配置好 Mojo 编程语言环境。
  • 推荐版本：Mojo 0.26.2.0 或更高版本。
  • 安装参考：请访问 Modular 官方文档 进行安装和配置。
• 网络环境：需要能够访问 Hugging Face 以下载模型权重文件（国内用户若访问受限，建议配置代理或使用镜像源）。

安装步骤

本项目无需复杂的编译安装过程，只需克隆仓库即可使用。

1. 打开终端，进入您存放项目的目录。
2. 克隆仓库：

   git clone https://github.com/tairov/llama2.mojo.git
   

3. 进入项目目录：

   cd llama2.mojo
   

基本使用

以下以运行轻量级的 stories15M 模型为例，展示最基础的文本生成流程。

1. 下载模型

从 Hugging Face 下载预训练的二进制模型文件：

wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin

(注：如果下载速度慢，请尝试使用国内镜像站或代理加速)

2. 运行推理

使用 mojo 命令运行脚本。以下示例将生成一段以 "Once upon a time" 开头的故事：

mojo llama2.mojo stories15M.bin -s 100 -n 256 -t 0.5 -i "Once upon a time"

参数说明

• -s <int>: 随机种子（默认：当前时间毫秒数）。
• -n <int>: 生成的步数/令牌数（默认：256，设为 0 则使用最大序列长度）。
• -t <float>: 温度值，范围 [0, 1.0]，控制生成的随机性（默认：0.9）。
• -i <string>: 输入提示词（Prompt）。
• -z <string>: 分词器路径（默认：tokenizer.bin，通常与模型同目录或自动处理）。
• -j <int>: 并行工作线程数（默认：性能核心数量）。

预期输出

运行成功后，终端将显示模型配置信息及生成的文本，例如：

num parallel workers: 4  SIMD width: 16
Total bytes read: 60816028 Estimated checkpoint size:  57 MB
n layers: 6 | vocab size: 32000
Once upon a time, there was a little girl named Lily. She loved to play with her toys and her friends...
achieved tok/s:  148.95

---

进阶提示：该项目也支持更大的模型（如 TinyLlama-1.1B），使用时需同时下载对应的模型文件和分词器文件，并通过 -z 参数指定分词器路径。