diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 7226a909b..cde900031 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -58,6 +58,8 @@ title: Neuron Trainer - local: package_reference/distributed title: Neuron Distributed + - local: package_reference/supported_models + title: Supported Architectures - local: package_reference/export title: Neuron Exporter - local: package_reference/modeling diff --git a/docs/source/package_reference/export.mdx b/docs/source/package_reference/export.mdx index 7f0102ecf..0d521e71c 100644 --- a/docs/source/package_reference/export.mdx +++ b/docs/source/package_reference/export.mdx @@ -28,40 +28,4 @@ exporting function according to the environment. Besides, you can check if the exported model is valid via [`~optimum.exporters.neuron.convert.validate_model_outputs`], which compares the compiled model's output on Neuron devices to the PyTorch model's output on CPU. -## Supported architectures - -| 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 | -| FlauBERT | feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification | -| GPT2 | text-generation | -| Llama, Llama 2 | text-generation | -| Mistral | 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 | -| 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 | -| 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 | -| 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 | - - - - -More details for checking supported tasks [here](https://huggingface.co/docs/optimum-neuron/guides/export_model#selecting-a-task). - - - -More architectures coming soon, stay tuned! 🚀 diff --git a/docs/source/package_reference/supported_models.mdx b/docs/source/package_reference/supported_models.mdx new file mode 100644 index 000000000..6f0ce94ca --- /dev/null +++ b/docs/source/package_reference/supported_models.mdx @@ -0,0 +1,71 @@ + + +# 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 | +| 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 | +| 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](https://huggingface.co/docs/optimum-neuron/guides/export_model#selecting-a-task). + + + +More architectures coming soon, stay tuned! 🚀 \ No newline at end of file