Neuron 模型推理

以下文件中提供的 API 適用於 inf2、trn1 和 inf1 上的推理。

NeuronModelForXXX 類有助於從 Hugging Face Hub 載入模型，並將其編譯為針對 Neuron 裝置最佳化的序列化格式。然後，您將能夠載入模型並使用 AWS Neuron 裝置提供的加速功能執行推理。

從 Transformers 切換到 Optimum

optimum.neuron.NeuronModelForXXX 模型類是與 Hugging Face Transformers 模型相容的 API。這意味著與 Hugging Face 的生態系統無縫整合。您只需將 AutoModelForXXX 類替換為 optimum.neuron 中對應的 NeuronModelForXXX 類即可。

如果您已經使用 Transformers，您只需替換模型類即可重複使用您的程式碼

from transformers import AutoTokenizer
-from transformers import AutoModelForSequenceClassification
+from optimum.neuron import NeuronModelForSequenceClassification

# PyTorch checkpoint
-model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")

+model = NeuronModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english",
+                                                             export=True, **neuron_kwargs)

如上所示，當您第一次使用 NeuronModelForXXX 時，您需要設定 export=True 以將模型從 PyTorch 編譯為 Neuron 相容格式。

您還需要傳遞 Neuron 特定引數來配置匯出。每個模型架構都有自己的一組引數，詳情如下所述。

模型匯出後，您可以將其儲存在本地或 Hugging Face 模型中心

# Save the neuron model
>>> model.save_pretrained("a_local_path_for_compiled_neuron_model")

# 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
... )

下次您想執行推理時，只需載入已編譯的模型，這將為您節省編譯時間

>>> from optimum.neuron import NeuronModelForSequenceClassification
>>> model = NeuronModelForSequenceClassification.from_pretrained("my-neuron-repo")

如您所見，無需傳遞匯出期間使用的 neuron 引數，因為它們已儲存在 config.json 檔案中，並將由 NeuronModelForXXX 類自動恢復。

首次執行推理時，首次執行管道會有一個預熱階段。此執行的延遲將比常規執行高 3-4 倍。

判別式 NLP 模型

如前一節所述，您只需對 Transformers 程式碼進行少量修改即可匯出和執行 NLP 模型

from transformers import AutoTokenizer
-from transformers import AutoModelForSequenceClassification
+from optimum.neuron import NeuronModelForSequenceClassification

# PyTorch checkpoint
-model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")

# Compile your model during the first time
+compiler_args = {"auto_cast": "matmul", "auto_cast_type": "bf16"}
+input_shapes = {"batch_size": 1, "sequence_length": 64}
+model = NeuronModelForSequenceClassification.from_pretrained(
+    "distilbert-base-uncased-finetuned-sst-2-english", export=True, **compiler_args, **input_shapes,
+)

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
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'

compiler_args 是編譯器的可選引數，這些引數通常控制編譯器如何在推理效能（延遲和吞吐量）和準確性之間進行權衡。在這裡，我們使用 Neuron 矩陣乘法引擎將 FP32 操作轉換為 BF16。

input_shapes 是您需要傳送給 Neuron 編譯器的強制性靜態形狀資訊。想知道您的模型哪些形狀是強制性的？使用以下程式碼檢視

>>> from transformers import AutoModelForSequenceClassification
>>> from optimum.exporters import TasksManager

>>> model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")

# Infer the task name if you don't know
>>> task = TasksManager.infer_task_from_model(model)  # 'text-classification'

>>> neuron_config_constructor = TasksManager.get_exporter_config_constructor(
...     model=model, exporter="neuron", task='text-classification'
... )
>>> print(neuron_config_constructor.func.get_mandatory_axes_for_task(task))
# ('batch_size', 'sequence_length')

請注意，用於編譯的輸入形狀應小於推理期間將饋入模型的輸入大小。

如果輸入大小小於編譯輸入形狀怎麼辦？

不用擔心，NeuronModelForXXX 類會將您的輸入填充到符合條件的形狀。此外，您可以在 from_pretrained 方法中設定 dynamic_batch_size=True 以啟用動態批處理，這意味著您的輸入可以具有可變批處理大小。

（請記住：動態性和填充不僅帶來靈活性，還會導致效能下降。這很公平！）

生成式 NLP 模型

如前所述，您只需對 Transformers 程式碼進行少量修改即可匯出和執行 NLP 模型

配置生成模型的匯出

與非生成模型一樣，可以將兩組引數傳遞給 from_pretrained() 方法，以配置如何將 Transformer 檢查點匯出為 Neuron 最佳化模型

compiler_args = { num_cores, auto_cast_type } 是編譯器的可選引數，這些引數通常控制編譯器如何在推理延遲和吞吐量以及準確性之間進行權衡。
input_shapes = { batch_size, sequence_length } 對應於模型輸入和 KV 快取（過去令牌的注意力鍵和值）的靜態形狀。
num_cores 是例項化模型時使用的神經核數量。每個神經核有 16 Gb 記憶體，這意味著更大的模型需要拆分到多個核上。預設為 1。
auto_cast_type 指定編碼權重的格式。它可以是 fp32 (float32)、fp16 (float16) 或 bf16 (bfloat16) 之一。預設為 fp32。
batch_size 是模型將接受的輸入序列的數量。預設為 1。
sequence_length 是輸入序列中的最大令牌數。預設為 max_position_embeddings（對於舊模型為 n_positions）。

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

# Instantiate and convert to Neuron a PyTorch checkpoint
+compiler_args = {"num_cores": 1, "auto_cast_type": 'fp32'}
+input_shapes = {"batch_size": 1, "sequence_length": 512}
-model = AutoModelForCausalLM.from_pretrained("gpt2")
+model = NeuronModelForCausalLM.from_pretrained("gpt2", export=True, **compiler_args, **input_shapes)

如前所述，這些引數只能在匯出期間配置。這意味著在推理期間尤其如此

輸入的 batch_size 應等於匯出期間使用的 batch_size，
輸入序列的 length 應低於匯出期間使用的 sequence_length，
最大令牌數（輸入 + 生成）不能超過匯出期間使用的 sequence_length。

文字生成推理

與原始 Transformer 模型一樣，使用 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", export=True)

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 方法的值仍然優先）。