加速 PyTorch 模型:使用 ROCm 在 AMD GPU 上应用 torch.compile

2024-09-01 08:28

本文主要是介绍加速 PyTorch 模型:使用 ROCm 在 AMD GPU 上应用 torch.compile,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

Accelerate PyTorch Models using torch.compile on AMD GPUs with ROCm — ROCm Blogs

介绍

PyTorch 2.0 引入了一个名为*torch.compile()*的工具,可以极大地加速 PyTorch 代码和模型。通过将 PyTorch 代码转换为高度优化的内核,`torch.compile` 在现有代码库上进行最小化修改即可提供显著的性能提升。此功能允许精确优化单个函数、整个模块以及复杂的训练循环,提供了一个多功能且强大的工具来提高计算效率。

在这篇博客中,我们将演示如何在 AMD GPU 上使用 ROCm 和 torch.compile 加速各种实际模型。

torch.compile 的工作原理

torch.compile 的执行涉及几个关键步骤:
1. 图获取:模型被分解并重写为子图。可以编译或优化的子图被扁平化。不能编译的子图会回退到eager mode(即时模式)。
2. 图降低:所有的 PyTorch 操作都会被分解成其选定的特定后端的内核。
3. 图编译:所有后端内核调用其对应的低级设备操作。torch.compile 的四项关键技术是:TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor。这些组件各自承担着使 torch.compile 功能得以实现的重要角色。

- TorchDynamo:可靠且快速地获取图。它通过符号解释 Python 字节码,将其转化为张量操作图。如果遇到无法解释的代码段,它会默认为常规的 Python 解释器。该方法确保可以处理多种程序,同时大幅提升性能。
- AOT Autograd:重新使用 Autograd 进行提前(AoT)计算图。AOT Autograd 是 PyTorch 2.0 的自动微分引擎。其功能是提前生成反向传递的跟踪,提升微分过程的效率。AOT Autograd 使用 PyTorch 的 torch_dispatch 机制来遍历现有的 PyTorch autograd 引擎,提前捕获反向传递。这使得前向传递和反向传递都能加速。
- PrimTorch:提供稳定的基础操作符。它将复杂的 PyTorch 操作分解为更简单的操作。
- TorchInductor:为加速器和后端生成高速代码。TorchInductor 是一个深度学习编译器,将中间表示转化为可执行代码。它获取 TorchDynamo 生成的计算图并将其转化为优化的低级内核。对于 NVIDIA 和 AMD GPU,它使用OpenAI Triton 作为基础组件。

torch.compile 函数具有多种编译模式,例如 default、`reduce-overhead` 和 max-autotune,它们在编译时间和推理开销上有所不同。通常,`max-autotune` 模式相对于 reduce-overhead 模式编译时间更长,但推理速度更快。`default` 模式编译最快,但相对于 reduce-overhead 模式推理效率较低。`torch.compile` 函数在第一次执行期间将模型编译为优化内核。因此,初次运行可能会因为编译时间而稍长,但随后的执行由于减少了 Python 开销和 GPU 读写操作展示了加速效果。最终的加速效果可能因模型架构和批处理大小而异。您可以在PyTorch 2.0 简介介绍和教程中了解更多关于 PyTorch 编译过程的内容。

在这篇博客中,我们通过评估以下模型在Eager-mode和不同torch.compile模式下的性能,展示了使用 torch.compile 可以在 AMD GPU 上加速实际模型:
- 使用卷积神经网络(ResNet-152)模型进行图像分类
- 使用视觉变压器模型进行图像分类
- 使用 Llama 2 7B 模型进行文本生成

在这篇博客中使用的完整代码可以在 [ROCm blogs repository](rocm-blogs/blogs/artificial-intelligence/torch_compile at release · ROCm/rocm-blogs · GitHub) 中找到。

前提条件

这篇博客是在以下环境中测试的。有关设置的详细支持信息,请参阅 [ROCm 文档](ROCm installation for Linux — ROCm installation (Linux))。- 硬件和操作系统:
  - [AMD Instinct GPU](AMD Instinct™ Accelerators)
  - Ubuntu 22.04.3 LTS- 软件:
  - [ROCm 6.0+](Quick start installation guide — ROCm installation (Linux))
  - [ROCm 2.0+ 版的 PyTorch](Installing PyTorch for ROCm — ROCm installation (Linux))- 库:
  - transformerssentencepiecenumpytabulatescipymatplotlib在这篇博客中,我们使用 Linux 设备上安装了 MI210 加速器的 [rocm/pytorch-nightly](https://hub.docker.com/r/rocm/pytorch-nightly/tags) Docker 镜像。建议使用 PyTorch 的 nightly 版本以实现更优化的加速效果。

安装依赖项

!pip install -q transformers==4.31 sentencepiece numpy tabulate scipy matplotlib sentencepiece huggingface_hub

检查 AMD GPU 和 PyTorch 版本(>2.0)。

import torch
print(f"number of GPUs: {torch.cuda.device_count()}")
print([torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])torch_ver = [int(x) for x in torch.__version__.split(".")[:2]]
assert torch_ver >= [2, 0], "Requires PyTorch >= 2.0"
print("PyTorch Version:", torch.__version__)

输出:

    number of GPUs: 1['AMD Instinct MI210']PyTorch Version: 2.4.0a0+git1f8177d

接下来,我们将定义一个辅助函数来测量给定函数的执行时间。

import time
def timed(fn, n_test: int, dtype: torch.dtype) -> tuple:"""Measure the execution time for a given function.Args:- fn (function): The function to be timed.- n_test (int): Number of times the function is executed to get the average time.- dtype (torch.dtype): Data type for PyTorch tensors.Returns:- tuple: A tuple containing the average execution time (in milliseconds) and the output of the function."""with torch.no_grad(), torch.autocast(device_type='cuda', dtype=dtype):           dt_loop_sum = []for _ in range(n_test):torch.cuda.synchronize()start = time.time()output = fn()torch.cuda.synchronize()end = time.time()dt_loop_sum.append(end-start)dt_test = sum(dt_loop_sum) / len(dt_loop_sum) return dt_test * 1000, output

通过 TorchDynamo 使用 torch.compile 需要一个将捕获的图转换为快速机器代码的后端。不同的后端可以带来不同的优化增益。您可以通过运行 torch.compiler.list_backends() 查看 TorchDynamo 支持的后端列表。

torch.compiler.list_backends() 

输出:

    ['cudagraphs', 'inductor', 'onnxrt', 'openxla', 'openxla_eval', 'tvm']

在这篇博客中,我们选择 inductor 作为后端,这是默认设置。这个后端将允许我们从原生 PyTorch 应用程序的操作中动态生成 Triton 内核。

使用 torch.compile 加速 ResNet-152

ResNet 是一种卷积神经网络,最初在论文 Deep Residual Learning for Image Recognition(He 等)中提出。在本次评估中,我们使用 ResNet-152 作为图像分类模型的骨干网络。我们在不同模式下测试并比较推理时间,包括 Eager 模式、`default`、`reduce-overhead` 和 max-autotune 模式。

验证模型和环境设置

首先,我们下载并显示用作分类模型输入的测试图像。

# Download an example image from the pytorch website
import urllib
import matplotlib.pyplot as plt
url, filename = ("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
try: urllib.URLopener().retrieve(url, filename)
except: urllib.request.urlretrieve(url, filename)from PIL import Image
input_image = Image.open(filename)
plt.imshow(input_image)
plt.axis('off')
plt.show()

导入图像预处理器和模型来处理上述图像。

import torch
import torchvision.transforms as transforms# create the image preprocessor
preprocess = transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model# load the resnet152 model
model = torch.hub.load('pytorch/vision:v0.17.2', 'resnet152', pretrained=True)
model.eval()# move the input and model to GPU for speed if available
if torch.cuda.is_available():input_batch = input_batch.to('cuda')model.to('cuda')
with torch.no_grad():output = model(input_batch)# Tensor of shape 1000, with confidence scores over ImageNet's 1000 classes
print(output.shape)

输出:

    torch.Size([1, 1000])

打印基于概率的 topk 标签的辅助函数。

def print_topk_labels(output, k):# The output has unnormalized scores. To get probabilities, you can run a softmax on it.probabilities = torch.nn.functional.softmax(output[0], dim=0)# Read the categorieswith open("imagenet_classes.txt", "r") as f:categories = [s.strip() for s in f.readlines()]# Show top categories per imagetopk_prob, topk_catid = torch.topk(probabilities, k)for i in range(topk_prob.size(0)):print(categories[topk_catid[i]], topk_prob[i].item())
# Download ImageNet labels
!wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt

显示前 5 个标签及其概率。

print_topk_labels(output, 5)

输出:

    Samoyed 0.7907489538192749Pomeranian 0.08977615833282471white wolf 0.03610273823142052keeshond 0.02681431733071804Arctic fox 0.022788070142269135

我们可以发现,模型效果很好。这表明环境是正确的,我们已经准备好使用 torch.compile 测试基于 ResNet-152 的模型。

ResNet-152 模型在 Eager Mode 下的性能评估

为了预热 GPU,我们在进行 20 次额外迭代以获取模型的平均推理时间之前,先运行 ResNet-152 模型 10 次。

n_warmup = 10
n_test = 20
dtype = torch.bfloat16
inference_time=[]
mode=[]
t_warmup, _ = timed(lambda:model(input_batch), n_warmup, dtype)
t_test, output = timed(lambda:model(input_batch), n_test, dtype)
print(f"Average inference time for resnet152(warmup): dt_test={t_warmup} ms")
print(f"Average inference time for resnet152(test): dt_test={t_test} ms")
print_topk_labels(output, 5)
inference_time.append(t_test)
mode.append("eager")

输出:

    Average inference time for resnet152(warmup): dt_test=164.6312952041626 msAverage inference time for resnet152(test): dt_test=18.761909008026123 msSamoyed 0.80078125Pomeranian 0.0791015625white wolf 0.037353515625keeshond 0.0257568359375Arctic fox 0.022705078125

ResNet-152 模型在 torch.compile(default) Mode 下的性能评估

要将 torch.compile 应用于 ResNet-152,我们可以按照下面的代码进行包装。
• mode:我们使用 default 编译模式,这是性能和开销之间的良好平衡。
• fullgraph:如果为 True,`torch.compile()` 要求整个函数都能捕获到一个单一的图中。如果这不可能,则会引发错误。

#clean up the workspace with torch._dynamo.reset().
torch._dynamo.reset()
model_opt1 = torch.compile(model, fullgraph=True)
t_compilation, _ = timed(lambda:model_opt1(input_batch), 1, dtype)
t_warmup, _ = timed(lambda:model_opt1(input_batch), n_warmup, dtype)
t_test, output = timed(lambda:model_opt1(input_batch), n_test, dtype)
print(f"Compilation time: dt_compilation={t_compilation} ms")
print(f"Average inference time for compiled resnet152(warmup): dt_test={t_warmup} ms")
print(f"Average inference time for compiled resnet152(test): dt_test={t_test} ms")
print_topk_labels(output, 5)
inference_time.append(t_test)
mode.append("default")

输出:

    Compilation time: dt_compilation=24626.18637084961 msAverage inference time for compiled resnet152(warmup): dt_test=15.319490432739258 msAverage inference time for compiled resnet152(test): dt_test=15.275216102600098 msSamoyed 0.80078125Pomeranian 0.0791015625white wolf 0.037353515625keeshond 0.0257568359375Arctic fox 0.022705078125

ResNet-152 模型在 torch.compile(reduce-overhead) Mode 下的性能评估

reduce-overhead 模式利用 CUDA graphs 来减少内核启动的开销,改善整体延迟。如果你想了解更多,可以关于CUDA graphs 的内容。

torch._dynamo.reset()
model_opt2 = torch.compile(model, mode="reduce-overhead", fullgraph=True)
t_compilation, _ = timed(lambda:model_opt2(input_batch), 1, dtype)
t_warmup, _ = timed(lambda:model_opt2(input_batch), n_warmup, dtype)
t_test, output = timed(lambda:model_opt2(input_batch), n_test, dtype)
print(f"Compilation time: dt_compilation={t_compilation} ms")
print(f"Average inference time for compiled resnet152(warmup): dt_test={t_warmup} ms")
print(f"Average inference time for compiled resnet152(test): dt_test={t_test} ms")
print_topk_labels(output, 5)
inference_time.append(t_test)
mode.append("reduce-overhead")

输出:

    Compilation time: dt_compilation=18916.11909866333 msAverage inference time for compiled resnet152(warmup): dt_test=39.9461030960083 msAverage inference time for compiled resnet152(test): dt_test=5.042397975921631 msSamoyed 0.80078125Pomeranian 0.0791015625white wolf 0.037353515625keeshond 0.0257568359375Arctic fox 0.022705078125

ResNet-152模型在torch.compile(max-autotune)模式下的性能评估

max-autotune模式利用基于Triton的矩阵乘法和卷积运算。它默认启用CUDA图。

torch._dynamo.reset()
model_opt3 = torch.compile(model, mode="max-autotune", fullgraph=True)
t_compilation, _ = timed(lambda:model_opt3(input_batch), 1, dtype)
t_warmup, _ = timed(lambda:model_opt3(input_batch), n_warmup, dtype)
t_test, output = timed(lambda:model_opt3(input_batch), n_test, dtype)
print(f"Compilation time: dt_compilation={t_compilation} ms")
print(f"Average inference time for compiled resnet152(warmup): dt_test={t_warmup} ms")
print(f"Average inference time for compiled resnet152(test): dt_test={t_test} ms")
print_topk_labels(output, 5)
inference_time.append(t_test)
mode.append("max-autotune")

输出:

    AUTOTUNE convolution(1x64x56x56, 256x64x1x1)triton_convolution_49 0.0238 ms 100.0%triton_convolution_48 0.0240 ms 99.3%convolution 0.0242 ms 98.7%triton_convolution_46 0.0325 ms 73.4%triton_convolution_52 0.0326 ms 73.0%triton_convolution_53 0.0331 ms 72.0%triton_convolution_47 0.0333 ms 71.6%triton_convolution_50 0.0334 ms 71.3%triton_convolution_51 0.0341 ms 70.0%triton_convolution_42 0.0360 ms 66.2%SingleProcess AUTOTUNE takes 64.3134 seconds...AUTOTUNE convolution(1x256x14x14, 1024x256x1x1)triton_convolution_538 0.0285 ms 100.0%triton_convolution_539 0.0290 ms 98.3%convolution 0.0299 ms 95.2%triton_convolution_536 0.0398 ms 71.5%triton_convolution_542 0.0400 ms 71.2%triton_convolution_543 0.0406 ms 70.1%triton_convolution_537 0.0411 ms 69.3%triton_convolution_540 0.0443 ms 64.3%triton_convolution_541 0.0464 ms 61.4%triton_convolution_532 0.0494 ms 57.6%SingleProcess AUTOTUNE takes 15.0623 seconds...AUTOTUNE addmm(1x1000, 1x2048, 2048x1000)bias_addmm 0.0240 ms 100.0%addmm 0.0240 ms 100.0%triton_mm_2176 0.0669 ms 35.9%triton_mm_2177 0.0669 ms 35.9%triton_mm_2174 0.0789 ms 30.4%triton_mm_2175 0.0789 ms 30.4%triton_mm_2180 0.0878 ms 27.3%SingleProcess AUTOTUNE takes 8.4102 secondsCompilation time: dt_compilation=820945.9936618805 msAverage inference time for compiled resnet152(warmup): dt_test=41.12842082977295 msAverage inference time for compiled resnet152(test): dt_test=5.32916784286499 msSamoyed 0.796875Pomeranian 0.083984375white wolf 0.037353515625keeshond 0.025634765625Arctic fox 0.0225830078125

基于输出,我们可以看到Triton正在自主优化矩阵乘法和卷积操作。相比于其他模式,这个过程需要极长的时间。你可以在这里将 编译时间 与之前测试的模式进行比较。

虽然使用了Triton调优,但在这种情况下,`max-autotune`模式并没有显著增强性能,和 reduce-overhead模式相比没有明显优势。这表明在我们的测试平台上,ResNet-152的瓶颈并不主要在于矩阵乘法或卷积操作。要进一步提高性能并应用高级设置,请参考torch._inductor.config。

比较从上述四种模式获得的推理时间

import matplotlib.pyplot as plt# Plotting the bar graph
plt.bar(mode, inference_time)
print(inference_time)
print(mode)# Adding labels and title
plt.xlabel('mode')
plt.ylabel('Inference time (ms)')
plt.title('ResNet-152')# Displaying the plot
plt.show()

输出:

    [18.761909008026123, 15.275216102600098, 5.042397975921631, 5.32916784286499]['eager', 'default', 'reduce-overhead', 'max-autotune']

从图表中可以看到,`torch.compile`显著提升了ResNet-152在AMD MI210与ROCm上的性能,达到了*3.5*倍以上的提升。

使用 torch.compile 加速 Vision Transformer

Vision Transformer(ViT)是一个类似 BERT 的 transformer 编码器模型,在大规模的图像集合上以有监督方式进行了预训练,具体来说是在分辨率为 224×224 像素的 ImageNet-21k 数据集上预训练的。以下是如何使用这个模型将 COCO 2017 数据集中的一张图像分类为 1,000 个 ImageNet 类别之一的示例,使用的是vit-base-patch16-224检查点。

from transformers import ViTImageProcessor, ViTForImageClassification
from PIL import Image
import requests
import matplotlib.pyplot as plturl = 'http://images.cocodataset.org/val2017/000000039769.jpg'
image = Image.open(requests.get(url, stream=True).raw)
plt.imshow(image)
plt.axis('off')  # Turn off axis
plt.show()# load the image processor and model
processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')inputs = processor(images=image, return_tensors="pt")if torch.cuda.is_available():inputs = inputs.to('cuda')model.to('cuda')outputs = model(**inputs)
logits = outputs.logits
# model predicts one of the 1000 ImageNet classes
predicted_class_idx = logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])

输出:

    Predicted class: Egyptian cat

模型和环境看起来都很好。接下来,我们将按照与 ResNet-152 相同的测试流程进行测试,包括在不同模式下测试模型,并在最后评估性能。在每种模式下,我们将进行 10 次迭代以进行预热,然后进行额外的 20 次迭代,以获得模型的平均推理时间。

n_warmup = 10
n_test = 20
dtype = torch.bfloat16
inference_time=[]
mode=[]

评估 Vision Transformer 模型在 Eager 模式下的性能

torch._dynamo.reset()
t_warmup, _ = timed(lambda:model(**inputs), n_warmup, dtype)
t_test, output = timed(lambda:model(**inputs), n_test, dtype)
print(f"Average inference time for ViT(warmup): dt_test={t_warmup} ms")
print(f"Average inference time for ViT(test): dt_test={t_test} ms")
inference_time.append(t_test)
mode.append("eager")
# model predicts one of the 1000 ImageNet classes
predicted_class_idx = output.logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])

输出:

    Average inference time for ViT(warmup): dt_test=8.17105770111084 msAverage inference time for ViT(test): dt_test=7.561385631561279 msPredicted class: Egyptian cat

评估 Vision Transformer 模型在 torch.compile(default) 模式下的性能

torch._dynamo.reset()
model_opt1 = torch.compile(model, fullgraph=True)
t_compilation, _ = timed(lambda:model_opt1(**inputs), 1, dtype)
t_warmup, _ = timed(lambda:model_opt1(**inputs), n_warmup, dtype)
t_test, output = timed(lambda:model_opt1(**inputs), n_test, dtype)
print(f"Compilation time: dt_compilation={t_compilation} ms")
print(f"Average inference time for ViT(warmup): dt_test={t_warmup} ms")
print(f"Average inference time for ViT(test): dt_test={t_test} ms")
inference_time.append(t_test)
mode.append("default")
# model predicts one of the 1000 ImageNet classes
predicted_class_idx = output.logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])

输出:

    Compilation time: dt_compilation=13211.912631988525 msAverage inference time for ViT(warmup): dt_test=7.065939903259277 msAverage inference time for ViT(test): dt_test=7.033288478851318 msPredicted class: Egyptian cat

评估 Vision Transformer 模型在 torch.compile(reduce-overhead) 模式下的性能

torch._dynamo.reset()
model_opt2 = torch.compile(model, mode="reduce-overhead", fullgraph=True)
t_compilation, _ = timed(lambda:model_opt2(**inputs), 1, dtype)
t_warmup, _ = timed(lambda:model_opt2(**inputs), n_warmup, dtype)
t_test, output = timed(lambda:model_opt2(**inputs), n_test, dtype)
print(f"Compilation time: dt_compilation={t_compilation} ms")
print(f"Average inference time for ViT(warmup): dt_test={t_warmup} ms")
print(f"Average inference time for ViT(test): dt_test={t_test} ms")
inference_time.append(t_test)
mode.append("reduce-overhead")
# model predicts one of the 1000 ImageNet classes
predicted_class_idx = output.logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])

输出:

    Compilation time: dt_compilation=10051.868438720703 msAverage inference time for ViT(warmup): dt_test=30.241727828979492 msAverage inference time for ViT(test): dt_test=3.2375097274780273 msPredicted class: Egyptian cat

评估 Vision Transformer 模型在 torch.compile(max-autotune) 模式下的性能

torch._dynamo.reset()
model_opt3 = torch.compile(model, mode="max-autotune", fullgraph=True)
t_compilation, _ = timed(lambda:model_opt3(**inputs), 1, dtype)
t_warmup, _ = timed(lambda:model_opt3(**inputs), n_warmup, dtype)
t_test, output = timed(lambda:model_opt3(**inputs), n_test, dtype)
print(f"Compilation time: dt_compilation={t_compilation} ms")
print(f"Average inference time for ViT(warmup): dt_test={t_warmup} ms")
print(f"Average inference time for ViT(test): dt_test={t_test} ms")
inference_time.append(t_test)
mode.append("max-autotune")
# model predicts one of the 1000 ImageNet classes
predicted_class_idx = output.logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])

输出:

    AUTOTUNE convolution(1x3x224x224, 768x3x16x16)convolution 0.0995 ms 100.0%triton_convolution_2191 0.2939 ms 33.9%triton_convolution_2190 0.3046 ms 32.7%triton_convolution_2194 0.3840 ms 25.9%triton_convolution_2195 0.4038 ms 24.6%triton_convolution_2188 0.4170 ms 23.9%...AUTOTUNE addmm(197x768, 197x768, 768x768)bias_addmm 0.0278 ms 100.0%addmm 0.0278 ms 100.0%triton_mm_2213 0.0363 ms 76.7%triton_mm_2212 0.0392 ms 71.0%triton_mm_2207 0.0438 ms 63.5%triton_mm_2209 0.0450 ms 61.9%triton_mm_2206 0.0478 ms 58.2%triton_mm_2197 0.0514 ms 54.2%triton_mm_2208 0.0533 ms 52.3%triton_mm_2196 0.0538 ms 51.8%...AUTOTUNE addmm(1x1000, 1x768, 768x1000)bias_addmm 0.0229 ms 100.0%addmm 0.0229 ms 100.0%triton_mm_4268 0.0338 ms 67.8%triton_mm_4269 0.0338 ms 67.8%triton_mm_4266 0.0382 ms 59.8%triton_mm_4267 0.0382 ms 59.8%triton_mm_4272 0.0413 ms 55.4%triton_mm_4273 0.0413 ms 55.4%triton_mm_4260 0.0466 ms 49.1%triton_mm_4261 0.0466 ms 49.1%SingleProcess AUTOTUNE takes 8.9279 secondsCompilation time: dt_compilation=103891.38770103455 msAverage inference time for ViT(warmup): dt_test=31.742525100708004 msAverage inference time for ViT(test): dt_test=3.2366156578063965 msPredicted class: Egyptian cat

比较在上述四种模式下获得的 ViT 推理时间

# Plotting the bar graph
plt.bar(mode, inference_time)
print(inference_time)
print(mode)# Adding labels and title
plt.xlabel('mode')
plt.ylabel('Inference time (ms)')
plt.title('ViT')# Displaying the plot
plt.show()

输出:

    [7.561385631561279, 7.033288478851318, 3.2375097274780273, 3.2366156578063965]['eager', 'default', 'reduce-overhead', 'max-autotune']

从图表中可以看出,`torch.compile` 显著提升了 ViT 的性能,在 AMD MI210 上通过 ROCm 提升了超过 2.3 倍。 

加速Llama 2 7B模型与torch.compile

Llama 2是一个大型语言模型,由一系列能够响应提示生成文本和代码的模型组成。PyTorch团队在PyTorch Labs GitHub仓库中提供了一个简单、高效的PyTorch原生实现的Transformer文本生成模型。在我们的评估中,我们简化了代码,仅用于应用`torch.compile`进行优化。具体代码可以在src文件夹中找到。

与之前的评估相比,我们这次重点评估的是Llama 2 7B模型的吞吐量(批次大小=1)。我们将进行20次迭代以得出模型的平均吞吐量。

下载openlm-research/open_llama_7b模型并转换为PyTorch格式。

gpt-fast文件夹可以在src文件夹中找到。

%%bash
pip install sentencepiece huggingface_hubcd gpt-fast
./scripts/prepare.sh openlm-research/open_llama_7b

输出:

    Model config {'block_size': 2048, 'vocab_size': 32000, 'n_layer': 32, 'n_head': 32, 'dim': 4096, 'intermediate_size': 11008, 'n_local_heads': 32, 'head_dim': 128, 'rope_base': 10000, 'norm_eps': 1e-05}Saving checkpoint to checkpoints/openlm-research/open_llama_7b/model.pth

在Eager模式下对Llama 2 7B模型进行性能评估

指定`--compile none`以使用Eager模式。
• --compile:设置为`none`以使用Eager模式
• --profile:启用torch.profiler的跟踪功能
• --checkpoint_path:检查点路径
• --prompt:输入提示
• --max_new_tokens:最大新的token数
• --num_samples:样本数量

%%bash
cd gpt-fast
python generate_simp.py --compile none --profile ./trace_compile_none --checkpoint_path checkpoints/openlm-research/open_llama_7b/model.pth --prompt "def quicksort(arr):"  --max_new_tokens 200  --num_samples 20

输出:

    __CUDNN VERSION: 3000000__Number CUDA Devices: 1Using device=cudaLoading model ...Time to load model: 3.94 secondsCompilation time: 11.18 secondsdef quicksort(arr):"""Quickly sorts a list."""arr = arr.sort()return arrdef fizzbuzz():"""Does the fizzbuzz algorithm."""return 'fizzbuzz'def reverse_string():"""Reverses a string."""return 'foobar'[::-1]if __name__ == "__main__":print(quicksort([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]))print(fizzbuzz())print(reverse_string())Vuetify, MUI, BEM, CSS, JavaScriptCSS / JavaScript / Vue##### vue2vuetifyThe vue2vuetify package contains a declarativeAverage tokens/sec: 28.61Memory used: 13.62 GB

输出结果将包括三个部分:
• 系统信息和编译时间,
• 基于给定提示的模型输出,以及
• 在模型执行期间收集的指标。

根据输出结果,我们观察到推理速度大约为每秒28个token,这并不算太差。需要注意的是,对于`def quicksort(arr):的响应质量可能并不令人满意。但在本博客中这是可以接受的,因为我们的重点是使用torch.compile`来提高推理吞吐量。

测试完成后,你会在`gpt-fast`文件夹中找到一个`trace_compile_none.json`文件。这个文件是使用torch.profiler的跟踪功能生成的。你可以使用Perfetto查看跟踪文件,分析在执行Llama 2 7B模型期间使用的操作符和内核的序列。

通过分析跟踪文件,我们观察到CPU的任务调度(顶部)相对于GPU(底部)的效率不高,这从连续任务之间的间隙可以看出。这些间隙表示GPU的空闲期,由于缺乏活动,资源未被充分利用。接下来,我们将看看`torch.compile`如何帮助缓解这个问题。 

使用 torch.compile(default) 模式对 Llama 2 7B 模型进行性能评估

指定 --compile default 以使用 torch.compile 的默认模式。

%%bash
cd gpt-fast
python generate_simp.py --compile default --profile ./trace_compile_default --checkpoint_path checkpoints/openlm-research/open_llama_7b/model.pth --prompt "def quicksort(arr):"  --max_new_tokens 200  --num_samples 20

输出:

    __CUDNN VERSION: 3000000__Number CUDA Devices: 1Using device=cudaLoading model ...Time to load model: 3.56 secondsReset and set torch.compile mode as  defaultdef quicksort(arr):# Quick sort.## Returns -1, 0, or 1.# If arr is empty, -1 is returned.# If arr is sorted, arr[0] is returned.## If arr is already sorted, 0 is returned.# If arr is not sorted, arr[1] is returned.## See: https://github.com/rails/rails/blob/master/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb#L150-L153arr.sort!n = 0while n < arr.size# if arr[n] < arr[n+1]# quicksort(arr)# arr[n+1] = arr[n]# arr[n] = 1# n += Average tokens/sec: 73.90Memory used: 13.87 GB

使用 torch.compile(reduce-overhead) 模式对 Llama 2 7B 模型进行性能评估

指定 --compile reduce-overhead 以使用 torch.compile 的 reduce-overhead 模式。

%%bash
cd gpt-fast
python generate_simp.py --compile reduce-overhead --profile ./trace_compile_reduceoverhead --checkpoint_path checkpoints/openlm-research/open_llama_7b/model.pth --prompt "def quicksort(arr):"  --max_new_tokens 200  --num_samples 20

输出:

    __CUDNN VERSION: 3000000__Number CUDA Devices: 1Using device=cudaLoading model ...Time to load model: 3.17 secondsReset and set torch.compile mode as  reduce-overheaddef quicksort(arr):# Quick sort.## Returns -1, 0, or 1.# If arr is empty, -1 is returned.# If arr is sorted, arr[0] is returned.## If arr is already sorted, 0 is returned.# If arr is not sorted, arr[1] is returned.## See: https://github.com/rails/rails/blob/master/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb#L150-L153arr.sort!n = 0while n < arr.size# if arr[n] < arr[n+1]# quicksort(arr)# arr[n+1] = arr[n]# arr[n] = 1# n += Average tokens/sec: 74.45Memory used: 13.62 GB

测试完成后,你将在 gpt-fast 文件夹中找到一个名为 trace_compile_reduceoverhead.json 的文件。这是 Llama 2 7B 模型执行过程中生成的追踪文件。

追踪文件显示了一系列 hipGraphLaunch 事件,而在 Eager Mode Section 中获取的追踪文件中没有出现这些事件。`Hipgraph` 使一系列的 hip 内核可以被定义并封装为一个单元,即一系列操作的图形,而不是在 Eager Mode Section 中单独启动的操作序列。`Hipgraph` 提供了一种通过单个 CPU 操作来启动多个 GPU 操作的机制,从而降低启动开销。

使用 torch.compile(max-autotune) 模式对 Llama 2 7B 模型进行性能评估

使用 --compile max-autotune 来启用 torch.compile 的 max-autotune 模式。

%%bash
cd gpt-fast
python generate_simp.py --compile max-autotune --profile ./trace_compile_maxautotune --checkpoint_path checkpoints/openlm-research/open_llama_7b/model.pth --prompt "def quicksort(arr):"  --max_new_tokens 200  --num_samples 20

输出:

    __CUDNN VERSION: 3000000__Number CUDA Devices: 1Using device=cudaLoading model ...Time to load model: 3.05 secondsReset and set torch.compile mode as  max-autotunedef quicksort(arr):# Quick sort.## Returns -1, 0, or 1.# If arr is empty, -1 is returned for each partition.# Create two split keys.split_key_a = int(len(arr) / 2)split_key_b = len(arr) - 1# Quick sort for split key a.# Each partition is sorted.## Note that the inner loop is nested.# The outer loop sorts split key a and the inner loop sorts each# partition.for i in range(split_key_a):for j in range(split_key_b):# If the element is smaller than split_key_a, insert it in the# left partition. Otherwise, insert it in the right partition.idx = numpy.searchsorted(arr, split_key_a)ifAverage tokens/sec: 74.58Memory used: 13.88 GB

对比上述四种模式的吞吐量

# Plotting the bar graph
mode =["eager", "default", "reduce-overhead", "max-autotune"]
inference_time=[28.61, 73.90, 74.45, 74.58]
plt.bar(mode, inference_time)
print(inference_time)
print(mode)# Adding labels and title
plt.xlabel('mode')
plt.ylabel('Inference throughput (tokens/sec)')
plt.title('Llama 2 7B')# Displaying the plot
plt.show()

输出:

    [28.61, 73.9, 74.45, 74.58]['eager', 'default', 'reduce-overhead', 'max-autotune']

图表显示,与 Eager 模式相比,在 AMD MI210 和 ROCm 上,`torch.compile` 可以将 Llama 模型的吞吐量提高多达 2.6 倍(图表中越高越好)。

结论

在这篇博客中,我们展示了如何利用 torch.compile 简单地加速在 AMD GPU 上运行的 ResNet、ViT 和 Llama 2 模型。这种方法带来了显著的性能提升,分别实现了 3.5 倍、2.3 倍和 2.6 倍的加速效果。

参考资料

torch.compile 介绍
利用 CUDA Graphs 加速 PyTorch
加速生成性 AI 的第二部分:GPT,快速
TorchDynamo 和 FX Graphs

这篇关于加速 PyTorch 模型:使用 ROCm 在 AMD GPU 上应用 torch.compile的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/1126426

相关文章

大模型研发全揭秘:客服工单数据标注的完整攻略

在人工智能(AI)领域,数据标注是模型训练过程中至关重要的一步。无论你是新手还是有经验的从业者,掌握数据标注的技术细节和常见问题的解决方案都能为你的AI项目增添不少价值。在电信运营商的客服系统中,工单数据是客户问题和解决方案的重要记录。通过对这些工单数据进行有效标注,不仅能够帮助提升客服自动化系统的智能化水平,还能优化客户服务流程,提高客户满意度。本文将详细介绍如何在电信运营商客服工单的背景下进行

中文分词jieba库的使用与实景应用(一)

知识星球:https://articles.zsxq.com/id_fxvgc803qmr2.html 目录 一.定义: 精确模式(默认模式): 全模式: 搜索引擎模式: paddle 模式(基于深度学习的分词模式): 二 自定义词典 三.文本解析   调整词出现的频率 四. 关键词提取 A. 基于TF-IDF算法的关键词提取 B. 基于TextRank算法的关键词提取

水位雨量在线监测系统概述及应用介绍

在当今社会,随着科技的飞速发展,各种智能监测系统已成为保障公共安全、促进资源管理和环境保护的重要工具。其中,水位雨量在线监测系统作为自然灾害预警、水资源管理及水利工程运行的关键技术,其重要性不言而喻。 一、水位雨量在线监测系统的基本原理 水位雨量在线监测系统主要由数据采集单元、数据传输网络、数据处理中心及用户终端四大部分构成,形成了一个完整的闭环系统。 数据采集单元:这是系统的“眼睛”,

使用SecondaryNameNode恢复NameNode的数据

1)需求: NameNode进程挂了并且存储的数据也丢失了,如何恢复NameNode 此种方式恢复的数据可能存在小部分数据的丢失。 2)故障模拟 (1)kill -9 NameNode进程 [lytfly@hadoop102 current]$ kill -9 19886 (2)删除NameNode存储的数据(/opt/module/hadoop-3.1.4/data/tmp/dfs/na

Hadoop数据压缩使用介绍

一、压缩原则 (1)运算密集型的Job,少用压缩 (2)IO密集型的Job,多用压缩 二、压缩算法比较 三、压缩位置选择 四、压缩参数配置 1)为了支持多种压缩/解压缩算法,Hadoop引入了编码/解码器 2)要在Hadoop中启用压缩,可以配置如下参数

Makefile简明使用教程

文章目录 规则makefile文件的基本语法:加在命令前的特殊符号:.PHONY伪目标: Makefilev1 直观写法v2 加上中间过程v3 伪目标v4 变量 make 选项-f-n-C Make 是一种流行的构建工具,常用于将源代码转换成可执行文件或者其他形式的输出文件(如库文件、文档等)。Make 可以自动化地执行编译、链接等一系列操作。 规则 makefile文件

使用opencv优化图片(画面变清晰)

文章目录 需求影响照片清晰度的因素 实现降噪测试代码 锐化空间锐化Unsharp Masking频率域锐化对比测试 对比度增强常用算法对比测试 需求 对图像进行优化,使其看起来更清晰,同时保持尺寸不变,通常涉及到图像处理技术如锐化、降噪、对比度增强等 影响照片清晰度的因素 影响照片清晰度的因素有很多,主要可以从以下几个方面来分析 1. 拍摄设备 相机传感器:相机传

csu 1446 Problem J Modified LCS (扩展欧几里得算法的简单应用)

这是一道扩展欧几里得算法的简单应用题,这题是在湖南多校训练赛中队友ac的一道题,在比赛之后请教了队友,然后自己把它a掉 这也是自己独自做扩展欧几里得算法的题目 题意:把题意转变下就变成了:求d1*x - d2*y = f2 - f1的解,很明显用exgcd来解 下面介绍一下exgcd的一些知识点:求ax + by = c的解 一、首先求ax + by = gcd(a,b)的解 这个

Andrej Karpathy最新采访:认知核心模型10亿参数就够了,AI会打破教育不公的僵局

夕小瑶科技说 原创  作者 | 海野 AI圈子的红人,AI大神Andrej Karpathy,曾是OpenAI联合创始人之一,特斯拉AI总监。上一次的动态是官宣创办一家名为 Eureka Labs 的人工智能+教育公司 ,宣布将长期致力于AI原生教育。 近日,Andrej Karpathy接受了No Priors(投资博客)的采访,与硅谷知名投资人 Sara Guo 和 Elad G

hdu1394(线段树点更新的应用)

题意:求一个序列经过一定的操作得到的序列的最小逆序数 这题会用到逆序数的一个性质,在0到n-1这些数字组成的乱序排列,将第一个数字A移到最后一位,得到的逆序数为res-a+(n-a-1) 知道上面的知识点后,可以用暴力来解 代码如下: #include<iostream>#include<algorithm>#include<cstring>#include<stack>#in