Hugging Face's logo

AWS Trainium 和 Inferentia 文档

神经元模型缓存

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

神经元模型缓存

神经元模型缓存是 neff 格式的已编译神经元模型的远程缓存。它集成到 NeuronTrainer 和 NeuronModelForCausalLM 类中，以支持从缓存加载预训练模型，而不是在本地编译它们。

注意：它不适用于使用任何其他 NeuronModelXX 类导出的模型，这些类使用不同的导出机制。

神经元模型缓存托管在 Hugging Face Hub 上，并包含所有流行的和受支持的 optimum-neuron 预训练模型的已编译文件。

在 Neuron 平台上训练 Transformers 或 Diffusion 模型或加载 NeuronModelForCausalLM 之前，需要使用 torch-neuronx 将其导出为 neuron 格式。

导出模型时，torch-neuronx 将会

  • 将其转换为一组 XLA 子图，
  • 使用 neuronx 编译器将每个子图编译成神经元可执行文件格式 (NEFF) 二进制文件。

第一步相对较快，但编译需要花费大量时间。为了避免每次在 NeuronX 主机上加载模型时都重新编译所有 NEFF 文件，torch-neuronx 将 NEFF 文件存储在本地目录中，通常为 /var/tmp/neuron-compile-cache。

但是，此本地缓存不在平台之间共享，这意味着每次在新主机上训练或导出模型时，都需要重新编译它。

我们创建了神经元模型缓存来解决此限制，方法是提供预编译模型图的公共存储库。

注意：我们还支持创建私有的、安全的远程模型缓存。

如何使用神经元模型缓存

当您使用 NeuronTrainer 或 NeuronModelForCausalLM 类时，将使用公共模型缓存。无需进行其他更改。

当将模型导出为 neuron 格式时，optimum-neuron 将在编译模型子图期间简单地在 hub 存储库中查找缓存的 NEFF 文件。

如果 NEFF 文件被缓存，它们将从 hub 中获取并直接加载，而不是重新编译。

缓存如何工作

Optimum Neuron 缓存构建在 NeuronX 编译器缓存之上。

重要的是要理解缓存操作的是 NEFF 二进制文件，而不是模型本身。

如前所述，使用 NeuronTrainer 或 NeuronModelForCausalLM 导出到 Neuron 的每个模型都由 XLA 子图组成。

每个子图都是唯一的，并且由以下各项的组合产生

  • transformers 或 transformers_neuronx python 建模代码，
  • transformers 模型配置，
  • 导出期间选择的 input_shapes，
  • 模型的精度，全精度、fp16 或 bf16。

当将子图编译为 NEFF 文件时，其他参数会影响结果

  • Neuron X 编译器的版本，
  • 使用的神经元核心数，
  • 编译参数（例如优化级别）。

所有这些参数组合在一起以创建唯一哈希，该哈希标识 NEFF 文件。

这有两个非常重要的结果

  • 只有在实际导出模型时，才能识别关联的 NEFF 文件，
  • 即使模型配置中的微小更改也会导致一组不同的 NEFF 文件。

因此，很难提前知道与特定模型配置关联的 NEFF 是否被缓存。

神经元模型缓存查找（仅限 Inferentia）

神经元缓存查找是一项功能，允许用户在导出模型以进行推理之前查找兼容的缓存模型配置。

它基于由存储的缓存配置组成的专用注册表。

缓存的模型配置作为条目存储在神经元模型缓存中的特定子文件夹下

neuronxcc-2.12.54.0+f631c2365
├── 0_REGISTRY
│   └── 0.0.18
│       └── llama
│           └── meta-llama
│               └── Llama-2-7b-chat-hf
│                   └── 54c1f6689cd88f246fce.json

每个条目对应于模型配置及其导出参数的组合：这尽可能接近唯一标识导出的模型。

您可以使用 optimum-cli 通过传递 hub model_id 或包含模型 config.json 的文件的路径来查找兼容的缓存条目。

$ optimum-cli neuron cache lookup meta-llama/Llama-2-7b-chat-hf

*** 1 entrie(s) found in cache for meta-llama/Llama-2-7b-chat-hf ***

task: text-generation
batch_size: 1
num_cores: 24
auto_cast_type: fp16
sequence_length: 2048
compiler_type: neuronx-cc
compiler_version: 2.12.54.0+f631c2365
checkpoint_id: meta-llama/Llama-2-7b-chat-hf
checkpoint_revision: c1b0db933684edbfe29a06fa47eb19cc48025e93

请注意，即使存在兼容的缓存条目，如果您修改了编译参数或更新了 neuronx 包，这也不能始终保证模型在导出期间不会被重新编译。

高级用法（仅限 Trainium）

如何使用私有神经元模型缓存（仅限 Trainium）

公共缓存的存储库是 aws-neuron/optimum-neuron-cache。此存储库包含常用模型的所有预编译文件，因此它是公开可用的，并且每个人都可以免费使用。但是有两个限制

  1. 您将无法在此存储库上推送您自己的已编译文件
  2. 它是公开的，您可能希望为私有模型使用私有存储库

为了缓解这种情况，您可以使用 optimum-cli 创建您自己的私有缓存存储库，或设置环境变量 CUSTOM_CACHE_REPO。

使用 Optimum CLI

Optimum CLI 为缓存创建和设置提供 2 个子命令

  • create：创建新的缓存存储库，您可以将其用作私有神经元模型缓存。
  • set：在本地设置神经元缓存存储库的名称，存储库需要存在，并且默认情况下将由 optimum-neuron 使用。

创建一个新的神经元缓存存储库

optimum-cli neuron cache create --help

usage: optimum-cli neuron cache create [-h] [-n NAME] [--public]

optional arguments:
  -h, --help            show this help message and exit
  -n NAME, --name NAME  The name of the repo that will be used as a remote cache for the compilation files.
  --public              If set, the created repo will be public. By default the cache repo is private.

-n / --name 选项允许您为神经元缓存存储库指定名称，如果未设置，将使用默认名称。--public 标志允许您将神经元缓存公开，因为它默认情况下将作为私有存储库创建。

示例

optimum-cli neuron cache create

Neuron cache created on the Hugging Face Hub: michaelbenayoun/optimum-neuron-cache [private].
Neuron cache name set locally to michaelbenayoun/optimum-neuron-cache in /home/michael/.cache/huggingface/optimum_neuron_custom_cache.

设置不同的 Trainium 缓存存储库

usage: optimum-cli neuron cache set [-h] name

positional arguments:
  name        The name of the repo to use as remote cache.

optional arguments:
  -h, --help  show this help message and exit

示例

optimum-cli neuron cache set michaelbenayoun/optimum-neuron-cache

Neuron cache name set locally to michaelbenayoun/optimum-neuron-cache in /home/michael/.cache/huggingface/optimum_neuron_custom_cache

当在新实例上工作以使用您自己的缓存时，optimum-cli neuron cache set 命令很有用。

使用环境变量 CUSTOM_CACHE_REPO

使用 CLI 并不总是可行的，对于小型测试也不是很实用。在这种情况下，您可以简单地设置环境变量 CUSTOM_CACHE_REPO。

例如，如果您的缓存存储库名为 michaelbenayoun/my_custom_cache_repo，您只需要执行

CUSTOM_CACHE_REPO="michaelbenayoun/my_custom_cache_repo" torchrun ...

或

export CUSTOM_CACHE_REPO="michaelbenayoun/my_custom_cache_repo"
torchrun ...

您必须登录到 Hugging Face Hub 才能从您的私有缓存存储库推送和拉取文件。

缓存系统流程

Cache system flow
缓存系统流程

在每个训练步骤的开始，NeuronTrainer 计算 NeuronHash 并检查 Hugging Face Hub 上的缓存存储库（官方和自定义），以查看是否存在与此哈希关联的已编译文件。如果是这种情况，文件将直接下载到本地缓存目录，并且不需要编译。否则，将执行编译。

就像下载已编译文件一样，NeuronTrainer 将跟踪每个训练步骤中新创建的编译文件，并在保存时或训练结束时将其上传到 Hugging Face Hub。这假设您具有对缓存存储库的写入权限，否则将不会推送任何内容。

Optimum CLI

Optimum CLI 可用于执行各种与缓存相关的任务，如 optimum-cli neuron cache 命令用法消息所述

usage: optimum-cli neuron cache [-h] {create,set,add,list} ...

positional arguments:
  {create,set,add,list,synchronize,lookup}
    create              Create a model repo on the Hugging Face Hub to store Neuron X compilation files.
    set                 Set the name of the Neuron cache repo to use locally (trainium only).
    add                 Add a model to the cache of your choice (trainium only).
    list                List models in a cache repo (trainium only).
    synchronize         Synchronize local compiler cache with the hub cache (inferentia only).
    lookup              Lookup the neuronx compiler hub cache for the specified model id (inferentia only).

optional arguments:
  -h, --help            show this help message and exit

向缓存添加模型（仅限 Trainium）

可以通过 optimum-cli neuron cache add 命令将模型编译文件添加到缓存存储库

usage: optimum-cli neuron cache add [-h] -m MODEL --task TASK --train_batch_size TRAIN_BATCH_SIZE [--eval_batch_size EVAL_BATCH_SIZE] [--sequence_length SEQUENCE_LENGTH]
                                    [--encoder_sequence_length ENCODER_SEQUENCE_LENGTH] [--decoder_sequence_length DECODER_SEQUENCE_LENGTH]
                                    [--gradient_accumulation_steps GRADIENT_ACCUMULATION_STEPS] --precision {fp,bf16} --num_cores
                                    {1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32} [--max_steps MAX_STEPS]

运行此命令时，将运行一个小型训练会话，并将推送生成的编译文件。

确保本地设置了要使用的神经元缓存存储库，这可以通过运行 optimum-cli neuron cache set 命令来完成。您还需要确保您已登录到 Hugging Face Hub，并且您具有指定缓存存储库的写入权限，这可以通过 huggingface-cli login 命令来完成。

如果至少其中一项要求未满足，则该命令将失败。

示例

optimum-cli neuron cache add \
  --model prajjwal1/bert-tiny \
  --task text-classification \
  --train_batch_size 16 \
  --eval_batch_size 16 \
  --sequence_length 128 \
  --gradient_accumulation_steps 32 \
  --num_cores 32 \
  --precision bf16

这将为 prajjwal1/bert-tiny 模型在为指定参数设置的神经元缓存存储库上推送编译文件。

列出缓存存储库

请求缓存存储库以了解哪些编译文件可用也很方便。这可以通过 optimum-cli neuron cache list 命令完成

usage: optimum-cli neuron cache list [-h] [-m MODEL] [-v VERSION] [name]

positional arguments:
  name                  The name of the repo to list. Will use the locally saved cache repo if left unspecified.

optional arguments:
  -h, --help            show this help message and exit
  -m MODEL, --model MODEL
                        The model name or path of the model to consider. If left unspecified, will list all available models.
  -v VERSION, --version VERSION
                        The version of the Neuron X Compiler to consider. Will list all available versions if left unspecified.

如您所见，可以

  • 列出所有编译器版本可用的所有模型。
  • 通过指定 -v / --version 参数，列出给定编译器版本可用的所有模型。
  • 通过指定 -m / --model 参数，列出给定模型的所有编译文件，对于不同的输入形状等等，可能会有很多。

示例

optimum-cli neuron cache list aws-neuron/optimum-neuron-cache

←使用 Amazon Sagemaker 进行训练和部署 使用 AWS Trainium 微调 Transformers→