人工智能最近取得的成功往往归功于 GPU 的出现和发展。GPU 的架构通常包括数千个多处理器、高速内存、专用张量内核等，特别适合满足 AI/ML 工作负载的密集需求。遗憾的是，人工智能开发的快速增长导致对 GPU 的需求激增，使其难以获得。因此，ML 开发人员越来越多地探索用于训练和运行模型的替代硬件选项。在这篇文章中，我们将回到老式的 CPU，重新审视它与 ML 应用的相关性。虽然与 GPU 相比，CPU 通常不太适合 ML 工作负载，但它们更容易获得。在 CPU 上运行（至少部分）工作负载的能力会对开发效率产生重大影响。

分析和优化人工智能/ML 工作负载运行时性能，这是加速开发和最小化成本的一种手段。虽然无论使用哪种计算引擎，这一点都至关重要，但不同平台的剖析工具和优化技术可能大相径庭。在本篇文章中，我们将讨论与 CPU 有关的一些性能优化选项。我们的重点是英特尔® 至强® CPU 处理器（采用英特尔® AVX-512）和 PyTorch（2.4 版）框架（尽管类似技术也可应用于其他 CPU 和框架）。更具体地说，我们将在带有 AWS 深度学习 AMI 的亚马逊 EC2 c7i 实例上运行我们的实验。

我们的目标是证明，虽然 CPU 上的 ML 开发可能不是我们的首选，但我们有办法 “减轻打击”，在某些情况下，甚至可以使其成为可行的替代方案。

注意：

我们在这篇文章中的意图只是展示 CPU 上可用的一些 ML 优化机会。与大多数以 CPU 上的 ML 优化为主题的在线教程相反，我们将重点关注训练工作负载而非推理工作负载。有许多专门针对推理的优化工具，我们不会涉及。
请勿将我们提及的任何工具或技术的官方文档的替代品。

示例--ResNet-50

我们将在一个以 ResNet-50 为骨干的简单图像分类模型（来自《图像识别深度残差学习》）上进行实验。我们将在一个假数据集上训练该模型。完整的训练脚本出现在下面的代码块中：

import torch
import torchvision
from torch.utils.data import Dataset, DataLoader
import time
# A dataset with random images and labels
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 % 10, dtype=torch.uint8)
        return rand_image, label
train_set = FakeDataset()
batch_size=128
num_workers=0
train_loader = DataLoader(
    dataset=train_set,
    batch_size=batch_size,
    num_workers=num_workers
)
model = torchvision.models.resnet50()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters())
model.train()
t0 = time.perf_counter()
summ = 0
count = 0
for idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    batch_time = time.perf_counter() - t0
    if idx > 10:  # skip first steps
        summ += batch_time
        count += 1
    t0 = time.perf_counter()
    if idx > 50:
        break
print(f'average step time: {summ/count}')
print(f'throughput: {count*batch_size/summ}')

在 c7i.2xlarge（带 8 个 vCPU）和 PyTorch 2.4 CPU 版本上运行此脚本，吞吐量为每秒 9.12 个样本。为了便于比较，我们注意到在 Amazon EC2 g5.2xlarge 实例（带 1 个 GPU 和 8 个 vCPU）上，相同（未优化脚本）的吞吐量为每秒 340 个样本。考虑到这两种实例类型的比较成本，我们发现在 GPU 实例上进行训练的性价比大约高出 11（！！）倍。基于这些结果，使用 GPU 训练 ML 模型的偏好是非常有道理的。让我们来评估一下缩小这一差距的可能性。

PyTorch 性能优化

在本节中，我们将探讨一些提高训练工作负载运行时性能的基本方法。虽然你可能会从关于 GPU 优化的文章中认识到其中的一些方法，但有必要强调一下 CPU 和 GPU 平台上训练优化的显著区别。在 GPU 平台上，我们主要致力于最大限度地提高 CPU（训练数据预处理）与 GPU（模型训练）之间的并行性。在 CPU 平台上，所有处理都在 CPU 上进行，我们的目标是最有效地分配 CPU 资源。

批量大小

增加训练批次大小可以降低模型参数更新的频率，从而提高性能。(在 GPU 上，它还能减少 CPU 和 GPU 之间的事务开销，如内核加载）。不过，在 GPU 上，我们的目标是最大限度地利用 GPU 内存，而在 CPU 上，同样的策略可能会损害性能。由于超出本篇文章讨论范围的原因，CPU 内存更为复杂，发现最佳批处理大小的最佳方法可能是反复试验。请记住，改变批次大小可能会影响训练的收敛性。

下表总结了在任意选择批次大小的情况下，我们训练工作负载的吞吐量：

与我们在 GPU 上的发现相反，在 c7i.2xlarge 实例类型上，我们的模型似乎更倾向于较低的批量大小。

多进程数据加载

GPU 上的一种常用技术是为数据加载器分配多个进程，以降低 GPU 饥饿的可能性。在 GPU 平台上，一般的经验法则是根据 CPU 内核的数量来设置工作进程的数量。但是，在 CPU 平台上，模型训练与数据加载器使用相同的资源，这种方法可能会适得其反。同样，选择最佳工作者数量的最佳方法可能是反复试验。下表显示了不同工人数选择下的平均吞吐量：

混合精度

另一种流行的技术是使用较低精度的浮点数据类型，如 torch.float16 或 torch.bfloat16，一般认为 torch.bfloat16 的动态范围更适合 ML 训练。当然，降低数据类型精度可能会对收敛性产生不利影响，因此应谨慎操作。PyTorch 自带 torch.amp，这是一个自动混合精度包，用于优化这些数据类型的使用。Intel® AVX-512 包含对 bfloat16 数据类型的支持。修改后的训练步骤如下所示：

for idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    with torch.amp.autocast('cpu',dtype=torch.bfloat16):
        output = model(data)
        loss = criterion(output, target)
    loss.backward()
    optimizer.step()

优化后的吞吐量为每秒 24.34 个采样点，提高了 86%！

通道最后内存格式

最后通道内存格式是一项测试级优化，主要与视觉模型有关，支持在内存中存储四维（NCHW）张量，将通道作为最后一个维度。这样，每个像素的所有数据都被存储在一起。这种优化主要适用于视觉模型。这种内存格式被认为对英特尔平台更 “友好”，据报告，它提高了英特尔® 至强® CPU 上 ResNet-50 的性能。调整后的训练步骤如下：

for idx, (data, target) in enumerate(train_loader):
    data = data.to(memory_format=torch.channels_last)
    optimizer.zero_grad()
    with torch.amp.autocast('cpu',dtype=torch.bfloat16):
        output = model(data)
        loss = criterion(output, target)
    loss.backward()
    optimizer.step()

由此产生的吞吐量为每秒 37.93 个采样点，与基线实验相比提高了 56%，总计提高了 415%。我们正在发挥作用

火炬编译

在此，我们将仅限于评估默认（TorchInductor）后端和 PyTorch 的英特尔® 扩展库（Intel® Extension for PyTorch）中的 ipex 后端，该库专门针对英特尔硬件进行了优化。更新后的模型定义如下：

import intel_extension_for_pytorch as ipex
model = torchvision.models.resnet50()
backend='inductor' # optionally change to 'ipex'
model = torch.compile(model, backend=backend)

在我们的玩具模型中，火炬编译的影响只有在禁用 “通道最后 ”优化时才会显现出来（每个后端都增加了约 27%）。当应用 “通道最后 ”时，性能实际上会下降。因此，我们在后续实验中放弃了这一优化。

内存和线程优化

有很多机会可以优化底层 CPU 资源的使用。其中包括根据底层 CPU 硬件结构优化内存管理和线程分配。可以通过使用高级内存分配器（如 Jemalloc 和 TCMalloc）和/或减少速度较慢的内存访问（如跨 NUMA 节点）来改进内存管理。通过适当配置 OpenMP 线程库和/或使用英特尔的 Open MP 库，可以改善线程分配。

一般来说，这类优化需要深入了解 CPU 架构及其支持的 SW 栈的特性。为了简化操作，PyTorch 提供了 torch.backends.xeon.run_cpu 脚本，用于自动配置内存和线程库，以优化运行时性能。下面的命令将导致使用专用内存和线程库。在讨论分布式培训选项时，我们将回到 NUMA 节点的话题。

我们将验证 TCMalloc（conda install conda-forge::gperftools）和英特尔 Open MP 库（pip install intel-openmp）的安装是否正确，并运行以下命令。

python -m torch.backends.xeon.run_cpu train.py

使用 run_cpu 脚本后，我们的运行性能进一步提高到每秒 39.05 个采样点。请注意，run_cpu 脚本包含许多可进一步调整性能的控件。

PyTorch 的英特尔扩展

PyTorch 的英特尔® 扩展通过其 ipex.optimize 函数为训练优化提供了更多机会。在此，我们将演示其默认用法。

 model = torchvision.models.resnet50()
 criterion = torch.nn.CrossEntropyLoss()
 optimizer = torch.optim.SGD(model.parameters())
 model.train()
 model, optimizer = ipex.optimize(
    model, 
    optimizer=optimizer,
    dtype=torch.bfloat16
 )

结合上文讨论的内存和线程优化，吞吐量达到每秒 40.73 个采样点。(请注意，在禁用 “通道最后 ”配置时也能达到类似的结果）。

CPU 上的分布式训练

英特尔® 至强® 处理器采用非统一内存访问（NUMA）设计，CPU 内存被划分为若干组，即 NUMA 节点，每个 CPU 内核被分配到一个节点。虽然任何 CPU 内核都可以访问任何 NUMA 节点的内存，但访问自己的节点（即本地内存）要快得多。这就产生了跨 NUMA 节点分布式训练的概念，其中分配给每个 NUMA 节点的 CPU 内核就像分布式进程组中的单个进程，而跨节点的数据分布则由英特尔® oneCCL（英特尔专用的集体通信库）管理。

我们可以使用 ipexrun 实用程序在 NUMA 节点间轻松运行数据分布训练。在下面的代码块中，我们将调整脚本以运行数据分布式训练（根据此处详细介绍的用法）：

import os, time
import torch
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
import torch.distributed as dist
import torchvision
import oneccl_bindings_for_pytorch as torch_ccl
import intel_extension_for_pytorch as ipex

os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = "29500"
os.environ["RANK"] = os.environ.get("PMI_RANK", "0")
os.environ["WORLD_SIZE"] = os.environ.get("PMI_SIZE", "1")
dist.init_process_group(backend="ccl", init_method="env://")
rank = os.environ["RANK"]
world_size = os.environ["WORLD_SIZE"]
batch_size = 128
num_workers = 0
# define dataset and dataloader
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 % 10, dtype=torch.uint8)
        return rand_image, label
train_dataset = FakeDataset()
dist_sampler = DistributedSampler(train_dataset)
train_loader = DataLoader(
    dataset=train_dataset, 
    batch_size=batch_size,
    num_workers=num_workers,
    sampler=dist_sampler
)
# define model artifacts
model = torchvision.models.resnet50()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters())
model.train()
model, optimizer = ipex.optimize(
    model, 
    optimizer=optimizer,
    dtype=torch.bfloat16
)
# configure DDP
model = torch.nn.parallel.DistributedDataParallel(model)
# run training loop
# destroy the process group
dist.destroy_process_group()

遗憾的是，截至本文撰写之时，Amazon EC2 c7i 实例系列不包括多 NUMA 实例类型。为了测试我们的分布式培训脚本，我们恢复到具有 64 个 vCPU 和 2 个 NUMA 节点的 Amazon EC2 c6i.32xlarge 实例。我们验证了 PyTorch 的英特尔® oneCCL 绑定的安装情况，并运行以下命令（如此处所记录）：

source $(python -c "import oneccl_bindings_for_pytorch as torch_ccl;print(torch_ccl.cwd)")/env/setvars.sh
# This example command would utilize all the numa sockets of the processor, taking each socket as a rank.
ipexrun cpu --nnodes 1 --omp_runtime intel train.py 

下表比较了 c6i.32xlarge 实例在进行和未进行分布式训练时的性能结果：

在我们的实验中，数据分布并没有提高运行性能。

使用 Torch/XLA 进行 CPU 训练

在之前我们讨论了 PyTorch/XLA 库及其对 XLA 编译的使用，以便在 TPU、GPU 和 CPU 等 XLA 设备上实现基于 PyTorch 的训练。与火炬编译类似，XLA 使用图编译生成针对目标设备优化的机器代码。随着 OpenXLA 项目的成立，其既定目标之一就是支持包括 CPU 在内的所有硬件后端实现高性能（包括 CPU RFC）。下面的代码块展示了使用 PyTorch/XLA 进行训练所需对原始（未优化）脚本进行的调整：

import torch
import torchvision
import timeimport torch_xla
import torch_xla.core.xla_model as xm

device = xm.xla_device()
model = torchvision.models.resnet50().to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters())
model.train()
for idx, (data, target) in enumerate(train_loader):
    data = data.to(device)
    target = target.to(device)
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    xm.mark_step()

遗憾的是（截至本文撰写之时），XLA 在我们的玩具模型上取得的结果似乎远不如我们上面看到的（未优化的）结果（--相差 7 倍之多）。我们希望随着 PyTorch/XLA CPU 支持的成熟，这种情况会有所改善。

测试结果

我们在下表中总结了部分实验结果。为了便于比较，我们添加了在亚马逊 EC2 g5.2xlarge GPU 实例上按照本文章讨论的优化步骤训练模型的吞吐量。每美元样本的计算基于亚马逊 EC2 按需定价页面（截至本文撰写时，c7i.2xlarge 每小时 0.357 美元，g5.2xlarge 每小时 1.212 美元）。

虽然我们成功地将玩具模型在 CPU 实例上的训练性能提高了相当大的幅度（446%），但仍然不如在 GPU 实例上的（优化）性能。根据我们的结果，在 GPU 上的训练成本要低 6.7 倍。通过额外的性能调整和/或应用额外的优化策略，我们有可能进一步缩小差距。我们再次强调，我们得出的性能比较结果是该模型和运行环境所独有的。

总结

鉴于 CPU 的普遍性，能否有效地利用 CPU 来训练和/或运行 ML 工作负载，对开发效率和最终产品部署策略都会产生巨大影响。虽然与 GPU 相比，CPU 架构的性质对许多 ML 应用程序不太友好，但仍有许多工具和技术可用于提升其性能，我们在本篇文章中讨论并演示了其中的一些精选工具和技术。