apple-aiml-research/ml-ane-transformers

Apple Neural Engine (ANE) Transformers

这是什么 – ane_transformers 是一个小型开源库，提供针对 Apple 的 Neural Engine (ANE) 优化的 Transformer 模型的参考 PyTorch 实现。它包含两个包：

• ane_transformers.reference – 一个独立、易于阅读的实现，展示了在 ANE 上高效运行所需的底层更改。
• ane_transformers.huggingface – 针对特定 Hugging Face 模型类（如 DistilBert）的即插即用替代品，保留相同 Python API 的同时集成这些优化。

为何重要 – Apple 的 A14、A15、M1、M2 等芯片均包含专用的神经处理单元。通过重构计算图（融合操作、将嵌入查找移至 CPU、针对 ANE 的指令集定制模型），该库声称在通过 Core ML 在设备上运行时，与原生 PyTorch/torch-script 模型相比，可实现 最高 10 倍的延迟降低 和 最高 14 倍的峰值内存降低。

如何使用 – README 提供了对 Hugging Face distilbert-base-uncased-finetuned-sst-2-english 模型的完整教程：

1. 使用 transformers.AutoModelForSequenceClassification 加载标准模型。
2. 替换为 ane_transformers.huggingface.distilbert.DistilBertForSequenceClassification，并复制原始 state dict。
3. 对输入进行分词，使用 torch.jit.trace 进行模型追踪，并使用 coremltools.convert 将追踪后的模型转换为 Core ML mlprogram 包。
4. 将生成的 .mlpackage 拖入 Xcode 项目，使用 Xcode 的性能标签页在 iPhone 或 Mac 上验证加速效果。

安装 – 该包已发布在 PyPI 上：

pip install ane_transformers          # 预构建的 wheel（最快）
# 或用于开发：
pip install -e .

如果在 tokenizers 依赖项上构建失败，README 指向一个涉及 Rust 编译器的已知修复方法。

支持的硬件 – 优化针对包含 ANE 的 Apple Silicon（iPhone 12/A14 或更新机型、A14+ 的 iPad、M1 或更新的 Mac）进行了调优。在旧硬件上运行时，单元测试会发出警告，但生成的 Core ML 模型仍可运行（只是无法获得宣传的加速效果）。

限制 / 注意事项

• 目前仅包装了少数 Hugging Face 模型（示例中为 DistilBERT）。扩展到其他架构需要额外工作。
• 优化后的模型增加了一次性编译开销；首次加载较慢。
• 606 个操作中有 4 个仍保留在 CPU 上（嵌入查找），这是针对该模型大小的有意设计。
• 性能提升取决于序列长度和批量大小；README 指出，对于较长序列（如 512 个标记，批量 8）可实现最高 10 倍的延迟减少。

谁应该关注此项目 – 希望在本地部署 NLP 模型并需要最佳设备内延迟和内存占用的 iOS/macOS 开发者，以及对如何将 Transformer 图适配到 Apple 的 ANE 感兴趣的研究人员。