快速入门

🤗 Optimum Neuron 的设计目标只有一个：让任何 🤗 Transformers 用户都能轻松进行训练和推理，同时充分利用 AWS 加速器的强大功能。

训练

有三个主要的类需要了解

NeuronArgumentParser：继承了 Transformers 中的原始 HfArgumentParser，并对参数值进行了额外的检查，以确保它们能与 AWS Trainium 实例良好配合。
NeuronTrainingArguments：扩展了原始的 TrainingArguments，添加了特定于 Trainium 的参数，以优化在 AWS Trainium 硬件上的性能。
NeuronTrainer：负责编译和分发模型以在 Trainium 芯片上运行，并执行训练和评估的 trainer 类。

NeuronTrainer 与 NeuronTrainer 非常相似，为 🤗 Transformers Trainer 实例调整脚本主要包括：

将 Trainer 类替换为 NeuronTrainer 类
使用 NeuronTrainingArguments 代替 TrainingArguments

-from transformers import TrainingArguments
+from optimum.neuron import NeuronTrainingArguments
-from transformers import Trainer
+from optimum.neuron import NeuronTrainer as Trainer
-training_args = TrainingArguments(
+training_args = NeuronTrainingArguments(
  # training arguments...
)
# A lot of code here
# Initialize our Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset if training_args.do_train else None,
    eval_dataset=eval_dataset if training_args.do_eval else None,
    compute_metrics=compute_metrics,
    tokenizer=tokenizer,
    data_collator=data_collator,
)

所有 Trainium 实例至少配备 2 个 Neuron Core。为了利用这些核心，我们需要使用 torchrun 启动训练。下面您将看到一个示例，说明如何在 trn1.2xlarge 实例上使用 bert-base-uncased 模型启动训练脚本。

torchrun --nproc_per_node=2 huggingface-neuron-samples/text-classification/run_glue.py \
--model_name_or_path bert-base-uncased \
--dataset_name philschmid/emotion \
--do_train \
--do_eval \
--bf16 True \
--per_device_train_batch_size 16 \
--learning_rate 5e-5 \
--num_train_epochs 3 \
--output_dir ./bert-emotion

推理

您可以在 Neuron 设备上进行推理之前，将您的 🤗 Transformers 模型编译并导出为序列化格式

optimum-cli export neuron \
  --model distilbert-base-uncased-finetuned-sst-2-english \
  --batch_size 1 \
  --sequence_length 32 \
  --auto_cast matmul \
  --auto_cast_type bf16 \
  distilbert_base_uncased_finetuned_sst2_english_neuron/

上面的命令将导出具有静态形状的 distilbert-base-uncased-finetuned-sst-2-english：batch_size=1 和 sequence_length=32，并将所有 matmul 操作从 FP32 转换为 BF16。查看导出器指南以获取更多编译选项。

然后，您可以使用 NeuronModelForXXX 类在 Neuron 设备上运行导出的 Neuron 模型，这些类类似于 🤗 Transformers 中的 AutoModelForXXX 类

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_sst2_english_neuron")

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
inputs = tokenizer("Hamilton is considered to be the best musical of past years.", return_tensors="pt")

logits = model(**inputs).logits
print(model.config.id2label[logits.argmax().item()])
# 'POSITIVE'

AWS Trainium 和 Inferentia

快速入门

训练

推理