神经元模型缓存

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

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

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

在训练 Transformer 或扩散模型或在 Neuron 平台上加载 NeuronModelForCausalLM 之前，需要使用 torch-neuronx 将其导出为神经元格式。

导出模型时，torch-neuronx 将

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

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

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

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

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

如何使用神经元模型缓存

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

将模型导出为神经元格式时，optimum-neuron 将在编译模型子图期间直接在 hub 存储库中查找缓存的 NEFF 文件。

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

缓存机制

Optimum Neuron Cache 基于 NeuronX 编译器缓存。

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

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

每个子图都是唯一的，并且是以下内容的组合结果：

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

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

Neuron X 编译器的版本，
使用的 Neuron 内核数量，
编译参数（例如优化级别）。

所有这些参数组合在一起，创建了一个唯一的哈希值，用于标识 NEFF 文件。

这有两个非常重要的后果

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

因此，很难预先知道与特定模型配置相关的 NEFF 文件是否已缓存。

Neuron 模型缓存查找（仅限 Inferentia）

Neuron 缓存查找是一个功能，允许用户在导出模型以供推断之前查找兼容的缓存模型配置。

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

缓存的模型配置被存储为 Neuron 模型缓存中特定子文件夹下的条目

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

每个条目对应于模型配置及其导出参数的组合：这是我们所能做到的最接近唯一标识已导出模型的方法。

您可以使用 optimum-cli 通过传递 Hub 模型_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）

如何使用私有 Neuron 模型缓存（仅限 Trainium）

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

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

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

使用 Optimum CLI

Optimum CLI 提供了两个用于缓存创建和设置的子命令

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

创建新的 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 选项允许您为 Neuron 缓存存储库指定名称，如果未设置，将使用默认名称。--public 标志允许您将 Neuron 缓存设为公开，因为它将默认情况下创建为私有存储库。

示例

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.

设置不同的 Trainiun 缓存存储库

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]

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

确保要使用的 Neuron 缓存存储库已在本地设置，这可以通过运行 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

这将把针对指定参数设置的 Neuron 缓存仓库中 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

AWS Trainium 和 Inferentia