又一个节省金钱的AI模型训练技巧
本文的主题是AWS自家研发的AI芯片AWS Inferentia，更具体地说是第二代AWS Inferentia2。这是去年我们关于AWS Trainium的文章的续集，并且属于一系列关于专用AI加速器的文章。与我们之前的系列文章中探索的芯片相反，AWS Inferentia是专为AI模型推理设计的，并且专门针对深度学习推理应用。然而，AWS Inferentia2和AWS Trainium都共享相同的底层NeuronCore-v2架构和相同的软件堆栈（AWS Neuron SDK），这就引发了一个问题：AWS Inferentia能够用于AI训练负载吗？
诚然，在与Amazon EC2 Inf2实例系列（由AWS Inferentia加速器提供支持）相比较时，存在一些元素可能使其在某些训练负载上不太适用，与Amazon EC2 Trn1实例系列相比较也是如此。例如，尽管Inf2和都支持高带宽和低延迟的NeuronLink-v2设备间互连，但Trainium设备采用2D Torus拓扑连接，而不是环形拓扑连接，这可能会对集体通信运算符的性能产生影响（详见此处了解更多详情）。然而，某些训练负载可能不需要Trn1架构的独特特性，而在Inf1和Inf2架构上表现同样出色。
事实上，能够在Trainium和Inferentia加速器上进行训练将极大地增加我们可供选择的训练实例的多样性，并且能够根据每个深度学习项目的具体需求灵活选择训练实例。在我们最近的文章《深度学习的实例选择》中，我们详细阐述了拥有各种多样的实例类型对于深度学习训练的价值。而Trn1系列仅包含两种实例类型，训练Inf2将增加四种额外的实例类型。再加上Inf1，将再增加四个实例类型。
本文的目的是演示在AWS Inferentia上进行训练的机会。我们将定义一个玩具视觉模型，并比较在Amazon EC2 Trn1和Amazon EC2 Inf2实例系列上进行训练时的性能。非常感谢Ohad Klein和Yitzhak Levi对本文的贡献。
免责声明
- 请注意，在撰写本文时，仍有一些DL模型架构不受Neuron SDK支持。例如，虽然支持CNN模型推理，但仍不支持训练CNN模型。SDK文档包括一个模型支持矩阵，详细说明每个模型架构、训练框架（如TensorFlow和PyTorch）和Neuron架构版本的支持功能。
- 我们描述的实验是在Amazon EC2上使用最新版本的Neuron Deep Learning AMI进行的，该版本为“Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 20.04) 20230720”，其中包括Neuron SDK的2.8版本。由于Neuron SDK仍在积极开发中，我们所获得的比较结果很可能会随时间变化。强烈建议您使用最新版本的底层库重新评估本文的研究结果。
- 本文的目的是演示在AWS Inferentia驱动的实例上进行训练的潜力。请不要将本文视为对这些实例或我们提到的其他任何产品的认可。选择训练环境涉及许多变量，可能会因您项目的具体情况而大不相同。特别是，不同的模型在两种不同实例类型上运行时可能会展现出完全不同的相对价格和性能结果。
玩具模型
与我们在之前的文章中描述的实验类似，我们定义了一个简单的Vision Transformer（ViT）支持的分类模型（使用timm Python包版本0.9.5），以及一个随机生成的数据集。
from torch.utils.data import Datasetimport time, osimport torchimport torch_xla.core.xla_model as xmimport torch_xla.distributed.parallel_loader as plfrom timm.models.vision_transformer import VisionTransformer
# 使用随机数据class FakeDataset(Dataset): def __len__(self): return 1000000 def __getitem__(self, index): rand_image = torch.randn([3, 224, 224], dtype=torch.float32) label = torch.tensor(data=[index % 1000], dtype=torch.int64) return rand_image, labeldef train(batch_size=16, num_workers=4): # 为torchrun初始化XLA进程组 import torch_xla.distributed.xla_backend torch.distributed.init_process_group('xla') # 多进程：确保每个worker具有相同的初始权重 torch.manual_seed(0) dataset = FakeDataset() model = VisionTransformer() # 将模型加载到XLA设备上 device = xm.xla_device() model = model.to(device) optimizer = torch.optim.Adam(model.parameters()) data_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=num_workers) data_loader = pl.MpDeviceLoader(data_loader, device) loss_function = torch.nn.CrossEntropyLoss() summ, tsumm = 0, 0 count = 0 for step, (inputs, target) in enumerate(data_loader, start=1): t0 = time.perf_counter() inputs = inputs.to(device) targets = torch.squeeze(target.to(device), -1) optimizer.zero_grad() outputs = model(inputs) loss = loss_function(outputs, targets) loss.backward() xm.optimizer_step(optimizer) batch_time = time.perf_counter() - t0 if idx > 10: # 跳过前几个步骤 summ += batch_time count += 1 t0 = time.perf_counter() if idx > 500: break print(f'平均步骤时间：{summ/count}')if __name__ == '__main__': os.environ['XLA_USE_BF16'] = '1' # 根据vCPUs的数量设置数据加载器的工作线程数 # 例如，trn1为4，inf2.xlarge为2，inf2.12xlarge和inf2.48xlarge为8 train(num_workers=4)# 初始化命令：# torchrun --nproc_per_node=2 train.py结果
在下面的表格中，我们比较了各种Amazon EC2 Trn1和Amazon EC2 Inf2实例类型的速度和价格性能。
显然，Trainium驱动的实例类型支持更好的绝对性能（即增加训练速度），而在Inferentia驱动的实例上训练可以获得约39%更好的价格性能（对于双核实例类型），而对于较大的实例类型则更高。
再次强调，不要仅仅根据这些结果来做任何设计决策。某些模型架构可能在Trn1实例上运行成功，但在Inf2上会发生故障。其他模型可能在两者上都成功，但其比较性能结果与此处显示的结果可能截然不同。
请注意，我们省略了编译DL模型所需的时间。虽然这只在第一次运行模型时需要，但编译时间可能非常长（例如，对于我们的玩具模型，可能超过十分钟）。减少模型编译开销的两种方法是并行编译和离线编译。重要的是确保您的脚本不包含会触发频繁重新编译的操作（或图形更改）。有关更多详细信息，请参阅Neuron SDK文档。
总结
尽管AWS Inferentia被宣传为AI推理芯片，但看起来它也为训练深度学习模型提供了另一种选择。在我们之前关于AWS Trainium的帖子中，我们强调了将模型调整为在新的AI ASIC上进行训练时可能遇到的一些挑战。在AWS Inferentia驱动的实例类型上训练相同的模型，也能增加您的努力的潜在回报。