支援的架構

訓練

在 AWS Trainium 例項（Trn1）上進行訓練可實現採用分散式並行策略的大規模模型訓練。

要求

• 模型必須與 Neuron SDK 相容。如果模型足夠小，可以容納在 16GB 記憶體內，那麼任何能夠成功編譯的架構都支援訓練。
• 記憶體限制： 每個加速器有 16GB 記憶體用於存放模型權重、梯度、最佳化器狀態和啟用值。
• 對於大型模型： 需要支援張量並行和/或流水線並行的自定義模型實現。

以下架構具有支援分散式訓練的自定義模型實現

架構	任務	張量並行	流水線並行
Llama、Llama 2、Llama 3	文字生成	✓	✓
Qwen3	文字生成	✓	✓
Granite	文字生成	✓	✗

推理

下表列出了 Optimum Neuron 在 Amazon EC2 Inf2 例項上支援的推理架構和任務。

Transformers

架構	任務
ALBERT	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
AST	特徵提取、音訊分類
BERT	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
Beit	特徵提取、影像分類
CamemBERT	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
CLIP	特徵提取、影像分類
ConvBERT	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
ConvNext	特徵提取、影像分類
ConvNextV2	特徵提取、影像分類
CvT	特徵提取、影像分類
DeBERTa (僅限 INF2)	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
DeBERTa-v2 (僅限 INF2)	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
Deit	特徵提取、影像分類
DistilBERT	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
DonutSwin	特徵提取
Dpt	特徵提取
ELECTRA	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
ESM	特徵提取、掩碼填充、文字分類、詞元分類
FlauBERT	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
Granite	文字生成
Hubert	特徵提取、自動語音識別、音訊分類
Levit	特徵提取、影像分類
Llama、Llama 2、Llama 3	文字生成
Mixtral	文字生成
MobileBERT	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
MobileNetV2	特徵提取、影像分類、語義分割
MobileViT	特徵提取、影像分類、語義分割
ModernBERT	特徵提取、掩碼填充、文字分類、詞元分類
MPNet	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
Phi3	文字生成
Phi	特徵提取、文字分類、詞元分類
Qwen2、Qwen 3	文字生成
RoBERTa	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
RoFormer	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
Swin	特徵提取、影像分類
T5	文字到文字生成
UniSpeech	特徵提取、自動語音識別、音訊分類
UniSpeech-SAT	特徵提取、自動語音識別、音訊分類、音訊幀分類、音訊 x-vector
ViT	特徵提取、影像分類
Wav2Vec2	特徵提取、自動語音識別、音訊分類、音訊幀分類、音訊 x-vector
WavLM	特徵提取、自動語音識別、音訊分類、音訊幀分類、音訊 x-vector
Whisper	自動語音識別
XLM	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
XLM-RoBERTa	特徵提取、掩碼填充、多項選擇、問答、文字分類、詞元分類
Yolos	特徵提取、目標檢測

Diffusers

架構	任務
Stable Diffusion	文字到影像、影像到影像、修復
Stable Diffusion XL Base	文字到影像、影像到影像、修復
Stable Diffusion XL Refiner	影像到影像、修復
SDXL Turbo	文字到影像、影像到影像、修復
LCM	文字到影像
PixArt-α	文字到影像
PixArt-Σ	文字到影像
Flux	文字到影像

Sentence Transformers

架構	任務
Transformer	特徵提取、句子相似度
CLIP	特徵提取、零樣本影像分類