支持的架构

训练

在 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	特征提取、零样本图像分类