

Implications
Transformers models based on the BERT (Bidirectional Encoder Representations from
Transformers)
architecture, or its variants such as
distilBERT and
roBERTa run best on
Inf1 for non-generative tasks such as extractive question answering, sequence
classification, and token classification. However, text generation tasks can still be
adapted to run on Inf1 according to this AWS Neuron MarianMT
tutorial.
More information about models that can be converted out of the box on Inferentia can be
found in the Model Architecture
Fit
section of the Neuron documentation.
Dependencies
Using AWS Neuron to convert models requires a Neuron SDK
environment
which comes preconfigured on AWS Deep Learning
AMI.
Converting a model for AWS Neuron
Convert a model for AWS NEURON using the same code from Using TorchScript in
Python to trace a BertModel. Import the
torch.neuron framework extension to access the components of the Neuron SDK through a
Python API:
python
from transformers import BertModel, BertTokenizer, BertConfig
import torch
import torch.neuron
You only need to modify the following line:
diff
- torch.jit.trace(model, [tokens_tensor, segments_tensors])
+ torch.neuron.trace(model, [token_tensor, segments_tensors])
This enables the Neuron SDK to trace the model and optimize it for Inf1 instances.
To learn more about AWS Neuron SDK features, tools, example tutorials and latest
updates, please see the AWS NeuronSDK
documentation.