Merlin：为Apple Silicon从零构建的高效小型语言模型

Merlin是一个专为Apple Silicon设备（MacBook Pro和iPhone）从零开始构建的高效小型语言模型项目，采用PyTorch训练、MLX推理和自定义Metal内核，在int4量化+KV缓存配置下实现625 TPS的推理速度，峰值内存仅188MB，完全适配iPhone的4GB内存预算。

随着大型语言模型在云端蓬勃发展，端侧AI推理的需求日益增长。然而，将LLM部署到消费级设备面临着严峻的挑战：内存限制、算力瓶颈、以及功耗约束。特别是对于iPhone这类移动设备，如何在有限的4GB内存预算内实现流畅的AI推理，是一个极具技术难度的工程问题。

Merlin项目正是为攻克这一难题而生。这是一个从零开始构建的小型语言模型，专为Apple Silicon生态（MacBook Pro和iPhone）优化，采用纯本地推理架构，既具有教育意义，又追求实际可用性，完全开源。

Merlin项目的核心目标清晰而聚焦：

1. 最大化Apple Silicon上的推理TPS（每秒Token数）：通过深度优化，充分释放M系列芯片和A系列芯片的神经网络引擎潜力。
2. 最小化内存占用：通过量化技术和高效的内存管理，确保模型能在资源受限的设备上流畅运行。
3. 从头训练真实数据：不依赖预训练权重，完全自主训练，确保模型的透明性和可控性。
4. 自定义Metal内核：不依赖框架默认实现，手写高性能Metal内核，榨干硬件最后一滴性能。
5. iPhone目标：通过int4量化，使模型能在约4GB内存预算内运行。

这种设计哲学体现了对端侧AI的深刻理解：与其追求参数量的大而全，不如在有限的资源约束下追求极致的效率和实用性。

Merlin在基准测试中展现了令人印象深刻的性能表现。以基础模型（约1.17亿参数）在M4 MacBook Pro上的测试为例：

配置	TPS（每秒Token数）	峰值内存
fp32，无KV缓存	38.7	1536 MB
fp32 + KV缓存	242.6	802 MB
int4 + KV缓存	625.3	188 MB

最引人注目的数据是int4量化配合KV缓存的配置：在仅188MB内存占用的情况下，实现了625.3 TPS的推理速度。这意味着模型完全处于iPhone的4GB内存预算之内，同时提供了足够快的响应速度，完全可以支撑实时交互式应用。

这一成绩的背后，是项目团队在量化算法、内存管理和内核优化上的深度工程投入。

Merlin采用GPT风格的仅解码器Transformer架构，提供了四种配置以适应不同的使用场景：

配置	参数量	嵌入维度	注意力头数	层数	上下文长度
sanity	约160万	32	2	2	64
experiment	约2100万	256	8	8	512
iphone	约31.7亿	3072	24	20	4096
macbook	约71.9亿	4096	32	26	4096

从160万参数的微型配置到71.9亿参数的桌面级配置，Merlin覆盖了从快速原型验证到生产部署的全场景需求。

Merlin在架构细节上做出了一系列精心权衡：

1. RMSNorm替代LayerNorm：去除均值减法操作，计算更快，更适合硬件加速。
2. SwiGLU MLP：相比GELU激活函数，在相同参数量下获得更好的损失表现。
3. 权重共享：token嵌入层和输出头共享权重，在基础配置上节省约3900万参数。
4. 线性层无偏置：减少参数量和计算量，现代Transformer的常见做法。
5. Pre-norm结构：将归一化层置于残差连接之前，训练更加稳定。

这些设计选择共同构成了一个既高效又稳定的训练-推理一体化架构。

Merlin采用了分层的技术栈设计，充分利用各框架的优势：

角色	工具
训练	PyTorch + CUDA + Triton（NVIDIA）
推理（Mac）	MLX + 自定义Metal内核
推理（iPhone）	CoreML（规划中）
数据	TinyStories via tiktoken GPT-2
可观测性	Weights & Biases

PyTorch作为训练的单一事实来源，确保训练过程的稳定性和可复现性。MLX专门用于推理，配合手写Metal内核实现极致性能。权重转换是显式的，并通过严格的数值一致性测试验证。

这种分工明确的架构设计，使得Merlin既能享受PyTorch生态的丰富工具链，又能充分利用Apple Silicon的专用推理加速。

Merlin的代码库结构清晰，模块化程度高：

• model.py：PyTorch Transformer实现（训练用）
• infer.py：MLX推理实现，包含KV缓存、int4量化和自定义内核
• train.py：训练循环（AdamW、梯度裁剪、W&B日志、HF Hub检查点）
• data.py：TinyStories数据集的token化和内存映射处理
• convert.py：PyTorch到MLX的权重转换
• bench.py：TPS和内存基准测试
• test_e2e.py：PyTorch/MLX数值一致性测试（atol=1e-6，贪婪Token匹配）
• docs/：架构、技术栈、训练、推理的详细文档

每个模块职责单一，接口清晰，便于理解和扩展。