图片作者：Flux.1

PyTorch 是一个非常重要的库，用于构建新的机器学习模型。Meta 公司开发的 PyTorch 库引领了动态编译的趋势，使得调试更加简便。当你从零开始构建架构时，PyTorch 是最易于使用的库。

本文探讨了 PyTorch 的 Tensor 类的内部工作原理，在 Edward Z. Yang 的基础博客文章和我在使用该库过程中的个人经验的基础上。由于几乎所有机器学习库都从 PyTorch 中汲取了灵感，理解这里的底层选择将有助于你理解其他所有库。

让我们动起来！

让我们从最抽象的角度来开始讨论。张量是一个数学概念，它将标量、向量和矩阵推广到n维空间。这意味着一维数组和2x2矩阵都是张量，同样，(b, t, c)这样的矩阵在Transformer中也是张量。

这种变异对计算机科学家而言有什么重要性？

一般来说，矩阵运算最耗内存。进行矩阵乘法所需的内存传输量远远超过任何标量操作所需的传输量。因此，虽然我们希望内存中对标量的表示尽可能高效，为了防止显著的减速，我们还需要高效地表示张量。

虽然有多种方式来布局这些张量，但我将专注于带步张量，因为这是大型语言模型（LLMs）中最常见的类型。如果大家好奇，我可以再写一篇博客文章，介绍在不同机器学习任务中使用的其他布局。

当用户在 PyTorch 中创建张量时，他们通常会想到类似下面的情况：一个特定形状和值的矩阵，他们可能还想以后访问其中的特定值。

作者的图片——用户对张量（tensor）的画图

默认情况下，PyTorch 将值连续存储在内存中，按照 按行顺序 的方式，如下所示（假设每个值都是一个 4 字节的整数，例如）。按行顺序意味着同一行中的元素是连续存储的（这也是 C 和 Rust 等语言默认的存储矩阵方式）。

作者提供的图片 — 其张量的连续内存布局

虽然这种布局很简单，但我们必须存储更多的数据以防止丢失信息。例如，对于上面绘制的布局，我们不知道所表示的张量的维度是(4x1)、(1x4)还是(2x2)。这对我们做的所有操作都有很大影响。我们下面在图中标出张量的维度。

图片由作者提供，附带元数据

有了这些形状和内存中的值，我们现在得想办法获取特定条目。为此，我们用步长来解决。

一个步长是与每个张量维度相对应的值。我们使用这些值来确定在连续内存中移动多远才能找到我们需要的精确值。继续以我们的例子为例：

作者的图片 — 图片中标注的元素 [1][0]。

有了合适的步长，我们将得到所有正确的值。问题是，如何得到合适的步长？这取决于张量的维度，我们从后往前进行计算。对于所有连续的张量，最后一个维度的步长总是为1。一旦完成这一步，接下来，我们需要计算下一个维度的元素跳跃数量。这可以通过将下一个维度的大小乘以最后一个步长值来计算。以(2x2)为例，下一个维度是2，而最后一个步长值为1，因此结果为2。一般情况下，如图。

作者提供——查找步长的一般方法

这种格式是动态的，允许我们通过仅仅改变元数据（metadata）来调整张量的形状。每当想要调整张量的形状而无需复制数据时，我们只需使用 tensor.view(x,y)，只要新维度可以容纳所有数据，我们现在就可以通过这些新的维度访问数据。

让我们再来看看之前的例子做个总结。运行 view 将张量转换成 (4x1) 矩阵之后，我们可以看到，值和内存地址保持不变，而元数据则被更新了。

图片由作者提供，张量变成新形状

有了数据集，是时候看看这些张量是怎么被操作的了。从大体上讲，PyTorch 有一些操作，比如 torch.mm，它可以进行矩阵乘法。虽然这看起来像是一个单一的函数，但实际上，根据设备（如 CPU、CUDA 等）和它持有的数据类型（如 fp32、int 等），它有不同的实现方式。

这种分离也是你通常不能在不同设备上操作张量的原因。这样做会导致运行时错误，因为不同设备上的张量之间无法进行操作。

图片由作者提供 — 因设备故障，出现异常了

仅看这段Python代码，你不会看到这种选择带来的复杂性和性能影响。实际上，在PyTorch内部，有两种主要方法来调用函数：静态调用方式和动态调用方式。其中一种是静态调用方式，另一种是动态调用方式。

在静态调用情况下，函数在编译时就已经确切地确定了——这使得调用速度变得很快，但要求PyTorch做出它并不打算做的假设（例如设备类型、张量类型、数据类型等信息）。换句话说，只有到了运行时，PyTorch才有足够的信息来决定调用哪个函数。因此，PyTorch依赖于动态调度机制。

尽管动态调度确实有其代价。虽然现在动态库的额外存储开销已不再是主要问题，但是管理向后兼容性却是个难题。运行旧模型时，找到合适的依赖组合可能非常痛苦。这意味着在配置 PyTorch 环境时，你需要仔细列出所有可能的依赖项 — 仅仅依赖默认的 pip 包可能会导致版本不匹配和难以追踪的 bug。

一旦PyTorch根据设备和数据类型决定了要运行的操作，就交给用C++编写的低级内核处理。接下来我们来看看这些内核的结构，以及如何自己动手编写内核。

我们谈到的那些调度程序最终链接到底层内核，这些内核非常了不起 (可以查看每个计算类型在此文件夹中的实现)。如果你想为 PyTorch 编写内核，有一些工具可以让你轻松上手。

首先，当你填写正确的模式定义（定义模式的代码如下链接：https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml）时，连接你的核心代码到高层的PyTorch函数的包装器代码会自动生成。这让你可以直接将你的实现与对应的高层PyTorch操作连接起来。

第二，PyTorch 提供了宏来处理常见的情况。例如，底层的 TORCH_CHECK 宏用于生成更准确的调试信息，并防止因非法内存访问而导致程序崩溃。TORCH_CHECK 本身可以理解为一个抽象的 if 和 throw，并包含自定义消息。

    TORCH_CHECK(self.dim() == 1, "期望是1维张量，但实际上得到了", self.dim(), "-D 张量");

最后一个要介绍的是 TensorAccessor 类。这个底层类在内核中传递数据、维度和数据类型。这段代码默认正确处理了步长，使得它比单纯的原始指针接口更加友好。内部正确处理了步长，即使底层内存不连续。

    auto 访问器 = B.accessor<float, 2>();  // 2D 浮点数访问器
    for (int i = 0; i < B.size(0); ++i) {  
      for (int j = 0; j < B.size(1); ++j) {  
        float 值 = 访问器[i][j];  
        std::cout << 值 << "\n";  
      }  
    }

这里我要介绍的最后一个功能是张量的自动梯度计算。

这或许是PyTorch对机器学习库来说最重要的特性之一。Autograd（自动求梯度的简称）解放了程序员，让他们不必手动指定如何计算模型中的梯度。只需清晰地定义前向传播，程序员就可以依靠autograd自动计算后向传播。PyTorch通过在前向传播过程中存储更多的元数据来跟踪操作，从而自动计算出后向传播所需的梯度。

此元数据用于创建一个有向无环动态依赖关系图，以向 PyTorch 展示操作之间的关系。值得注意的是，默认情况下计算图是在运行时生成的（而不是在编译时）。虽然这使得调试和操作非常方便，但这会带来速度和性能上的损失。尽管已经有一些尝试来减少这种代价（特别是在 Torch 2.0 版本中的 torch.compile），但仍存在一些重大问题和挑战 [参见这份 Google 文档，概述了其中的“注意事项”]。

注意，autograd的数据量其实不小。如果你只是用模型来做推理，保留这些元数据会无谓地拖慢你的程序运行速度。为了避免这种情况，我们可以将不需要追踪的张量用with torch.no_grad():包裹起来。

自动微分示例

我将通过一个简单的例子来向你展示为什么自动求导是PyTorch的一个重要特性。

以下是我们创建一个简单的多层感知机（MLP）并对其进行快速随机梯度下降优化的简短演示。

    import torch  
    import torch.nn as nn  

    # 设置随机种子以保证结果可重现  
    torch.manual_seed(0)  

    # 示例数据  
    x = torch.randn(5, 2)     # 5 个样本，每个样本有 2 个特征  
    y = torch.randn(5, 1)     # 目标值  

    # 一个简单的 MLP 模型：2 -> 3 -> 1  
    model = nn.Sequential(  
        nn.Linear(2, 3),  
        nn.ReLU(),  
        nn.Linear(3, 1)  
    )  

    # 损失函数和优化器设置  
    loss_fn = nn.MSELoss()  
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)   # 学习率为 0.01  

    # 一轮训练  
    y_pred = model(x)  
    loss = loss_fn(y_pred, y)  
    loss.backward()           # autograd 自动完成了所有反向传播的计算  
    optimizer.step()

由于autograd会跟踪所有的操作，我们不需要手动计算ReLU或线性层（乘法和加法）的导数。如果不使用PyTorch实现，可能像下面这样：

$
\text{需要手动计算每个操作的导数}
$

    import numpy as np  

    # 假数据输入和目标输出  
    x = np.random.randn(5, 2)  # 5 个样本，2 个特征  
    y = np.random.randn(5, 1)  # 5 个样本，1 个输出  

    # 初始化权重  
    W1 = np.random.randn(2, 3)  # (输入维度, 隐藏层维度)  
    W2 = np.random.randn(3, 1)  # (隐藏层维度, 输出维度)  

    # 前向传播  
    z1 = x @ W1                # 形状: (5, 3)  
    a1 = np.maximum(0, z1)     # ReLU 激活函数  
    y_pred = a1 @ W2           # 形状: (5, 1)  

    # 计算均方误差损失  
    loss = np.mean((y_pred - y)**2)  
    print(f"损失前：{loss:.4f}")  

    # 反向传播过程 (手动计算梯度)  

    # dL/dy_pred  
    grad_y_pred = 2 * (y_pred - y) / y.shape[0]  # 形状: (5, 1)  

    # dL/dW2 = a1^T @ grad_y_pred  
    grad_W2 = a1.T @ grad_y_pred                 # 形状: (3, 1)  

    # dL/da1 = grad_y_pred @ W2^T  
    grad_a1 = grad_y_pred @ W2.T                # 形状: (5, 3)  

    # dL/dz1 = grad_a1 * ReLU'(z1)  
    grad_z1 = grad_a1 * (z1 > 0).astype(float)   # 形状: (5, 3)  

    # dL/dW1 = x^T @ grad_z1  
    grad_W1 = x.T @ grad_z1                      # 形状: (2, 3)  

    # 梯度下降更新  
    lr = 0.01  
    W1 -= lr * grad_W1  
    W2 -= lr * grad_W2  

    # 更新后的前向传播  
    z1 = x @ W1  
    a1 = np.maximum(0, z1)  
    y_pred = a1 @ W2  
    loss = np.mean((y_pred - y)**2)  
    print(f"损失后：{loss:.4f}")

我们很快就会变得难以跟踪这些操作，这既复杂又容易出错。

PyTorch的灵活性使其成为原型设计尖端模型的首选工具库。但当模型从研究进入生产时，新的挑战也随之出现：依赖冲突问题、动态图引起的运行时开销，以及需要针对特定硬件进行手动优化的内核。

这是Luminal（注：我为Luminal贡献代码）这类项目探索不同方法的地方。作为开源框架，Luminal在运行前编译整个模型，通过分析模型的完整计算图来消除冗余操作，自动融合内核，并优化内存布局，从而提高效率。通过预先看到整个模型，Luminal可以一次性应用优化，比如重写矩阵乘法或者静态分配内存，从而在数千次推理中获得回报。

例如，尽管 PyTorch 依赖于手写的 CUDA 内核以获得性能（例如 FlashAttention），Luminal 则根据模型的具体操作和目标硬件自动生成内核。这避免了维护特定硬件内核库的复杂性，但牺牲了一部分底层控制。结果是生成了一个单一、轻量级的二进制文件，没有任何外部依赖——非常适合降低推理成本，同时加快响应速度。

PyTorch最大的优势在于其灵活性和生态系统。但对于大规模部署模型的团队来说，如Luminal这样的工具突显了编译驱动的机器学习的力量：牺牲一些互动性体验以获得随着时间推移而累积的效率提升。

未来没有一种机器学习框架适合所有情况。通过深入了解 PyTorch 的内部工作原理，就能更好地选择或构建适合你下一个挑战的正确工具。

杨, E., “PyTorch 内部机制” (2019), blog.ezyang.com