AWS Trainium & Inferentia documentation

Supported architectures

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Supported architectures

Transformers

Architecture Task
ALBERT feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
BERT feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
BLOOM text-generation
CamemBERT feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
ConvBERT feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
DeBERTa (INF2 only) feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
DeBERTa-v2 (INF2 only) feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
DistilBERT feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
ELECTRA feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
ESM feature-extraction, fill-mask, text-classification, token-classification
FlauBERT feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
GPT2 text-generation
Llama, Llama 2 text-generation
Mistral text-generation
Mixtral text-generation
MobileBERT feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
MPNet feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
OPT text-generation
Phi feature-extraction, text-classification, token-classification
RoBERTa feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
RoFormer feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
T5 text2text-generation
XLM feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification
XLM-RoBERTa feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification

Diffusers

Architecture Task
Stable Diffusion text-to-image, image-to-image, inpaint
Stable Diffusion XL Base text-to-image, image-to-image, inpaint
Stable Diffusion XL Refiner image-to-image, inpaint
SDXL Turbo text-to-image, image-to-image, inpaint
LCM text-to-image

Sentence Transformers

Architecture Task
Transformer feature-extraction, sentence-similarity
CLIP feature-extraction, zero-shot-image-classification

More details for checking supported tasks here.

More architectures coming soon, stay tuned! 🚀