最近人工智能的成功往往归功于GPU的出现和发展。GPU的架构通常包括数千个多处理器、高速内存、专用张量核心等，特别适合满足AI/ML工作负载的密集需求。然而，人工智能开发的迅速增长导致了对GPU需求的激增，使得GPU变得难以获得。因此，ML开发者越来越多地探索替代硬件选项来训练和运行他们的模型。在之前的帖子中，我们讨论了在专用AI ASIC（如Google Cloud TPU，Haban Gaudi，和AWS Trainium）上进行训练的可能性。虽然这些选项提供了显著的成本节约机会，但它们并不适合所有ML模型，并且也可能像GPU一样面临供应不足的问题。在这篇文章中，我们回到传统的CPU，并重新审视其在ML应用中的相关性。尽管与GPU相比，CPU在处理ML工作负载方面通常不太适合，但它们更容易获得。能够在CPU上运行（至少部分）我们的工作负载，可能对开发生产力产生重大影响。
在之前的帖子（例如，这里）中，我们强调了分析和优化AI/ML工作负载的运行时性能对于加速开发和降低成本的重要性。虽然这一点无论使用哪种计算引擎都很重要，但不同平台的分析工具和优化技术差异很大。在本文中，我们将讨论一些与CPU相关的性能优化选项。我们的重点是Intel® Xeon® CPU处理器（带有Intel® AVX-512）以及PyTorch（版本2.4）框架（尽管类似的技术也可以应用于其他CPU和框架）。更具体地说，我们将在Amazon EC2 c7i实例上运行实验，该实例使用AWS Deep Learning AMI。请不要将我们选择的云平台、CPU版本、机器学习框架或其他任何工具或库视为对其替代品的推荐。
我们的目标是证明尽管在CPU上进行机器学习开发可能不是我们的首选，但有一些方法可以“缓解这种影响”，在某些情况下，甚至可能使其成为一个可行的替代方案。
声明注意事项本文的目的是展示在CPU上可用的一些机器学习优化机会。与大多数关于CPU上机器学习优化的在线教程不同，我们将重点放在训练工作负载上，而不是推理工作负载。有许多专门针对推理的优化工具，我们不会涉及（例如，参见这里和这里）。
请不要将此帖子视为我们提到的任何工具或技术的官方文档的替代品。请记住，鉴于AI/ML开发的快速步伐，您阅读本文时，我们提到的一些内容、库和/或说明可能会过时。请务必参考最新版本的文档。
重要的是，我们讨论的优化对运行时性能的影响可能会因模型和环境的细节（例如，参见官方 PyTorch TouchInductor CPU 推理性能仪表板中不同模型之间的高度差异）而有很大不同。我们将分享的比较性能数据仅适用于我们使用的示例模型和运行时环境。请务必在自己的模型和运行时环境中重新评估所有建议的优化。
最后，我们的重点将仅放在吞吐量性能（以每秒样本数衡量）上——而不是训练收敛性。然而，需要注意的是，一些优化技术（例如批量大小调整、混合精度等）可能会对某些模型的收敛性产生负面影响。在某些情况下，可以通过适当的超参数调整来克服这一点。
示例 — ResNet-50我们将使用一个简单的图像分类模型进行实验，该模型具有 ResNet-50 主干（来自 Deep Residual Learning for Image Recognition）。我们将使用一个假数据集来训练该模型。完整的训练脚本如下所示（大致基于 这个示例）：
import torch
import torchvision
from torch.utils.data import Dataset, DataLoader
import time
# 一个包含随机图像和标签的数据集
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: # 跳过前几步
summ += batch_time
count += 1
t0 = time.perf_counter()
if idx > 100:
break
print(f'平均步时间: {summ/count}')
print(f'吞吐量: {count*batch_size/summ}')运行此脚本在带有8个vCPU的c7i.2xlarge实例上，并使用PyTorch 2.4的CPU版本，每秒处理9.12个样本。为了便于比较，我们注意到在带有1个GPU和8个vCPU的Amazon EC2 g5.2xlarge实例上运行相同的（未优化的脚本），每秒处理340个样本。考虑到这两种实例类型的相对成本（截至撰写本文时，c7i.2xlarge每小时0.357美元，g5.2xlarge每小时1.212美元），我们发现使用GPU实例进行训练在性价比上大约提高了十一倍（!!）。根据这些结果，使用GPU训练机器学习模型的偏好非常有根据。让我们评估一些减少这种差距的可能性。
PyTorch 性能优化在本节中，我们将探讨一些基本方法来提高我们的训练工作负载的运行时性能。虽然你可能从我们关于GPU优化的文章帖子中认识到了一些内容，但重要的是要强调在CPU和GPU平台上进行训练优化之间存在显著差异。在GPU平台上，我们大部分的努力都集中在最大化CPU和GPU之间的并行化（即在CPU上进行训练数据预处理和在GPU上进行模型训练）。而在CPU平台上，所有的处理都在CPU上进行，我们的目标将是尽可能有效地分配其资源。
批量大小增加训练批次大小可以潜在地通过减少模型参数更新的频率来提高性能。（在GPU上，这还有减少内核加载等CPU-GPU事务开销的好处）。然而，虽然在GPU上我们希望找到最大化利用GPU内存的批次大小，同样的策略在CPU上可能会降低性能。由于超出本文范围的原因，CPU内存更复杂，找到最优化的批次大小的最佳方法可能是通过试错。请记住，改变批次大小可能会影响训练的收敛性。
下表总结了我们训练工作负载在几个（任意选择的）批次大小下的吞吐量：
批量大小的训练吞吐量函数（作者提供）
与我们在GPU上的发现相反，在c7i.2xlarge实例类型上，我们的模型似乎更喜欢使用较小的批次大小。
多进程数据加载在GPU上，一种常见的技术是为数据加载器分配多个进程，以减少GPU饥饿的可能性。在GPU平台上，一般的经验法则是根据CPU核心数设置进程数量。然而，在CPU平台上，由于模型训练和数据加载器使用相同的资源，这种方法可能会适得其反。再次强调，选择最优进程数的最佳方法可能是试错。下表显示了不同 _numworkers 选择的平均吞吐量：
训练吞吐量作为数据加载工人数目函数（作者提供）
混合精度另一种流行的技术是使用较低精度的浮点数据类型，例如 torch.float16 或 torch.bfloat16，其中 torch.bfloat16 的动态范围通常被认为更适合机器学习训练。自然，降低数据类型精度可能会对收敛产生不利影响，因此应谨慎操作。PyTorch 提供了 torch.amp，这是一个自动混合精度包，用于优化这些数据类型的使用。Intel® AVX-512 包括对 bfloat16 数据类型的 支持。修改后的训练步骤如下所示：
for idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
with torch.amp.autocast(device_type='cpu', dtype=torch.bfloat16):
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()优化后的吞吐量为每秒24.34个样本，提高了86%!!
后通道内存格式通道在后的内存格式 是一个处于测试阶段的优化（截至撰写本文时），主要针对视觉模型，支持将四维（NCHW）张量在内存中存储为通道在后的方式。这样可以将每个像素的所有数据一起存储。这种优化主要针对视觉模型。被认为更“适合Intel平台”，这种内存格式据报道 在Intel® Xeon® 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(device_type='cpu', dtype=torch.bfloat16):
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()得到的吞吐量为每秒37.93个样本——相比基线实验又提高了56%，总提升达到了415%。我们太棒了！！
PyTorch 编译在之前的文章中，我们介绍了PyTorch对图编译的支持及其对运行时性能的潜在影响。与默认的即时执行模式（即，每次操作都独立运行）相反，compile API将模型转换为中间计算图，然后将其JIT编译为底层训练引擎最优化的低级机器代码。该API支持通过不同的后端库进行编译，并提供了多种配置选项。在这里，我们将评估范围限制在默认（TorchInductor）后端和来自Intel® Extension for PyTorch的ipex后端，这是一个针对Intel硬件进行了专门优化的库。请参阅文档以获取适当的安装和使用说明。更新后的模型定义如下所示：
import intel_extension_for_pytorch as ipex
model = torchvision.models.resnet50()
backend='inductor' # 可选地改为 'ipex'
model = torch.compile(model, backend=backend)在我们的示例模型中，当禁用“channels last”优化时，torch 编译的影响才会显现（每个后端大约增加 ~27%）。而当启用“channels last”优化时，性能实际上会下降。因此，我们在后续的实验中放弃了这一优化。
内存和线程优化有许多机会可以优化底层CPU资源的使用。这包括优化内存管理和线程分配到底层CPU硬件结构。内存管理可以通过使用高级内存分配器（例如Jemalloc和TCMalloc）和/或减少较慢的内存访问（即跨NUMA节点）来改进。线程分配可以通过适当配置OpenMP线程库和/或使用Intel的OpenMP库来改进。
一般来说，这类优化需要深入了解CPU架构及其支持的软件栈特性。为了简化操作，PyTorch 提供了 _torch.backends.xeon.runcpu 脚本，用于自动配置内存和线程库，以优化运行时性能。下面的命令将使用专用的内存和线程库。当我们讨论分布式训练选项时，我们将再次讨论 NUMA 节点的话题。
我们验证了 TCMalloc (conda install conda-forge::gperftools) 和 Intel 的 OpenMP 库 (pip install intel-openmp) 的正确安装，并运行以下命令。
python -m torch.backends.xeon.run_cpu train.py使用 _runcpu 脚本进一步提升了我们的运行时性能，达到了每秒39.05个样本。请注意，_runcpu 脚本包含了许多用于进一步调整性能的控制选项。务必查阅相关文档以充分利用其功能。
Intel® Extension for PyTorchIntel® Extension for 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上进行分布式训练Intel® 酷睿® 处理器设计时考虑了非一致内存访问（NUMA），在这种设计中，CPU 内存被划分为若干组，即 NUMA 节点，每个 CPU 核心都被分配到一个节点。虽然任何 CPU 核心都可以访问任何 NUMA 节点的内存，但访问其自身的节点（即本地内存）要快得多。这导致了在 NUMA 节点之间分配训练的概念，其中分配给每个 NUMA 节点的 CPU 核心作为一个单独的进程在分布式进程组中运行，并且节点之间的数据分布由Intel® oneCCL，Intel 的专用集体通信库进行管理。
我们可以轻松地使用 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
# 定义数据集和数据加载器
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
)
# 定义模型组件
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
)
# 配置DDP
model = torch.nn.parallel.DistributedDataParallel(model)
# 运行训练循环
# 销毁进程组
dist.destroy_process_group()
不幸的是，截至撰写本文时，Amazon EC2 c7i 实例系列不包括多NUMA实例类型。为了测试我们的分布式训练脚本，我们退回到一个 Amazon EC2 c6i.32xlarge 实例，该实例具有 64 个 vCPU 和 2 个 NUMA 节点。我们验证了 Intel® oneCCL Bindings for PyTorch 的 安装，并运行以下命令（如 这里 所述）：
source $(python -c "import oneccl_bindings_for_pytorch as torch_ccl;print(torch_ccl.cwd)")/env/setvars.sh
# 该示例命令将利用处理器的所有NUMA插槽，每个插槽作为一个rank。
ipexrun cpu --nnodes 1 --omp_runtime intel train.py 以下表格比较了在带有和不带分布式训练的c6i.32xlarge实例上的性能结果：
跨NUMA节点的分布式训练（由作者提供）
在我们的实验中，数据分布并没有提升运行时性能。请参见ipexrun 文档以获取更多性能调优选项。
使用Torch/XLA进行CPU训练在之前的帖子（例如，这里）中，我们讨论了 PyTorch/XLA 库及其使用 XLA 编译 来实现在 XLA 设备（如 TPU、GPU 和 CPU）上基于 PyTorch 的训练。类似于 torch 编译，XLA 使用图编译来生成针对目标设备优化的机器代码。随着 OpenXLA 项目 的建立，其中一个目标是支持所有硬件后端的高性能，包括 CPU（参见 CPU RFC 这里）。下面的代码块展示了调整我们原始（未优化）脚本所需的更改，以使用 PyTorch/XLA 进行训练：
import torch
import torchvision
import time
import 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支持逐渐成熟，这种情况会得到改善。
结果我们将在下面的表格中总结我们实验中的一部分结果。为了便于比较，我们在Amazon EC2 g5.2xlarge GPU实例上训练我们的模型，并按照这篇文章中讨论的优化步骤进行了操作。_每美元样本数_是根据Amazon EC2 按需定价页面计算得出的（截至撰写本文时，c7i.2xlarge每小时0.357美元，g5.2xlarge每小时1.212美元）。
性能优化结果（作者提供）
尽管我们在CPU实例上成功地将我们的玩具模型的训练性能提高了相当大的幅度（446%），但它仍然不如在GPU实例上的（优化后）性能。根据我们的结果，在GPU上进行训练大约会便宜6.7倍。很有可能通过进一步的性能调优和/或应用额外的优化策略，我们可以进一步缩小差距。再次强调，我们所达到的比较性能结果仅适用于此模型和运行时环境。
Amazon EC2 Spot 实例折扣与基于云的GPU实例类型相比，云中基于CPU的实例类型越来越多，这可能意味着以折扣价获得计算能力的机会更大，例如，通过使用Spot实例。Amazon EC2 Spot Instances是从多余的云服务容量中提供的实例，可以享受高达90%的折扣。作为交换，AWS保留随时在没有任何或几乎没有警告的情况下终止实例的权利。由于对GPU的需求很高，您可能会发现获取CPU Spot实例比GPU实例更容易。截至撰写本文时，c7i.2xlarge Spot Instance价格为$0.1291，这将使我们的每美元样本数提高到1135.76，并进一步缩小了优化后的GPU和CPU价格性能之间的差距（缩小到2.43倍）。
虽然我们优化的CPU训练的玩具模型（以及我们选择的环境）的运行时性能低于GPU的结果，但将相同的优化步骤应用于其他模型架构（例如，包含GPU不支持的组件的模型）可能会使CPU性能与GPU相匹配甚至超越GPU。即使在性能差距无法弥补的情况下，由于GPU计算能力不足，也可能有理由将某些机器学习工作负载运行在CPU上。
概要鉴于CPU的普及性，能够有效地利用它们进行训练和/或运行机器学习工作负载，对于提高开发生产力和最终产品的部署策略可能具有巨大的影响。虽然与GPU相比，CPU架构对于许多机器学习应用来说不太友好，但有许多工具和技术可以提升其性能——我们在本文中讨论并展示了其中的几种。
在这篇文章中，我们专注于在CPU上优化训练。请务必查看我们在 Medium 上发布的许多其他文章，这些文章涵盖了与机器学习工作负载的性能分析和优化相关的各种主题。
随时随地看视频
