将模型导出到 Inferentia

摘要

将 PyTorch 模型导出到 Neuron 模型就像

optimum-cli export neuron \
  --model bert-base-uncased \
  --sequence_length 128 \
  --batch_size 1 \
  bert_neuron/

查看帮助以获取更多选项

optimum-cli export neuron --help

为什么要编译成 Neuron 模型？

AWS 提供了两代为机器学习推理而构建的 Inferentia 加速器，它们具有更高的吞吐量、更低的延迟，但成本更低：inf2（NeuronCore-v2）和 inf1（NeuronCore-v1）。

在生产环境中，要在 Neuron 设备上部署 🤗 Transformers 模型，您需要在推理之前编译您的模型并将其导出为序列化格式。通过使用 Neuron 编译器（neuronx-cc 或 neuron-cc）进行提前（AOT）编译，您的模型将被转换为序列化和优化的 TorchScript 模块。

虽然预编译避免了推理期间的开销，但编译后的 Neuron 模型有一些限制

编译期间使用的输入形状和数据类型无法更改。
Neuron 模型针对每个硬件和 SDK 版本进行了专门设计，这意味着
- 使用 Neuron 编译的模型无法在非 Neuron 环境中执行。
- 为 inf1（NeuronCore-v1）编译的模型与 inf2（NeuronCore-v2）不兼容，反之亦然。
- 为特定 SDK 版本编译的模型（通常）与其他 SDK 版本不兼容。

在本指南中，我们将向您展示如何将模型导出为针对 Neuron 设备优化的序列化模型。

🤗 Optimum 通过利用配置对象来支持 Neuron 导出。这些配置对象为许多模型架构预先准备好了，并且设计为易于扩展到其他架构。

要检查支持的架构，请转到配置参考页面。

使用 CLI 将模型导出到 Neuron

要将 🤗 Transformers 模型导出到 Neuron，您首先需要安装一些额外的依赖项

适用于 Inf2

pip install optimum[neuronx]

适用于 Inf1

pip install optimum[neuron]

Optimum Neuron 导出可以通过 Optimum 命令行使用

optimum-cli export neuron --help

usage: optimum-cli export neuron [-h] -m MODEL [--task TASK] [--atol ATOL] [--cache_dir CACHE_DIR] [--trust-remote-code]
                                 [--compiler_workdir COMPILER_WORKDIR] [--disable-validation] [--auto_cast {none,matmul,all}]
                                 [--auto_cast_type {bf16,fp16,tf32}] [--dynamic-batch-size] [--num_cores NUM_CORES] [--unet UNET]
                                 [--output_hidden_states] [--output_attentions] [--batch_size BATCH_SIZE]
                                 [--sequence_length SEQUENCE_LENGTH] [--num_beams NUM_BEAMS] [--num_choices NUM_CHOICES]
                                 [--num_channels NUM_CHANNELS] [--width WIDTH] [--height HEIGHT]
                                 [--num_images_per_prompt NUM_IMAGES_PER_PROMPT] [-O1 | -O2 | -O3]
                                 output

optional arguments:
  -h, --help            show this help message and exit
  -O1                   Enables the core performance optimizations in the compiler, while also minimizing compile time.
  -O2                   [Default] Provides the best balance between model performance and compile time.
  -O3                   May provide additional model execution performance but may incur longer compile times and higher host
                        memory usage during model compilation.

Required arguments:
  -m MODEL, --model MODEL
                        Model ID on huggingface.co or path on disk to load model from.
  output                Path indicating the directory where to store generated Neuronx compiled TorchScript model.

Optional arguments:
  --task TASK           The task to export the model for. If not specified, the task will be auto-inferred based on the model.
                        Available tasks depend on the model, but are among: ['audio-classification', 'audio-frame-
                        classification', 'audio-xvector', 'automatic-speech-recognition', 'conversational', 'depth-estimation',
                        'feature-extraction', 'fill-mask', 'image-classification', 'image-segmentation', 'image-to-image',
                        'image-to-text', 'mask-generation', 'masked-im', 'multiple-choice', 'object-detection', 'question-
                        answering', 'semantic-segmentation', 'text-to-audio', 'text-generation', 'text2text-generation', 'text-
                        classification', 'token-classification', 'zero-shot-image-classification', 'zero-shot-object-detection',
                        'stable-diffusion', 'stable-diffusion-xl'].
  --atol ATOL           If specified, the absolute difference tolerance when validating the model. Otherwise, the default atol
                        for the model will be used.
  --cache_dir CACHE_DIR
                        Path indicating where to store cache.
  --trust-remote-code   Allow to use custom code for the modeling hosted in the model repository. This option should only be set
                        for repositories you trust and in which you have read the code, as it will execute on your local machine
                        arbitrary code present in the model repository.
  --compiler_workdir COMPILER_WORKDIR
                        Path indicating the directory where to store intermediary files generated by Neuronx compiler.
  --disable-validation  Whether to disable the validation of inference on neuron device compared to the outputs of original
                        PyTorch model on CPU.
  --auto_cast {none,matmul,all}
                        Whether to cast operations from FP32 to lower precision to speed up the inference. Can be `"none"`,
                        `"matmul"` or `"all"`.
  --auto_cast_type {bf16,fp16,tf32}
                        The data type to cast FP32 operations to when auto-cast mode is enabled. Can be `"bf16"`, `"fp16"` or
                        `"tf32"`.
  --dynamic-batch-size  Enable dynamic batch size for neuron compiled model. If this option is enabled, the input batch size can
                        be a multiple of the batch size during the compilation, but it comes with a potential tradeoff in terms
                        of latency.
  --num_cores NUM_CORES
                        The number of cores on which the model should be deployed (text-generation only).
  --unet UNET           UNet model ID on huggingface.co or path on disk to load model from. This will replace the unet in the
                        original Stable Diffusion pipeline.
  --output_hidden_states
                        Whether or not for the traced model to return the hidden states of all layers.
  --output_attentions   Whether or not for the traced model to return the attentions tensors of all attention layers.

Input shapes:
  --batch_size BATCH_SIZE
                        Batch size that the Neuronx-cc compiler exported model will be able to take as input.
  --sequence_length SEQUENCE_LENGTH
                        Sequence length that the Neuronx-cc compiler exported model will be able to take as input.
  --num_beams NUM_BEAMS
                        Number of beams for beam search that the Neuronx-cc compiler exported model will be able to take as
                        input.
  --num_choices NUM_CHOICES
                        Only for the multiple-choice task. Num choices that the Neuronx-cc compiler exported model will be able
                        to take as input.
  --num_channels NUM_CHANNELS
                        Image tasks only. Number of channels that the Neuronx-cc compiler exported model will be able to take as
                        input.
  --width WIDTH         Image tasks only. Width that the Neuronx-cc compiler exported model will be able to take as input.
  --height HEIGHT       Image tasks only. Height that the Neuronx-cc compiler exported model will be able to take as input.
  --num_images_per_prompt NUM_IMAGES_PER_PROMPT
                        Stable diffusion only. Number of images per prompt that the Neuronx-cc compiler exported model will be
                        able to take as input.

导出标准（非 LLM）模型

Hugging Face 中心上的大多数模型都可以使用 torch trace 直接导出，然后转换为序列化和优化的 TorchScript 模块。

NEFF：Neuron 可执行文件格式，它是 Neuron 设备上的二进制可执行文件。

导出模型时，必须传递两组导出参数

compiler_args 是编译器的可选参数，这些参数通常控制编译器如何在推理性能（延迟和吞吐量）和准确性之间进行权衡，
input_shapes 是您需要发送到 Neuron 编译器的强制静态形状信息。

请键入以下命令以查看所有导出参数

optimum-cli export neuron -h

导出标准 NLP 模型可以按如下方式完成

optimum-cli export neuron --model distilbert-base-uncased-distilled-squad \
                          --batch_size 1 --sequence_length 16 \
                          --auto_cast matmul --auto_cast_type fp16 \
                          distilbert_base_uncased_squad_neuron/

这里模型使用 (1, 16) 的静态输入形状导出，并且编译器参数指定矩阵乘法运算必须以 float16 精度执行以实现更快的推理。

导出后，您应该会看到以下日志，这些日志通过将 Neuron 设备上的模型与 CPU 上的 PyTorch 模型进行比较来验证模型。

Validating Neuron model...
        -[✓] Neuron model output names match reference model (last_hidden_state)
        - Validating Neuron Model output "last_hidden_state":
                -[✓] (1, 16, 32) matches (1, 16, 32)
                -[✓] all values close (atol: 0.0001)
The Neuronx export succeeded and the exported model was saved at: distilbert_base_uncased_squad_neuron/

这会导出由 --model 参数定义的检查点的 Neuron 编译的 TorchScript 模块。

如您所见，任务已自动检测到。这是因为模型位于 Hub 上。对于本地模型，需要提供 --task 参数，否则它将默认为没有任何特定任务头的模型架构

optimum-cli export neuron --model local_path --task question-answering --batch_size 1 --sequence_length 16 --dynamic-batch-size distilbert_base_uncased_squad_neuron/

请注意，为 Hub 上的模型提供 --task 参数将禁用自动任务检测。然后，可以加载生成的 model.neuron 文件并在 Neuron 设备上运行。

对于每个模型架构，您可以通过 ~exporters.tasks.TasksManager 找到支持的任务列表。例如，对于 DistilBERT，对于 Neuron 导出，我们有

>>> from optimum.exporters.tasks import TasksManager
>>> from optimum.exporters.neuron.model_configs import *  # Register neuron specific configs to the TasksManager

>>> distilbert_tasks = list(TasksManager.get_supported_tasks_for_model_type("distilbert", "neuron").keys())
>>> print(distilbert_tasks)
['feature-extraction', 'fill-mask', 'multiple-choice', 'question-answering', 'text-classification', 'token-classification']

然后，您可以将其中一项任务传递到 optimum-cli export neuron 命令中的 --task 参数，如上所述。

导出后，可以使用 NeuronModelForXXX 类直接对 Neuron 模型进行推理

>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronModelForSequenceClassification

>>> tokenizer = AutoTokenizer.from_pretrained("./distilbert-base-uncased-finetuned-sst-2-english_neuron/")
>>> model = NeuronModelForSequenceClassification.from_pretrained("./distilbert-base-uncased-finetuned-sst-2-english_neuron/")

>>> inputs = tokenizer("Hamilton is considered to be the best musical of human history.", return_tensors="pt")
>>> logits = model(**inputs).logits
>>> print(model.config.id2label[logits.argmax().item()])
'POSITIVE'

如您所见，无需传递导出过程中使用的 Neuron 参数，因为它们保存在 config.json 文件中，并且将由 NeuronModelForXXX 类自动恢复。

请注意，输入始终填充到用于编译的形状，并且填充会带来计算开销。将静态形状调整为高于推理过程中将馈送到模型的输入的形状，但不要高出太多。

将 Stable Diffusion 导出到 Neuron

使用 Optimum CLI，您可以编译 Stable Diffusion 管道中的组件，以便在推理期间在 Neuron 设备上获得加速。

到目前为止，我们支持导出管道中的以下组件

CLIP 文本编码器
U-Net
VAE 编码器
VAE 解码器

“选择这些块是因为它们代表了管道中大部分计算，并且性能基准测试表明，在 Neuron 上运行它们会带来显著的性能提升。”

此外，请随时调整编译配置以找到用例中性能与准确性之间的最佳权衡。默认情况下，我们建议将 FP32 矩阵乘法运算转换为 BF16，这可以在适度牺牲准确性的前提下提供良好的性能。查看来自 AWS Neuron 文档的指南，以更好地了解编译选项。

可以使用 CLI 导出 Stable Diffusion 检查点

optimum-cli export neuron --model stabilityai/stable-diffusion-2-1-base \
  --task stable-diffusion \
  --batch_size 1 \
  --height 512 `# height in pixels of generated image, eg. 512, 768` \
  --width 512 `# width in pixels of generated image, eg. 512, 768` \
  --num_images_per_prompt 4 `# number of images to generate per prompt, defaults to 1` \
  --auto_cast matmul `# cast only matrix multiplication operations` \
  --auto_cast_type bf16 `# cast operations from FP32 to BF16` \
  sd_neuron/

将 Stable Diffusion XL 导出到 Neuron

与 Stable Diffusion 类似，您可以使用 Optimum CLI 编译 SDXL 管道中的组件以在 Neuron 设备上进行推理。

我们支持导出管道中的以下组件以提高速度

文本编码器
第二个文本编码器
U-Net（比 Stable Diffusion 管道中的 U-Net 大三倍）
VAE 编码器
VAE 解码器

“Stable Diffusion XL 特别适用于 768 到 1024 之间的图像。”

可以使用 CLI 导出 SDXL 检查点

optimum-cli export neuron --model stabilityai/stable-diffusion-xl-base-1.0 \
  --task stable-diffusion-xl \
  --batch_size 1 \
  --height 1024 `# height in pixels of generated image, eg. 768, 1024` \
  --width 1024 `# width in pixels of generated image, eg. 768, 1024` \
  --num_images_per_prompt 4 `# number of images to generate per prompt, defaults to 1` \
  --auto_cast matmul `# cast only matrix multiplication operations` \
  --auto_cast_type bf16 `# cast operations from FP32 to BF16` \
  sd_neuron/

将 LLM 导出到 Neuron

LLM 模型不是使用 Torch tracing 导出，而是直接转换为 Neuron 图，然后可以加载 Transformers 检查点权重。

就像标准 NLP 模型一样，您需要在导出 LLM 模型时指定静态参数

batch_size 是模型将接受的输入序列的数量。默认为 1，
sequence_length 是输入序列中令牌的最大数量。默认为 max_position_embeddings（对于旧模型，为 n_positions）。
auto_cast_type 指定编码权重的格式。它可以是 fp32（float32）、fp16（float16）或 bf16（bfloat16）之一。默认为 fp32。
num_cores 是实例化模型时使用的 Neuron 核心的数量。每个 Neuron 核心拥有 16 Gb 的内存，这意味着更大的模型需要拆分到多个核心上。默认为 1，

optimum-cli export neuron --model meta-llama/Meta-Llama-3-8B \
  --batch_size 1 \
  --sequence_length 4096 \
  --auto_cast_type fp16 `# cast operations from BF16 to FP16` \
  --num_cores 2 \
  llama3_neuron/

一个重要的限制是，LLM 模型只能在 Neuron 平台上导出，因为它们是在导出期间量身定制以适合实际设备的。

LLM 模型的导出可能比标准模型花费更长的时间（有时超过一个小时）。

如前所述，Neuron 模型参数是静态的。这意味着在推理期间

输入的 batch_size 应该小于导出期间使用的 batch_size，
输入序列的 length 应该小于导出期间使用的 sequence_length，
令牌的最大数量（输入 + 生成的）不能超过导出期间使用的 sequence_length。

导出后，可以使用 NeuronModelForCausalLM 类简单地重新加载 Neuron 模型。与原始 Transformers 模型一样，使用 generate() 而不是 forward() 生成文本序列。

from transformers import AutoTokenizer
-from transformers import AutoModelForCausalLM
+from optimum.neuron import NeuronModelForCausalLM

# Instantiate and convert to Neuron a PyTorch checkpoint
-model = AutoModelForCausalLM.from_pretrained("gpt2")
+model = NeuronModelForCausalLM.from_pretrained("./gpt2-neuron")

tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token_id = tokenizer.eos_token_id

tokens = tokenizer("I really wish ", return_tensors="pt")
with torch.inference_mode():
    sample_output = model.generate(
        **tokens,
        do_sample=True,
        min_length=128,
        max_length=256,
        temperature=0.7,
    )
    outputs = [tokenizer.decode(tok) for tok in sample_output]
    print(outputs)

生成是高度可配置的。有关详细信息，请参阅 https://huggingface.co/docs/transformers/generation_strategies。

请注意

对于每个模型架构，都会为所有参数提供默认值，但传递给 generate 方法的值将优先，
生成参数可以存储在 generation_config.json 文件中。当模型目录中存在此类文件时，它将被解析以设置默认参数（传递给 generate 方法的值仍然优先）。

通过 NeuronModel 以编程方式将模型导出到 Neuron

作为 optimim-cli 的替代方案，您也可以在自己的 Python 脚本或笔记本中使用 optimum.neuron.NeuronModelForXXX 模型类将模型导出到 Neuron。

以下是一个示例

>>> from optimum.neuron import NeuronModelForSequenceClassification

>>> input_shapes = {"batch_size": 1, "sequence_length": 64}  # mandatory shapes
>>> model = NeuronModelForSequenceClassification.from_pretrained(
...   "distilbert-base-uncased-finetuned-sst-2-english", export=True, **input_shapes
... )

# Save the model
>>> model.save_pretrained("./distilbert-base-uncased-finetuned-sst-2-english_neuron/")

# Push the neuron model to HF Hub
>>> model.push_to_hub(
...     "a_local_path_for_compiled_neuron_model", repository_id="my-neuron-repo", use_auth_token=True
... )

此示例可以使用与 optimum-cli 相同的导出参数适用于其他模型类型。

使用 NeuronX TGI 导出 Neuron 模型

NeuronX TGI 镜像不仅包含 NeuronX 运行时，还包含导出 Neuron 模型所需的所有包和工具。

使用以下命令使用 TGI 镜像将模型导出到 Neuron

docker run --entrypoint optimum-cli \
       -v $(pwd)/data:/data \
       --privileged \
       ghcr.io/huggingface/neuronx-tgi:latest \
       export neuron \
       --model <organization>/<model> \
       --batch_size 1 \
       --sequence_length 4096 \
       --auto_cast_type fp16 \
       --num_cores 2 \
       /data/<neuron_model_path>

导出的模型将保存到 ./data/<neuron_model_path> 下。

AWS Trainium & Inferentia