PyTorch模型性能分析与优化——第2部分

如何使用 PyTorch Profiler 和 TensorBoard 识别和降低训练步骤中的 CPU 计算

Photo by Denise Chan on Unsplash

这是一系列有关在 GPU 上运行的 PyTorch 模型进行分析和优化的帖子的第二部分。在第一篇帖子中,我们演示了使用 PyTorch Profiler 和 TensorBoard 迭代地分析和优化 PyTorch 模型的过程和潜力。在这篇帖子中,我们将重点关注一种特定类型的性能问题,这种问题在 PyTorch 中特别普遍,因为它使用了急切执行(eager execution):对于模型执行的某些部分,依赖于 CPU。识别这些问题的存在和来源可能会非常困难,通常需要使用专用性能分析器。在本文中,我们将分享一些使用 PyTorch Profiler 和 PyTorch Profiler TensorBoard 插件识别此类性能问题的技巧。

急切执行(Eager Execution)的利与弊

PyTorch 的主要吸引力之一是其急切执行模式。在急切模式下,每个形成模型的 PyTorch 操作都会独立执行,直到达到。这与图模式相反,在此模式下,整个模型被预编译成一个单独的图,以最佳方式在 GPU 上运行并作为整体执行。通常,这种预编译结果会导致更好的性能(例如,请参见此处)。在急切模式下,编程上下文在每个操作后都会返回应用程序,因此允许我们访问和评估任意张量。这使得构建、分析和调试 ML 模型更加容易。另一方面,它也使我们的模型更容易受到(有时是意外的)插入次优代码块的影响。正如我们将演示的那样,了解如何识别和修复此类代码块可能会对您的模型速度产生重大影响。

玩具示例

在以下块中,我们介绍将用于演示的玩具示例。代码非常松散地基于我们先前帖子中的示例和在此 PyTorch 教程中定义的损失函数。

首先,我们定义了一个简单的分类模型。它的架构对本文不重要。

import torchimport torch.nn as nnimport torch.nn.functional as Fimport torch.optimimport torch.profilerimport torch.utils.dataimport torchvision.modelsimport torchvision.transforms as Tfrom torchvision.datasets.vision import VisionDatasetimport numpy as npfrom PIL import Image# sample modelclass Net(nn.Module):    def __init__(self):        super().__init__()        self.conv1 = nn.Conv2d(3, 8, 3, padding=1)        self.conv2 = nn.Conv2d(8, 12, 3, padding=1)        self.conv3 = nn.Conv2d(12, 16, 3, padding=1)        self.conv4 = nn.Conv2d(16, 20, 3, padding=1)        self.conv5 = nn.Conv2d(20, 24, 3, padding=1)        self.conv6 = nn.Conv2d(24, 28, 3, padding=1)        self.conv7 = nn.Conv2d(28, 32, 3, padding=1)        self.conv8 = nn.Conv2d(32, 10, 3, padding=1)        self.pool = nn.MaxPool2d(2, 2)    def forward(self, x):        x = self.pool(F.relu(self.conv1(x)))        x = self.pool(F.relu(self.conv2(x)))        x = self.pool(F.relu(self.conv3(x)))        x = self.pool(F.relu(self.conv4(x)))        x = self.pool(F.relu(self.conv5(x)))        x = self.pool(F.relu(self.conv6(x)))        x = self.pool(F.relu(self.conv7(x)))        x = self.pool(F.relu(self.conv8(x)))        x = torch.flatten(x, 1) # flatten all dimensions except batch        return x

接下来,我们定义一个相当标准的交叉熵损失函数。该损失函数将是我们讨论的主要焦点。

def log_softmax(x):    return x - x.exp().sum(-1).log().unsqueeze(-1)def weighted_nll(pred, target, weight):    assert target.max() < 10    nll = -pred[range(target.shape[0]), target]    nll = nll * weight[target]    nll = nll / weight[target].sum()    sum_nll = nll.sum()    return sum_nll# custom loss definitionclass CrossEntropyLoss(nn.Module):    def forward(self, input, target):        pred = log_softmax(input)        loss = weighted_nll(pred, target, torch.Tensor([0.1]*10).cuda())        return loss

最后,我们定义数据集和训练循环:

# dataset with random images that mimics the properties of CIFAR10class FakeCIFAR(VisionDataset):    def __init__(self, transform):        super().__init__(root=None, transform=transform)        self.data = np.random.randint(low=0,high=256,size=(10000,32,32,3),dtype=np.uint8)        self.targets = np.random.randint(low=0,high=10,size=(10000),dtype=np.uint8).tolist()    def __getitem__(self, index):        img, target = self.data[index], self.targets[index]        img = Image.fromarray(img)        if self.transform is not None:            img = self.transform(img)        return img, target    def __len__(self) -> int:        return len(self.data)transform = T.Compose(    [T.Resize(256),     T.PILToTensor()])train_set = FakeCIFAR(transform=transform)train_loader = torch.utils.data.DataLoader(train_set, batch_size=1024,                               shuffle=True, num_workers=8, pin_memory=True)device = torch.device("cuda:0")model = Net().cuda(device)criterion = CrossEntropyLoss().cuda(device)optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)model.train()# training loop wrapped with profiler objectwith torch.profiler.profile(        schedule=torch.profiler.schedule(wait=1, warmup=4, active=3, repeat=1),        on_trace_ready=torch.profiler.tensorboard_trace_handler(’./log/example’),        record_shapes=True,        profile_memory=True,        with_stack=True) as prof:    for step, data in enumerate(train_loader):        inputs = data[0].to(device=device, non_blocking=True)        labels = data[1].to(device=device, non_blocking=True)        inputs = (inputs.to(torch.float32) / 255. - 0.5) / 0.5        if step >= (1 + 4 + 3) * 1:            break        outputs = model(inputs)        loss = criterion(outputs, labels)        optimizer.zero_grad(set_to_none=True)        loss.backward()        optimizer.step()        prof.step()

一个有经验的PyTorch开发者可能已经注意到,我们的示例中包含了一些低效的代码行。同时,我们的实现中并没有明显的问题,这种低效性也很常见。如果您想测试自己的PyTorch熟练程度,请看看能否在阅读之前找到我们实现的交叉熵损失函数的三个问题。在接下来的章节中,我们将假设我们无法自己找到这些问题,并展示如何使用PyTorch Profiler及其相关的TensorBoard插件来识别它们。

与我们先前的帖子一样,我们将迭代地运行实验,识别性能问题,并尝试修复它们。我们将在一个Amazon EC2 g5.2xlarge实例上运行我们的实验(包含一个NVIDIA A10G GPU和8个vCPU),并使用官方的AWS PyTorch 2.0 Docker镜像。我们选择的训练环境有些随意,不应被视为其任何组件的认可。

初始性能结果

下图显示了上述脚本的性能报告的概述选项卡。

Performance Overview of Baseline Model (Captured by Author)

从图中可以看出,我们的GPU利用率相对较高,为92.04%,步骤时间为216毫秒。(与我们先前的帖子一样,torch-tb-profiler版本0.4.1中的概述将所有三个训练步骤的步骤时间相加。)仅从此报告中,您可能不认为我们的模型存在任何问题。但是,性能报告的Trace View告诉了一个完全不同的故事:

Trace View of Baseline Model (Captured by Author)

如上所述,我们的交叉熵损失函数的前向传递单独占用了216毫秒训练步骤中的211毫秒!这明显表明出了问题。相对于模型,我们的损失函数包含了很少的计算,肯定不应该占据98%的步骤时间。仔细查看调用栈,我们可以看到一些函数调用,包括“to”,“copy_”和“cudaStreamSynchronize”。这种组合通常表明正在将数据从CPU复制到GPU – 这不是我们想在损失计算中间发生的事情。在这种情况下,我们的性能问题也与GPU利用率的短暂下降相一致,如图所示。但是,这并不总是如此。 GPU利用率下降通常不会与性能问题对齐,或者根本看不到。

我们现在知道我们的损失函数存在性能问题,很可能与从主机复制张量到GPU有关。但是,这可能不足以确定导致问题的精确代码行。为了方便我们的搜索,我们将每行代码都包装在带标签的torch.profiler.record_function上下文管理器中,并重新运行性能分析。

# 自定义损失定义class CrossEntropyLoss(nn.Module):    def forward(self, input, target):        with torch.profiler.record_function('log_softmax'):            pred = log_softmax(input)        with torch.profiler.record_function('define_weights'):            weights = torch.Tensor([0.1]*10).cuda()        with torch.profiler.record_function('weighted_nll'):            loss = weighted_nll(pred, target, torch.Tensor([0.1]*10).cuda())        return loss

添加标签有助于我们确定权重定义,更准确地说是将权重复制到GPU中的有问题的代码行。

Performance Issue of Weights Definition as Seen in Trace View (Captured by Author)

优化 #1:从训练步骤中删除冗余的主机到GPU复制

一旦我们确定了第一个问题,修复它就相当简单了。在下面的代码块中,我们在损失init函数中仅将权重向量复制到GPU一次:

class CrossEntropyLoss(nn.Module):    def __init__(self):        super().__init__()        self.weight = torch.Tensor([0.1]*10).cuda()    def forward(self, input, target):        with torch.profiler.record_function('log_softmax'):            pred = log_softmax(input)        with torch.profiler.record_function('weighted_nll'):            loss = weighted_nll(pred, target, self.weight)        return loss

下面的图像显示了完成此修复后的性能分析结果:

Performance Overview Following Optimization #1 (Captured by Author)

令人失望的是,我们的第一个优化对步骤时间影响非常微小。如果我们查看Trace View报告,我们可以看到我们有一个新的严重性能问题需要解决。

Trace View Following Optimization #1 (Captured by Author)

我们的新报告指出了来自我们的weighted_nll函数的问题。与以前一样,我们使用了torch.profiler.record_function来确定有问题的代码行。在这种情况下,它是assert调用。

def weighted_nll(pred, target, weight):    with torch.profiler.record_function('assert'):        assert target.max() < 10    with torch.profiler.record_function('range'):        r = range(target.shape[0])    with torch.profiler.record_function('index'):        nll = -pred[r, target]    with torch.profiler.record_function('nll_calc'):        nll = nll * weight[target]        nll = nll/ weight[target].sum()        sum_nll = nll.sum()    return sum_nll

请注意,这个问题也存在于基础实验中,但是由于我们之前的性能问题而被隐藏了。在性能优化的过程中,经常会出现之前被其他问题隐藏的严重问题以这种方式突然出现,这并不罕见。

对调用栈的进一步分析显示了对“item”、“_local_scalar_dense”和“cudaMemcpyAsync”的调用。这通常表明正在从GPU复制数据到主机。确实,我们的assert调用在CPU上执行,需要访问驻留在GPU上的目标张量,从而引发高度低效的数据复制。

优化 #2:从训练步骤中删除冗余的GPU到主机复制

虽然验证输入标签的合法性可能是必要的,但应该以不对我们的训练性能产生太大负面影响的方式进行。在我们的情况下,修复问题只需将assert移动到数据输入管道中,在标签被复制到GPU之前执行。在assert被删除后,我们的性能仍然基本不变:

Performance Overview Following Optimization #2 (Captured by Author)

重要提示:尽管我们的目标通常是尝试减少在前向传递中主机和GPU之间的复制,但有时这是不可能的(例如,如果我们需要一个GPU不支持的内核)或不可取的(例如,如果在CPU上运行特定的内核将提高性能)。

分析Trace View介绍了我们接下来的性能问题:

Trace View Following Optimization #2 (Captured by Author)

我们再次看到,我们之前的优化揭示了一个新的严重性能问题,这次是在索引我们的pred张量时。索引由r和target张量定义。虽然目标张量已经驻留在GPU上,但是在上一行中定义的r张量没有。这又一次触发了低效的主机到GPU数据复制。

优化 #3:用torch.arange替换range

Python的range函数在CPU上输出一个列表。在训练步骤中存在任何列表都应该引起警惕。在下面的代码块中,我们用torch.arange替换了range的使用,并将其配置为直接在GPU上创建输出张量:

def weighted_nll(pred, target, weight):    with torch.profiler.record_function('range'):        r = torch.arange(target.shape[0], device="cuda:0")    with torch.profiler.record_function('index'):        nll = -pred[r, target]    with torch.profiler.record_function('nll_calc'):        nll = nll * weight[target]        nll = nll/ weight[target].sum()        sum_nll = nll.sum()    return sum_nll

这种优化的结果如下所示:

Performance Overview Following Optimization #3 (Captured by Author)

现在我们说得来了!!我们的步骤时间已经降至5.8毫秒,性能提高了惊人的3700%。

更新后的Trace View显示,损失函数已经降至一个非常合理的0.5毫秒。

Trace View Following Optimization #3 (Captured by Author)

但是还有改进的空间。让我们更仔细地看一下 weighted_nll 函数的 Trace View,它占据了大部分损失计算。

weighted_nll 函数的 Trace View(作者所拍摄)

从跟踪中可以看出,该函数由多个小块组成,每个小块最终映射到一个单独的 CUDA 内核,该内核通过 CudaLaunchKernel 调用加载到 GPU。理想情况下,我们希望尽可能减少 GPU 内核的总数,以减少 CPU 和 GPU 之间的交互量。一种方法是尽可能使用更高级的 PyTorch 运算符,例如 torch.nn.NLLLoss。这样的函数被假定“融合”了底层操作,因此需要更少的总内核数。

优化 #4:将自定义 NLL 替换为 torch.nn.NLLLoss

下面的代码块包含我们更新后的损失定义,现在使用 torch.nn.NLLLoss。

class CrossEntropyLoss(nn.Module):    def __init__(self):        super().__init__()        self.weight = torch.Tensor([0.1]*10).cuda()    def forward(self, input, target):        pred = log_softmax(input)        nll = torch.nn.NLLLoss(self.weight)        loss = nll(pred, target)        return loss

在这里,我们已经引入了另一个常见的错误,我们将继续演示。

使用更高级的函数进一步将我们的步骤时间降至 5.3 毫秒(从 5.8 毫秒)。

优化 #4 后的性能概述(作者所拍摄)

然而,如果我们更仔细地查看 Trace View,我们可以看到损失函数的大部分时间现在都花在初始化 torch.nn.NLLLoss 对象上了!

优化 #4 后的 Trace View(作者所拍摄)

回顾我们的损失函数,我们可以看到在训练步骤的每次迭代中都初始化了一个新的 NLLLoss 对象。自然而然,对象初始化发生在 CPU 上,虽然(在我们的情况下)它相对较快,但它是我们希望在训练步骤中避免执行的操作。

优化 #5:避免在训练步骤中初始化对象

在下面的代码块中,我们修改了我们的损失实现,以便在 init 函数中创建 torch.nn.NLLLoss 的单个实例。

class CrossEntropyLoss(nn.Module):    def __init__(self):        super().__init__()        self.weight = torch.Tensor([0.1]*10).cuda()        self.nll = torch.nn.NLLLoss(self.weight)     def forward(self, input, target):        pred = log_softmax(input)        loss = self.nll(pred, target)        return loss

结果显示,步骤时间进一步提高,现在为 5.2 毫秒。

优化 #6:使用 torch.nn.CrossEntropyLoss 替代自定义损失

PyTorch 包括内置的 torch.nn.CrossEntropyLoss,我们现在评估并将其与我们的自定义损失实现进行比较。

criterion = torch.nn.CrossEntropyLoss().cuda(device)

结果步骤时间创下了新低,仅为5毫秒,整体性能提升了4200%(与216毫秒相比)。

损失计算的前向传递性能改善更加显著:从最初的211毫秒,降至仅79秒(!!),如下图所示:

优化 #7:编译损失函数

最后一次优化尝试中,我们将配置损失函数以使用torch.compile API在图形模式下运行。正如我们在本文中详细讨论并在本文前篇文章中证明的那样,torch.compile将使用诸如内核融合和乱序执行等技术将损失函数映射到底层训练加速器中的低级计算内核,从而优化训练过程。

criterion = torch.compile(torch.nn.CrossEntropyLoss().cuda(device))

下图显示了此实验的跟踪视图结果。

我们首先可以看到包含“OptimizedModule”和“dynamo”术语的出现,这表明使用了torch.compile。我们还可以看到,实际上,模型编译并没有减少损失函数加载的内核数量,这意味着它没有识别任何额外的内核融合机会。事实上,在我们的情况下,损失编译实际上使损失函数的前向传递时间从79微秒增加到154微秒。似乎CrossEntropyLoss没有足够的重量来受益于此优化。

您可能会想知道为什么我们不能将torch编译应用于最初的损失函数,并依赖它以最优方式编译我们的代码。这可以节省我们所述的逐步优化的所有麻烦。这种方法的问题在于,尽管PyTorch 2.0编译(截至本文撰写时)确实优化了某些类型的GPU到CPU交叉,但某些类型会导致图形编译崩溃,而其他类型则会导致创建多个小图而不是单个大图。最后一类会导致图形中断,从而限制了torch.compile功能提高性能的能力。(解决此问题的一种方法是调用torch.compile,并将fullgraph标志设置为True。)有关使用此选项的更多详细信息,请参见我们以前的文章。

结果

在下表中,我们总结了我们运行的实验的结果:

Optimization Experiments Results (By Author)

我们的连续优化带来了惊人的4143%性能提升!!请记住,我们最初使用的是一个看起来非常单纯的损失函数。如果没有对我们应用程序的行为进行深入分析,我们可能永远不会知道有什么问题,而会继续支付比所需支付的高出41倍(!!)的费用。

您可能已经注意到,在我们的最终试验中,GPU利用率显著下降。这表明进一步性能优化的重大潜力。尽管我们的演示已经接近尾声,但我们的工作还没有完成。请参见我们以前的文章,获取关于如何从这里继续的一些想法。

结论

让我们总结一些我们所学到的内容。我们将总结分为两部分。首先,我们描述了可能影响训练性能的一些编码习惯。其次,我们推荐一些性能分析的技巧。请注意,这些结论基于我们在本文中分享的示例,可能不适用于您自己的用例。机器学习模型在性质和行为上差异很大。因此,强烈建议您根据您项目的详细信息评估这些结论。

编程技巧

您实现模型前向传递的方式会对其性能产生重大影响。在此,我们列出了一些建议,基于本文中所涉及的示例。

  1. 避免在前向传递中初始化常量张量。将其放在构造函数中。
  2. 避免在前向传递中对存储在 GPU 上的张量使用断言。将它们移到数据输入管道中和/或检查 PyTorch 是否有任何内置方法来执行所需的数据验证。
  3. 避免使用列表。请检查是否使用 torch.arange 在设备上直接创建张量可以是更好的选择。
  4. 使用 PyTorch 运算符,如 torch.nn.NLLLoss 和 torch.nn.CrossEntropyLoss,而不是创建自己的损失实现。
  5. 避免在前向传递中初始化对象。将其放在构造函数中。
  6. 在适当的情况下,请考虑使用 torch.compile。

性能分析技巧

正如我们所演示的,Tensorboard PyTorch Profiler 插件的 Trace 视图对于识别模型中的性能问题至关重要。以下是我们示例的一些主要收获:

  1. 高 GPU 利用率不一定意味着您的代码运行得最佳。
  2. 注意代码执行时间比预期长的部分。
  3. 使用 torch.profiler.record_function 来确定性能问题。
  4. GPU 利用率下降并不一定与性能问题的源头对齐。
  5. 注意来自主机到 GPU 的意外数据复制。这些通常通过在 Trace 视图中搜索“to”、“copy_”和“cudaStreamSynchronize”来识别。
  6. 注意来自 GPU 到主机的意外数据复制。这些通常通过在 Trace 视图中搜索“item”和“cudaStreamSynchronize”来识别。

总结

在本文中,我们重点关注了训练应用程序中性能问题,这些问题是由训练步骤的前向传递中 CPU 和 GPU 之间的冗余交互引起的。我们演示了如何使用 PyTorch Profiler 及其关联的 TensorBoard 插件识别此类问题并实现显著的性能提升。

与我们之前的文章一样,我们强调成功优化的路径将根据训练项目的细节(包括模型架构和训练环境)而大不相同。在实践中,达成您的目标可能比我们在此处提供的示例更加困难。我们所描述的一些技术可能对您的性能影响很小,甚至可能使其更糟。我们还注意到,我们选择的精确优化和应用它们的顺序在某种程度上是任意的。强烈鼓励您根据项目的具体细节开发自己的工具和技术,以达到您的优化目标。