相关说明

这篇文章的大部分内容参考自我的新书《解构大语言模型：从线性回归到通用人工智能》，欢迎有兴趣的读者多多支持。
本文涉及到的代码链接如下：regression2chatgpt/ch07_autograd/gpu.ipynb

本文将讨论如何利用PyTorch实现GPU计算。本文讨论的内容将在后续的模型搭建中应用非常广泛：

• 利用神经网络学习语言（五）——长短期记忆网络（LSTM）
• 理解大语言模型（二）——从零开始实现GPT-2

关于GPU更多的使用技巧请参考：大语言模型的工程技巧（二）——混合精度训练

一、概述

在人工智能领域，模型运算通常涉及大量可以并行执行的操作。以简单的线性回归模型为例，如图1所示，在单个数据点的模型运算中，相应的计算节点之间并没有依赖关系，每个乘法操作都可以并行执行。从数学角度来看，这样的运算可以表示为张量的矩阵乘法，产生单一的数值输出。

图1

对于只有一个输出值的矩阵乘法，都可以使用并行计算来提高执行速度。更进一步，如果考虑多个模型对多个数据点的运算，相应的矩阵乘法输出是一个张量，而这个张量中的每个元素的运算都是相互独立的，这使得并行执行更加可行。类似地，张量的其他运算，例如加法和减法，也可以通过并行执行来加速计算。在神经网络领域，这一点尤为重要，因为神经网络中包含多个线性模型。整个运算过程中涉及大量的张量运算，如何有效加速这些运算过程成为这个领域工程实现中要考虑的首要问题。

二、CPU v.s. GPU

默认情况下，在提到计算机运算时，通常指的是由中央处理器（CPU）作为计算核心执行的运算任务。然而，CPU的设计原则使它并不擅长处理并行计算。举个例子，考虑矩阵乘法，如图2所示，CPU会按顺序逐个元素相乘，导致整个运算过程时间过长。虽然我们可以利用多核CPU进行并行计算以加速任务，但多核CPU价格昂贵，在实际应用中难以普及。

图2

与CPU不同，GPU（图形处理器）最初是为图形处理而设计的，属于边缘计算组件。它通常拥有大量的计算核心（有时高达上千个），可以并行地执行大量的简单计算任务。因此，使用GPU进行张量运算能够显著缩短运算时间。对于神经网络而言，这一点具有重要意义，它使得构建更大、更复杂的模型成为可能（使用GPU，模型训练时间可以从几年缩短到几天）。事实上，GPU的大规模应用是神经网络，尤其是深度学习得以迅猛发展的原因之一。因此，在涉及神经网络模型的实际应用中，总是优先考虑使用GPU作为计算的核心（这是编程领域一次重大的范式转变）。随着神经网络的迅猛发展，业界甚至在考虑重新设计计算机架构来提升GPU在硬件中的地位，以便更好地适应这一快速发展的领域。

三、代码实现

借助PyTorch（基于CUDA¹）在GPU计算方面的卓越工作，GPU计算的代码实现并不复杂。如程序清单7-9所示（完整代码），只需简单地将数据移动到相应的GPU中即可。与CPU的架构设计不同，GPU无法直接读取在内存（RAM）中的数据进行运算，因此需要将数据复制到GPU专用的内存（GPU Dedicated Memory）中。同理，GPU也不支持跨计算核心的运算。这包括跨CPU和GPU的计算，因为数据一部分存放在内存中，另一部分存放在GPU专用内存中；以及跨不同GPU的计算，因为数据分布在不同的GPU专用内存中。在实际使用中，面对这样的需求，需要在进行计算之前将数据移动到同一个计算核心上。

程序清单1 GPU计算

 1 |  # 检查是否有GPU2 |  torch.cuda.is_available()        # True3 |  # GPU的个数4 |  torch.cuda.device_count()        # 15 |  # 默认情况下，创建的张量存放在内存中，使用CPU进行计算6 |  x = torch.randn(2, 3)7 |  print(x.is_cuda)                 # False8 |  # 可以使用张量提供的函数，将数据移到GPU上9 |  # 当有n个GPU时，相应的设备id是cuda:0, cuda:0,... ,cude:n-1
10 |  print(x.to('cuda:0').is_cuda)    # True
11 |  # 在创建张量时，通过指定device将张量移到GPU上 
12 |  y = torch.randn(2, 3, device='cuda:0')
13 |  print(y.is_cuda)                 # True
14 |  print(y.to('cpu').is_cuda)       # False
15 |  # 不支持跨计算核心运算
16 |  x + y                            # error

想要正确使用GPU进行计算，关键是要充分利用其并行计算的特性，而不是仅仅进行串行计算，图3所示的例子很好地阐述了这个观点。当计算任务主要是串行的时，CPU往往比GPU更具优势。当计算任务能够被有效地并行化时，GPU的优势才能充分显现出来。

图3

---

1. CUDA（Compute Unified Device Architecture）是由NVIDIA开发的并行计算平台和编程模型。它允许开发者利用NVIDIA的GPU进行通用目的的并行计算，包括科学计算、机器学习、深度学习等领域的任务。CUDA的出现彻底改变了GPU的角色，使其从原本的图形渲染加速变为通用计算的利器。使用CUDA，开发者可以将计算任务分配到GPU上，并利用其大量的计算核心实现并行处理，从而大幅提升计算速度。
   为了使用CUDA，开发者需要编写基于CUDA的GPU核心代码，然后通过相应的编译工具将其与CPU代码结合起来，以充分利用GPU的计算能力。同时，许多深度学习框架（如TensorFlow、PyTorch等）也提供了对CUDA的支持，使开发者能够更方便地在GPU上进行模型训练和推断。 ↩︎