PyTorch · aiwiki.page

PyTorch 是一个主要使用 C++ 与 Python 编程语言开发的开源机器学习与深度学习框架。它最初由 Facebook 人工智能研究院（FAIR）创建并于 2016 年发布，通过 CUDA 借助图形处理器（GPU）提供高性能的多维张量运算，并配备了一个能在运行时动态构建计算图的自动微分引擎。

该框架的核心架构原则——“即时运行”（define-by-run）的动态图即时执行模式——标志着对早期框架静态计算图的重大革新，允许开发者以动态方式执行神经网络，并使用标准调试工具实时检查中间状态。这种灵活性使 PyTorch 成为学术研究领域的首选框架，在探索 Transformer架构、卷积神经网络和扩散模型等新型架构方面尤为突出。随着时间推移，PyTorch 从学术原型设计逐步拓展至工业级生产部署，为现代基础模型和大语言模型的训练与推理提供了核心支持。

自 2022 年起，该项目由 Linux 基金会旗下的 PyTorch 基金会负责管理；其软件生态系统涵盖核心神经网络原语、模块化数据加载流水线、高度优化的分布式训练库，以及在 PyTorch 2.0 中引入的原生编译器技术栈。PyTorch 支持包括中央处理器（CPU）、GPU 和张量处理单元（TPU）在内的多种专用硬件后端，已成为现代人工智能工程的基础关键技术。

起源与历史演进

PyTorch 的血统可直接追溯至 Torch（尤其是 Torch7），后者是一个主要基于 Lua 和 C 语言开发的开源机器学习库。2010 年代初，Torch 凭借其运行速度、低内存占用以及模块化的 C/CUDA 张量引擎（TH 与 THC），在计算机视觉和自然语言处理实验室中广受欢迎。然而，Lua 语言在很大程度上与由 NumPy、SciPy 和 pandas 等 Python 工具构成的庞大科学计算生态相脱节。早期的 Torch 与 Python 衔接方案（例如 torch-python 和 lutorpy）受制于跨运行时序列化延迟、外部函数接口（FFI）开销，以及跨语言运行时笨拙的内存所有权语义。这些痛点表明，有必要设计一个原生以 Python 为中心、同时保留高性能原生算子的框架。

2016 年，Facebook 人工智能研究院（FAIR）的研究人员与工程师——主要包括 Adam Paszke、Soumith Chintala 和 Ronan Collobert，以及 Sam Gross、Gregory Chanan 和 Edward Yang——启动了 PyTorch 项目。他们重新架构了底层的张量和自动求导逻辑，将其封装为 C/C++ 后端绑定，并通过符合 Python 惯用法的应用程序编程接口（API）对外暴露。与 TensorFlow 1.x 和 Theano 等需要先构建不可变声明式图再执行的传统主流库（静态“先定义后运行”设计）不同，PyTorch 采用了命令式的“运行即定义”动态计算图。各项运算在标准 Python 执行流中即时求值，同时底层引擎记录所有操作以用于动态反向传播。这种动态工作流消除了复杂的符号调试抽象层，并极大简化了可变长序列模型和树形结构网络的构建。该框架迅速在各大主流学术界获得自下而上的广泛采纳；2017 至 2019 年间，PyTorch 在 CVPR、NeurIPS 和 ICML 等顶级会议的接收论文数量上超越了竞争对手。

尽管 PyTorch 迅速占领了学术研究领域，但其动态分发机制带来了解释器开销，给生产环境部署带来了困难，尤其是在资源受限的移动端硬件和高吞吐的服务端后端。为了弥合研究原型与生产效率之间的鸿沟，Facebook 于 2018 年 5 月宣布将 PyTorch 与 Caffe2 合并，后者是由贾扬清主导开发、针对工业级规模和边缘部署优化的执行引擎。2018 年底随着 PyTorch 1.0 的发布，这一整合正式确立，将 Caffe2 的模块化算子和图级优化流水线统一整合到底层 C++ 核心库 ATen 和 c10 中。PyTorch 1.0 还引入了 TorchScript，这是一种中间表示框架，允许开发者将模型从动态 Python 执行转化为可序列化、可优化的计算图格式，从而能够在脱离 Python 解释器的独立 C++ 运行时（LibTorch）中运行。

随着 PyTorch 成为训练现代基础模型和大语言模型的事实标准平台，行业利益相关方希望建立一种中立于供应商的治理模式，以防止该平台被单一商业实体控制。2022 年 9 月，Meta 宣布将 PyTorch 的治理权移交给设立于 Linux 基金会之下的新机构——PyTorch 基金会（PyTorch Foundation）。该基金会设立了一个理事会，成员代表来自 Meta、AMD、亚马逊云科技（AWS）、Google Cloud、微软和英伟达（NVIDIA）。此章程使该项目的开源管理正式化，为跨越现代加速器和 CPU 的版本发布管理、运行稳定性以及跨硬件互操作性确立了多方参与的标准。

核心张量架构与内存模型

PyTorch 的基础抽象是 torch.Tensor，它是一个 $n$ 维数组包装器，底层基于 ATen C++ 张量库管理的扁平连续内存段构建。在 PyTorch 的内部设计中，张量将其数学表示与物理数据存储解耦：一个 torch.Tensor 实例包含结构元数据——包括形状（shape）、步长向量（stride）、字节偏移量和数据类型，同时引用包含原始字节数组的 c10::Storage 缓冲区。步长定义了在任意指定维度上从一个元素跨越到下一个元素所必须遍历的内存元素数量。对于位于多维索引 $(i_0, i_1, \dots, i_{d-1})$ 处的元素，其相对于存储起始地址的标量偏移量计算公式为：

\text{offset} = \sum_{k=0}^{d-1} i_k \cdot s_k

其中 $s_k$ 表示沿维度 $k$ 的步长。这种分离架构实现了免复制的结构变换：诸如切片、.view()、.transpose() 和 .permute() 等重塑形态的操作，仅生成具有变更后步长和形状的新张量头信息，底层仍指向现有的存储缓冲区。如果索引操作导致的内存布局脱离了标准的行主序（row-major order），张量就会变为非连续的；此时若执行要求致密连续物理缓冲区的操作，就需要调用 .contiguous()，该方法会分配新的存储空间并执行内存复制。PyTorch 支持多种标量数据类型，涵盖标准 IEEE 浮点精度（torch.float32、torch.float64、torch.float16）、Google 的 16 位脑浮点数（torch.bfloat16）、复数类型、整数和布尔型容器，以及专门的量化类型（例如用于低精度推理的 torch.qint8 和 torch.quint8）。

PyTorch 通过 torch.device 结构抽象了硬件执行目标，统一了不同计算架构下的设备放置逻辑。每个张量都封装了一个设备属性，用于指示其驻留的内存及内核执行环境。该框架提供了跨宿主机 CPU、基于 CUDA 的 NVIDIA GPU、基于 ROCm（提供兼容 CUDA 的编程接口）的 AMD GPU 以及基于 Metal Performance Shaders（MPS）框架的 Apple Silicon GPU 的原生派发能力。借助 ATen 动态派发系统，调用张量操作会自动映射到特定设备的内核实现，而无需用户编写条件分支代码。设备间的数据传输和跨设备操作通过 tensor.to(device) 等 API 或构造函数参数显式管理，使开发者能够以一致的方式面向各类硬件加速器编程。

对于驻留在 CPU 上的数据，PyTorch 与 NumPy 之间的内存互操作设计以零复制指针映射为核心。torch.from_numpy() 和 torch.Tensor.numpy() 等函数可以构建对应的对象，两者共享相同的物理内存地址，无需克隆或重新分配内存缓冲区。只要底层数组满足内存对齐和字节序要求，PyTorch 就会将 NumPy 数组指针直接导入其 Storage 子系统中。因此，在 PyTorch CPU 张量上执行的原位修改会立即同步反映至关联的 NumPy 数组中，反之亦然。这种零复制桥接仅限于 CPU 宿主机内存；驻留在 CUDA 或 MPS 等非 CPU 设备上的张量，在导出到 NumPy 之前必须先通过 .cpu() 显式传输至宿主机系统内存中。

为了消除同步 GPU 驱动分配（如 cudaMalloc 和 cudaFree）带来的性能开销，PyTorch 实现了专用的 CUDA 缓存分配器（THCCachingAllocator）。该分配器预先向操作系统和 GPU 驱动申请大块内存，并将释放的块组织到划分为小块和大块分配类别的内存池中。当张量超出作用域并被 Python 运行时回收时，其占用的 GPU 内存会归还给 PyTorch 的内部缓存，而不是触发驱动级的释放操作，从而使后续相同或更小尺寸的内存分配能够瞬间完成。为了减少长期运行的训练循环中的内存碎片，分配器会执行内存块分割，并将相邻未分配的内存块合并为更大的连续块，其行为可通过 PYTORCH_CUDA_ALLOC_CONF 等环境变量进行配置。针对高吞吐量的宿主机到设备传输，PyTorch 通过 tensor.pin_memory() 提供锁页（pinned/page-locked）宿主机内存，允许直接内存访问（DMA）硬件引擎在 PCIe 总线上进行异步数据传输，与 GPU 内核执行流并发运行。

动态自动微分（Autograd）

PyTorch 实现了一套基于磁带（tape-based）的动态反向模式自动微分引擎，即 Autograd。与在运行前定义并编译执行图的静态计算框架不同，PyTorch 会在执行前向传播的过程中动态构建一个有向无环图（DAG）。对 requires_grad 属性设为 True 的张量执行的每项操作，都会被记录到内部执行磁带中。该图的叶节点为用户创建的输入张量（is_leaf=True），而中间张量和输出张量则通过其 grad_fn 属性引用生成它们的操作节点。图中的节点对应 C++ 的 Node 对象（例如 AddBackward0 或 MulBackward0），有向边则跟踪连续操作之间的数据流和梯度依赖关系。

该引擎通过反向累积向量-雅可比积（VJP）来计算梯度。对于将 $\mathbb{R}^n$ 映射到 $\mathbb{R}^m$ 的向量值操作 $\mathbf{y} = f(\mathbf{x})$ ，其完整的雅可比矩阵 $J \in \mathbb{R}^{m \times n}$ 包含以下偏导数：

J = \begin{bmatrix} \frac{\partial y_1}{\partial x_1} & \cdots & \frac{\partial y_1}{\partial x_n} \\ \vdots & \ddots & \vdots \\ \frac{\partial y_m}{\partial x_1} & \cdots & \frac{\partial y_m}{\partial x_n} \end{bmatrix}

在处理高维张量时，直接计算或存储 $J$ 是不可行的。Autograd 通过接收传入的上游梯度向量 $\mathbf{v} = \nabla_{\mathbf{y}} \ell \in \mathbb{R}^m$ （其中 $\ell$ 通常是一个标量损失函数），并计算转置向量-雅可比积来避免实例化 $J$ ：

\mathbf{v}^T J = \sum_{i=1}^{m} v_i \frac{\partial y_i}{\partial \mathbf{x}}

当在标量张量上调用 y.backward()，或在带有显式向量 $\mathbf{v}$ 的非标量张量上调用 y.backward(gradient=v) 时，引擎会对计算图执行反向拓扑遍历。计算得到的梯度会沿着每个节点的局部变换规则向后传播，并通过就地相加累积到叶张量的 .grad 属性中。

动态图的生命周期在设计上是瞬态的。在反向传播期间，前向传播中存储的中间激活值缓冲区会在其各自的反向方法执行后立即被消耗并从内存中释放，计算图本身也会被销毁。若要对同一个前向图执行多次连续的梯度计算，需要在调用 .backward() 时传入 retain_graph=True；而计算高阶导数则需要设置 create_graph=True，以便为反向传播本身构建计算图。为了消除在计算预测值时的跟踪开销，PyTorch 提供了相应的上下文管理器：torch.no_grad() 会在保留张量版本计数器的同时禁用梯度跟踪；而 torch.inference_mode() 则会彻底禁用 Autograd 跟踪和张量版本跟踪，从而在神经网络推理阶段实现更低的执行延迟和更小的内存占用。

用户可以通过继承 torch.autograd.Function 并定义两个静态方法 forward() 和 backward() 来扩展 Autograd。forward(ctx, *args, **kwargs) 方法负责执行张量计算，其中求导所需的中间张量和非张量元数据可以使用 ctx.save_for_backward() 保存，或直接作为属性附加到上下文对象 ctx 上。backward(ctx, *grad_outputs) 方法通过 ctx.saved_tensors 解包保存的缓冲区，并应用解析梯度推导，为每个需要求导的输入参数返回向量-雅可比积，对不可微参数则返回 None。通过 Function.apply(*args) 调用该操作，即可将自定义算子直接嵌入到执行图中，与 PyTorch 的原生操作无缝协作。

神经网络构建与优化

在 PyTorch 中，构建计算架构的基础抽象是 torch.nn.Module。深度学习模型通常继承 nn.Module，在构造函数（__init__）中定义内部层，并通过重写 forward 方法来指定执行流程。在 __init__ 中进行 Python 属性赋值时，会自动触发内部钩子，将 torch.nn.Parameter 实例（用于封装表示可训练权重的 torch.Tensor）以及嵌套的子模块注册到私有有序字典（_parameters 和 _modules）中。不可训练的运行变量（例如归一化层中的运行时统计量）则通过 register_buffer() 进行追踪，从而确保在通过 .to(device) 进行设备转移时与模型参数一同迁移，同时不会接收梯度更新。可以通过 .parameters()、.named_children() 和 .modules() 等方法递归遍历模块的参数层次结构。结构序列化依赖于模块的状态字典（state_dict()），它将字符串标识符映射到参数和缓冲区张量，从而通过 torch.save 和 torch.load 实现持久化存储和模型检查点保存。

PyTorch 在 torch.nn 包中提供了完备的预定义层与函数式原语库。线性变换由 nn.Linear 封装，而空间和时序结构则通过多维卷积算子（nn.Conv1d、nn.Conv2d、nn.Conv3d）与循环神经网络单元（nn.RNN、nn.LSTM、nn.GRU）建模。归一化层包含批量归一化（nn.BatchNorm2d）、层归一化（nn.LayerNorm）、nn.GroupNorm，以及诸如 nn.RMSNorm 的现代替代方案。标准激活函数（包括 nn.ReLU、nn.GELU 和 nn.SiLU）既可以作为模块化的有状态对象使用，也可以作为 torch.nn.functional 中的无状态函数式调用使用。模块的行为在不同的执行阶段会有所不同：调用 model.train() 和 model.eval() 会在所有子模块中切换布尔标志，从而改变诸如随机失活（nn.Dropout，在评估模式下会被跳过）和批量归一化（使用运行时的总体统计量而非批次统计量进行归一化）等层的运行机制。

PyTorch 中的损失函数继承自 torch.nn.modules.loss._Loss，用于评估预测值与目标值之间的差异以启动梯度传播。对于回归任务，nn.MSELoss 计算批次各元素的均方误差：

\mathcal{L}_{\text{MSE}} = \frac{1}{N} \sum_{i=1}^{N} (y_i - \hat{y}_i)^2

对于分类任务，nn.CrossEntropyLoss 计算多分类交叉熵。为了保持数值稳定性，并避免连续运算引发的灾难性抵消或算术下溢，PyTorch 将 Softmax 变换与负对数似然计算直接融合成单一算子：

\mathcal{L}_{\text{CE}} = - \sum_{k=1}^{C} y_k \log \left( \frac{\exp(\hat{y}_k)}{\sum_{j=1}^{C} \exp(\hat{y}_j)} \right)

类似地，二分类任务使用 nn.BCEWithLogitsLoss，该模块利用 Log-Sum-Exp 技巧将 Sigmoid 激活函数与二元交叉熵直接结合，以确保数值稳健性。

参数优化由 torch.optim 包统筹处理，该包对随机优化的底层机制进行了抽象。优化器实例持有对参数张量的引用，并维护一个内部状态映射，用于追踪动量缓冲区、历史梯度移动平均等动态元数据。在规范的训练步骤中，累积在参数张量 .grad 属性中的梯度通过 optimizer.zero_grad(set_to_none=True) 进行清空——与填零相比，将梯度设置为 None 可以释放张量缓冲区并降低内存开销。随后，optimizer.step() 方法应用参数更新公式。常见的优化器包括随机梯度下降（torch.optim.SGD，支持结合 Nesterov 动量）以及诸如Adam优化器（torch.optim.Adam）等自适应矩估计方法。PyTorch 通过 torch.optim.AdamW 提供了去耦权重衰减，直接从权重向量中减去衰减项：

\theta_t = \theta_{t-1} - \eta \lambda \theta_{t-1} - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}

其中 $\eta$ 为学习率， $\lambda$ 为权重衰减系数， $\hat{m}_t$ 为偏差修正后的一阶矩， $\hat{v}_t$ 为偏差修正后的二阶矩。这避免了标准 Adam 中将 $L_2$ 正则化与自适应梯度缩放耦合时所导致的有效权重衰减率失真问题。

复杂的优化方案可通过参数组进行配置，通过向优化器传入字典列表，为任意参数子集指定独立的超参数。这使得逐层设置学习率（在迁移学习和微调中很常见）以及将偏置项和归一化缩放参数排除在权重衰减之外成为可能。在不同轮次或迭代步之间的动态速率调整由 torch.optim.lr_scheduler 中的学习率调度器负责管理，包括 StepLR、CosineAnnealingLR、OneCycleLR 以及 ReduceLROnPlateau 等自适应监控调度器。调度器追踪步数递增，并在无需重新初始化优化器的情况下，依次更新每个参数组的 lr 属性。

数据摄入与流水线管理

PyTorch 通过 torch.utils.data 中的两个核心类将数据集存储与批次采样解耦：数据集抽象和 DataLoader。该框架支持两种不同的数据集范式：

映射式数据集（Map-style datasets，torch.utils.data.Dataset）：实现 __getitem__() 和 __len__() 协议。这类数据集通过整数索引提供对数据样本的 $O(1)$ 随机访问，因而成为适合存放在本地随机存取存储介质中的数据集的标准选择。数据打乱通过在提取数据前对索引序列进行置换来实现。
可迭代式数据集（Iterable-style datasets，torch.utils.data.IterableDataset）：实现 __iter__() 协议以流式读取连续的数据记录。该范式专为随机跳转寻址开销过大或无法实现的顺序访问场景而设计，例如流式网络套接字、动态数据库游标或数以 TB 计的未压缩网络归档文件，且不需要预先计算数据长度。在多工作进程（multi-worker）配置下，每个工作进程都会运行迭代器的一个独立副本；开发者必须使用 torch.utils.data.get_worker_info() 对数据流偏移量进行分片，以避免跨进程重复读取。

torch.utils.data.DataLoader 负责协调并行数据提取、批处理以及内存暂存。为了在 CPU 密集型数据预处理阶段避开 Python 全局解释器锁（GIL）带来的执行瓶颈，数据加载器通过设置 num_workers > 0 来利用多进程并发。工作进程通过 fork 或 spawn 多进程原语进行实例化，生成的批次张量则通过 POSIX 共享内存（/dev/shm）回传给主进程，从而避免了进程间序列化的额外开销。prefetch_factor 参数决定了每个工作进程提前预加载的批次数量。启用 persistent_workers=True 可以在连续的训练轮次（epoch）之间保持工作进程池存活，减少重复销毁和初始化进程的开销。当目标硬件为图形处理器（GPU）时，pin_memory=True 会使用 CUDA 分配器绑定将张量分配在分页锁定（锁页）主机内存中。随后，专用的锁页内存线程会将经由 PCIe 总线向设备内存的异步直接内存访问（DMA）拷贝流水线化，与 CPU 处理和 GPU 核函数执行并发进行。

批次的构建依赖于整理（collation）机制。默认情况下，DataLoader 会调用 default_collate，将单个样本字典或元组的列表转换为堆叠在新建的前导批次维度 $B$ 上的批量 torch.Tensor 对象。在处理自然语言处理和音频任务中常见的变长数据时（例如文本 token 序列中每个样本的序列长度 $L_i$ 各不相同），由于张量维度不匹配，标准堆叠操作会失败。为解决这一问题，用户可以向 collate_fn 传递自定义可调用对象。该自定义整理器可以使用 torch.nn.utils.rnn.pad_sequence 等函数将变长序列动态填充至当前批次的最大长度 $\max(L_i)$ ，构建布尔注意力掩码，或将序列序列化为 torch.nn.utils.rnn.PackedSequence 容器。动态分批和基于长度的分桶采样器（例如将长度相近的样本组合在一起）常与自定义整理操作相结合，以减少无用的填充计算开销。

数据预处理和数据增强流水线与这些抽象原生集成。生态系统中的相关库提供了直接作用于张量或领域原语的模块化变换算子。在计算机视觉领域，torchvision.transforms.v2 提供了批处理、GPU 加速的操作，涵盖几何变换、色彩抖动、仿射扭曲以及 CutMix 和 Mixup 等样本混合增强技术。对于本地磁盘容量不足的基座模型大规模分布式训练，WebDataset 和 Hugging Face Datasets 等流式框架可以挂载到 IterableDataset 流水线中，通过对象存储流式传输分片归档文件（如 .tar 文件），并支持实时解压与缓存。为了在高带宽 I/O 环境下获得极致吞吐量，NVIDIA DALI 等外部硬件加速引擎可以完全绕过受限于 Python 的标准流水线，直接在 CUDA 流中调度端到端解码与数据增强流水线。

PyTorch 2.0 与图编译

PyTorch 在历史上一直采用即时执行（eager-mode）范式（即“即时运行”，define-by-run），在 Python 中顺序执行各项操作。这种动态方法虽然提供了灵活性和符合原生 Python 习惯的调试体验，但由于每个张量操作都要执行独立的 GPU 核函数，因而面临运行时解释开销以及显存带宽饱和的问题。早期的图方案（如 TorchScript）曾试图弥合这一差距，但必须依托严格受限的 Python 语法子集进行静态跟踪或代码标注。2022 年 12 月，PyTorch 宣布推出 PyTorch 2.0（于 2023 年 3 月正式发布），通过统一的 torch.compile API 将框架重心转向即时编译（JIT）。该架构通过包含四个组件的编译器技术栈将动态用户代码与底层硬件执行解耦，这四个组件分别是：TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor。

TorchDynamo 作为图捕获前端，通过帧评估 API（在 PEP 523 中实现标准化）直接在 CPython 运行时层工作，在字节码执行之前拦截评估帧。它解析 CPython 字节码，追踪涉及 PyTorch 张量的数据流，并将直线式执行块提取为表示为 torch.fx 图模块的计算图。为了保持对任意 Python 代码的完全兼容，Dynamo 引入了不变量检查守卫（guard）系统。当守卫条件满足时，Dynamo 会执行预编译好的图产物；如果遇到不支持的动态 Python 语法结构（例如不受支持的第三方 C 扩展或无法在编译期确定的运行时控制流），Dynamo 会触发“图中断”（graph break），使该代码段退回至标准的 Python 即时执行模式，而不会导致程序崩溃中断。

AOTAutograd 解决了传统自动微分编译器的局限性——后者通常仅优化前向传播，而将反向传播推迟至动态运行时进行分发。AOTAutograd 在执行前拦截 PyTorch 的函数式调用，同时记录前向执行及其反向模式微分的求导公式。在下放至后端编译器之前，它会生成一个统一的前向与反向联合计算图。为了在 PyTorch 包含 2,000 多个算子的庞大接口表面上简化图变换，PrimTorch 子系统将高层 ATen 操作分解为由约 180 至 250 个不可约基元算子构成的核心集合。这种规范化使得编译器后端能够专注于针对极小的算子集制定优化规则，同时又不失对复杂上层 API 的覆盖。

TorchInductor 是 PyTorch 2.0 的默认代码生成后端。Inductor 并没有直接生成对预编译算子库的调用，而是执行循环变换、缓冲区分配和激进的算子融合，以缓解现代深度学习工作负载中典型的内存/显存带宽瓶颈。针对 NVIDIA CUDA 和 AMD ROCm 加速器，TorchInductor 会生成由 Triton 编写的高性能代码，并直接编译为 GPU 机器码。针对 x86 和 ARM CPU，Inductor 则生成通过 OpenMP 并行化并利用 SIMD 指令集（如 AVX-512 和 AVX2）进行向量化的 C++ 源代码。通过将逐元素操作、激活函数和规约操作链融合成单个核函数执行通路，Inductor 大幅减少了在高速设备显存之间的中间往返读写。

该编译技术栈还集成了专门的显存优化原语与融合算法核函数。缩放点积注意力原语（torch.nn.functional.scaled_dot_product_attention，简称 SDPA）就是针对 Transformer 工作负载的典型集成。SDPA 可以将注意力计算动态路由至 FlashAttention、基于 CUTLASS 的核函数或高效显存注意力（memory-efficient attention）等后端实现，将注意力计算相对于序列长度 $N$ 的算法空间复杂度从 $O(N^2)$ 降低至 $O(N)$ 。与此同时，选择性激活检查点技术与 AOTAutograd 相结合以计算最优的重计算（rematerialization）策略，在前向传播过程中系统性地丢弃低成本的中间激活值以尽可能降低峰值显存占用，并在反向传播计算梯度时按需重新计算这些值。

分布式训练与可扩展性

PyTorch 为分布式训练提供了多种抽象层，涵盖了从单机基础多 GPU 执行到跨多节点集群的大规模并行计算。最早的抽象层 torch.nn.DataParallel（DP）在单台机器上采用单进程、多线程模型。在 DP 中，模型参数会被复制到各个 GPU 上，输入数据沿批次维度切分，输出则汇总到主设备（rank 0）上以计算损失。这种架构受制于 Python 全局解释器锁（GIL）的开销、线程竞争，且由于主 GPU 需要执行损失聚合与反向梯度分发，还会导致固有的显存负载不均衡。为解决这些局限，PyTorch 引入了 torch.nn.parallel.DistributedDataParallel（DDP），它为每个 GPU 分配一个独立的 Python 进程，彻底消除了 GIL 竞争。DDP 将梯度组织到大小可配置的内存桶（默认 25 MB）中，并实现了通信与计算的重叠：随着反向传播在网络层间逐层向后推进，装满的内存桶会在整个反向传播结束前提前触发异步的环形全规约（ring-all-reduce）集合通信。

DDP 和多节点工作负载中的进程间协调由 torch.distributed 包通过 torch.distributed.init_process_group 进行管理。该运行时支持针对不同硬件拓扑优化的多种通信后端。NVIDIA 集体通信库（NCCL）是 GPU 加速集群的标准后端，它利用 NVLink 和 InfiniBand 互连技术，最大化 all_reduce、all_gather 和 reduce_scatter 等集合通信操作的带宽。Gloo 后端为 CPU 张量提供跨平台通信及回退支持，而消息传递接口（MPI）后端则兼容传统的高性能计算（HPC）基础设施。开发者还可以通过 torch.distributed.new_group 将进程划分为不同的通信组，从而将张量通信隔离在特定节点或加速卡子集之间。

当扩展 Transformer 架构和大语言模型时，若模型参数、梯度及优化器状态超出了单张加速卡的物理显存容量，标准的数据并行便不再适用。为解决该问题，PyTorch 集成了完全分片数据并行（Fully Sharded Data Parallel，即 torch.distributed.fsdp.FullyShardedDataParallel，在 PyTorch 1.11 中作为原生功能引入）。FSDP 基于微软 DeepSpeed 等库推广开来的零冗余优化器（ZeRO）理念，将模型参数、优化器状态和梯度分散分片到各个数据并行进程中。在执行过程中，它在前向传播时按需使用 all_gather 收集缺失的层权重，使用完毕后立即释放以腾出设备显存；反向传播期间重复这一收集过程，并使用 reduce_scatter 来同步分片梯度。对于参数量超过数百亿的模型，FSDP 经常与张量并行（torch.distributed.tensor.parallel）结合使用，后者利用分布式张量原语（DTensor），沿行或列维度将单个权重矩阵（如注意力投影头和前馈层）切分到多个设备上。

多种互补的执行策略可在大规模训练中最大化吞吐量和内存效率。流水线并行将模型各层切分到分配给不同设备的连续流水线阶段中。为了减少 GPU 空闲等待时间（流水线气泡），运行时采用了微批次调度方案，例如 1F1B（一步前向紧接一步反向）。当硬件限制导致无法增大物理批次大小时，梯度累积可在执行 optimizer.step() 和 optimizer.zero_grad() 之前，通过多次调用 loss.backward() 在多个微批次上累积梯度。自动混合精度（torch.amp）进一步提升了计算速度并降低了内存占用，可通过 torch.autocast 上下文管理器与 torch.cuda.amp.GradScaler 调用。Autocast 会将适用的操作（如矩阵乘法和卷积）动态路由至更低精度（torch.float16 或 torch.bfloat16）以利用专用硬件执行单元，同时将数值敏感的操作保留在 torch.float32 中。为了防止 16 位 IEEE 浮点表示在反向传播期间发生下溢，GradScaler 会在反向传播前动态缩放损失值，并在检测到无效梯度（如 NaN 或 Inf）时放弃参数更新。

生产部署与生态系统

为了弥合动态即时模式（eager mode）研究与静态生产环境需求之间的差距，PyTorch 发展出了专门的模型导出途径。在早期，PyTorch 主要依赖 TorchScript，这是一种能够被编译为中间表示的 Python 静态类型子集。TorchScript 计算图可以通过两种互补的机制来实例化：跟踪（torch.jit.trace），即通过传入虚拟输入执行模块，记录张量级操作并丢弃非张量的 Python 控制流；以及脚本化（torch.jit.script），即直接解析 Python 抽象语法树（AST）以保留条件分支和循环。随着 PyTorch 2.0 的发布，该框架引入了基于 TorchDynamo 和 AOTAutograd 构建的提前编译（AOT）图捕获机制 torch.export。与 TorchScript 不同，torch.export 会生成一个由标准规范 ATen 算子（torch.ops.aten）组成的标准化函数式计算图 ExportedProgram。它将动态输入维度建模为符号整数约束（SymInt），从而施加合理的数学边界，彻底消除了运行时的 Python 解释器开销。

在生产服务器的后端部署方面，通过 TorchScript 或 torch.export 导出的模型可集成至 LibTorch C++ 应用程序编程接口。LibTorch 提供了原生运行时绑定，能够在脱离 Python 解释器依赖的情况下执行计算图，从而消除了高吞吐量多线程服务架构中的全局解释器锁（GIL）竞争。PyTorch 还可以通过原生开放神经网络交换（ONNX）导出器（torch.onnx.export），充当跨平台模型表示的主要生成源。这些导出的 ONNX 图可以直接输入到各硬件厂商专用的运行时中，例如面向 NVIDIA 架构的 TensorRT、面向英特尔处理器的 OpenVINO 以及高通神经处理 SDK（Qualcomm Neural Processing SDK），从而支持算子融合、FP16/INT8 精度量化以及针对低延迟神经网络推理的算子内核自动调优。

边缘设备和嵌入式系统的内存与计算预算非常有限，标准桌面运行时难以满足其要求。PyTorch 最初通过 PyTorch Mobile 来应对这些限制，该方案包含一个经过精简的 LibTorch 运行时，并支持算子选择性编译，以缩减 Android 和 iOS 上的二进制文件体积。2023年10月，PyTorch 基金会推出了 ExecuTorch，旨在移动设备、数字信号处理器（DSP）和微控制器领域全面替代 PyTorch Mobile。ExecuTorch 摒弃了标准的 PyTorch C++ 运行时，转而采用一种轻量级、模块化且在执行期间实现零动态内存分配的 C++ 运行时。ExecuTorch 直接在 torch.export 生成的计算图上运行，通过专用委托接口（包括 Apple Core ML/MPS、Arm Ethos-U 和高通 Hexagon QNN）将张量操作路由至各类硬件后端。

PyTorch 的第一方领域库及其更广泛的开源生态系统进一步巩固了其生产应用地位。PyTorch 基金会维护着多个领域专用库，对数据流水线、模型主干网络和自定义 CUDA 内核进行了标准化：

TorchVision：提供经典与现代计算机视觉架构的实现（如 ResNet、Mask R-CNN 和 Vision Transformer）、图像变换流水线，以及非极大值抑制（NMS）和 RoIAlign 等加速空间操作。
TorchAudio：提供与 FFmpeg 和 SoX 的原生 I/O 绑定、音频数字信号处理原语（如梅尔频率倒谱系数和短时傅里叶变换），以及用于语音识别和语音合成的预训练声学模型。
TorchText：曾提供用于自然语言处理的分词程序、数据集加载器和基线词表，后转为由社区主导维护。

除官方软件包外，PyTorch 还是诸多行业标准第三方框架的基础计算引擎。诸如 Hugging Face Transformers、Hugging Face Diffusers 等工具包，以及 vLLM 和 TensorRT-LLM 等推理引擎，均构建于 PyTorch 的核心原语之上，用于在异构数据中心基础设施间分发和部署大语言模型及扩散模型。