苹果神经引擎（ANE）Transformers

如果您正考虑在搭载 A14 或更高版本以及 M1 或更高版本芯片的苹果设备上部署您的 Transformer 模型，以实现相比基准实现最高快 10 倍且峰值内存消耗降低14 倍的效果，可以将 ane_transformers 作为参考 PyTorch 实现。

ane_transformers.reference 包含一个独立的参考实现，而 ane_transformers.huggingface 则包含了对 Hugging Face 模型类（如 distilbert）的优化版本，旨在展示我们研究论文中提出的优化原则在现有第三方实现中的应用。

请参阅我们的研究论文，其中详细解释了各项优化，并提供了交互式图表，帮助您探索来自案例研究的延迟和峰值内存消耗数据：在不同设备及操作系统版本上部署 Hugging Face distilbert 模型。以下图片为研究论文中关于安装 iOS 16.0 的 iPhone 13 的非交互式截图：

教程：优化部署 Hugging Face distilbert

本教程是针对我们研究论文中案例研究的模型部署流程的分步指南。上述图表中的 Hugging Face distilbert 性能数据正是使用相同代码生成的。

为了开始优化，我们首先初始化基准模型如下：

import transformers
model_name = "distilbert-base-uncased-finetuned-sst-2-english"
baseline_model = transformers.AutoModelForSequenceClassification.from_pretrained(
    model_name,
    return_dict=False,
    torchscript=True,
).eval()

随后，我们初始化一个数学上等价但经过优化的模型，并将其参数从基准模型中加载过来：

from ane_transformers.huggingface import distilbert as ane_distilbert
optimized_model = ane_distilbert.DistilBertForSequenceClassification(
    baseline_model.config).eval()
optimized_model.load_state_dict(baseline_model.state_dict())

接下来，我们为模型创建示例输入：

tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
tokenized = tokenizer(
    ["用于追踪模型的示例输入文本"],
    return_tensors="pt",
    max_length=128,  # token 序列长度
    padding="max_length",
)

然后，我们对优化后的模型进行追踪，以获得 coremltools 转换工具所需的预期输入格式（Torchscript）。

import torch
traced_optimized_model = torch.jit.trace(
    optimized_model,
    (tokenized["input_ids"], tokenized["attention_mask"])
)

最后，我们使用 coremltools 生成 Core ML 模型包文件并将其保存：

import coremltools as ct
import numpy as np
ane_mlpackage_obj = ct.convert(
    traced_optimized_model,
    convert_to="mlprogram",
    inputs=[
        ct.TensorType(
                f"input_{name}",
                    shape=tensor.shape,
                    dtype=np.int32,
                ) for name, tensor in tokenized.items()
            ],
            compute_units=ct.ComputeUnit.ALL,
)
out_path = "HuggingFace_ane_transformers_distilbert_seqLen128_batchSize1.mlpackage"
ane_mlpackage_obj.save(out_path)

为验证性能，开发者现在可以启动 Xcode，并简单地将该模型包文件作为资源添加到项目中。点击“Performance”选项卡后，开发者便可在本地可用设备上生成性能报告，例如在运行 Xcode 的 Mac 上，或连接到该 Mac 的其他苹果设备上。下图展示了在安装 iOS 16.0 的 iPhone 13 Pro Max 上为此模型生成的性能报告。

根据上图所示，对于本教程所选用的序列长度 128、批次大小 1 的配置，延迟改善了 2.84 倍。若采用更高的序列长度（如 512）和批次大小（如 8），则可实现延迟降低多达 10 倍、峰值内存消耗降低 14 倍的效果。有关详细且交互式的性能数据，请参阅我们研究论文中的图 2。

请注意，由于优化后模型中的操作数量增加，加载和编译时间会有所延长。不过，这些成本仅发生一次，如果以异步方式加载模型，则不会影响用户体验。

此外，优化后的模型中有 606 个操作中的 4 个是在 CPU 上执行的。这些是与嵌入查找相关的操作，对于当前的模型配置而言，在 CPU 上执行更为高效。

关于单元测试的说明

单元测试除其他内容外，还会测量 ANE 的加速倍数。鉴于此参考实现的设备规格要求为：Mac 需配备 M1 或更高版本芯片，iPhone 和 iPad 则需搭载 A14 或更高版本芯片，因此，若在不符合此规格的设备上运行速度提升相关的单元测试，系统将显示警告信息。即便模型是在不符合规格的 Mac 上生成的，它仍应在符合规格的设备上正常工作。

安装与故障排除

• 最快方式：pip install ane_transformers
• 可本地编辑：pip install -e .
• 若安装过程中出现 ERROR: Failed building wheel for tokenizers 或 error: can't find Rust compiler 错误，请按照此解决方案进行处理。

ml-ane-transformers 快速上手指南

本指南帮助中国开发者在搭载 A14/M1 及以上芯片的 Apple 设备上，利用 Apple Neural Engine (ANE) 加速部署 Hugging Face Transformer 模型（如 DistilBERT），实现最高 10 倍 的推理速度提升和 14 倍 的峰值内存降低。

环境准备

在开始之前，请确保满足以下硬件和软件要求：

• 硬件要求：
  • Mac: M1 芯片或更新版本。
  • iPhone/iPad: A14 仿生芯片或更新版本。
• 软件要求：
  • macOS 12.0+ (Monterey) 或 iOS 16.0+。
  • Python 3.8+。
  • Xcode（用于后续在真机上验证性能报告）。
• 前置依赖：
  • torch (PyTorch)
  • transformers (Hugging Face)
  • coremltools (Apple Core ML 转换工具)
  • numpy

注意：虽然模型可以在非指定设备（如 Intel Mac）上生成，但性能加速效果仅在符合上述芯片规格的设备上生效。

安装步骤

推荐使用 pip 进行安装。国内开发者若遇到网络问题，可配置清华或阿里镜像源加速。

标准安装：

pip install ane_transformers

使用国内镜像源加速安装（推荐）：

pip install ane_transformers -i https://pypi.tuna.tsinghua.edu.cn/simple

开发模式安装（如需修改源码）：

git clone https://github.com/apple/ml-ane-transformers.git
cd ml-ane-transformers
pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple

故障排除：如果安装过程中出现 ERROR: Failed building wheel for tokenizers 或 error: can't find Rust compiler，请先安装 Rust 编译器，参考解决方案：HuggingFace Tokenizers Issue #2831。

基本使用

以下示例演示如何将 Hugging Face 的 distilbert 模型转换为优化的 Core ML 格式，以便在 Apple 设备上利用 ANE 加速。

1. 初始化基线模型与优化模型

首先加载标准的 Hugging Face 模型，然后初始化学术等效但经过 ANE 优化的模型结构，并载入基线模型的参数。

import transformers
from ane_transformers.huggingface import distilbert as ane_distilbert

model_name = "distilbert-base-uncased-finetuned-sst-2-english"

# 初始化基线模型
baseline_model = transformers.AutoModelForSequenceClassification.from_pretrained(
    model_name,
    return_dict=False,
    torchscript=True,
).eval()

# 初始化优化模型并加载参数
optimized_model = ane_distilbert.DistilBertForSequenceClassification(
    baseline_model.config).eval()
optimized_model.load_state_dict(baseline_model.state_dict())

2. 准备输入数据与模型追踪 (Tracing)

创建样本输入数据，并使用 PyTorch JIT 追踪优化后的模型，以便 coremltools 能够识别输入输出格式。

import torch

# 初始化分词器并处理文本
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
tokenized = tokenizer(
    ["Sample input text to trace the model"],
    return_tensors="pt",
    max_length=128,  # token 序列长度
    padding="max_length",
)

# 追踪优化模型以获取 Torchscript 格式
traced_optimized_model = torch.jit.trace(
    optimized_model,
    (tokenized["input_ids"], tokenized["attention_mask"])
)

3. 转换为 Core ML 模型

使用 coremltools 将追踪后的模型转换为 .mlpackage 格式，指定所有计算单元（包括 ANE）参与计算。

import coremltools as ct
import numpy as np

ane_mlpackage_obj = ct.convert(
    traced_optimized_model,
    convert_to="mlprogram",
    inputs=[
        ct.TensorType(
                f"input_{name}",
                    shape=tensor.shape,
                    dtype=np.int32,
                ) for name, tensor in tokenized.items()
            ],
            compute_units=ct.ComputeUnit.ALL,
)

# 保存模型文件
out_path = "HuggingFace_ane_transformers_distilbert_seqLen128_batchSize1.mlpackage"
ane_mlpackage_obj.save(out_path)

4. 验证性能

转换完成后，将生成的 .mlpackage 文件添加到 Xcode 项目的资源文件中。在 Xcode 中打开该模型，点击 Performance 标签页，即可连接本地 Mac 或已连接的 iOS 设备生成性能报告。

• 预期效果：在序列长度 128、Batch Size 1 的配置下，延迟可降低约 2.84 倍；若增加序列长度（如 512）或 Batch Size，延迟最高可降低 10 倍，峰值内存消耗降低 14 倍。
• 注意：优化模型的操作数量略有增加，导致加载和编译时间稍长，建议异步加载模型以提升用户体验。少量嵌入查找操作仍会在 CPU 上执行以获得最佳效率。