PyTorch 是一款开放源代码机器学习框架,在现代 AI 应用中以其灵活性、易用性和性能而闻名。这在一定程度上是因为其与机器学习开发者和 AI 研究者所青睐的热门 Python 高级编程语言相兼容。
PyTorch 支持两大主要 AI 工作负载:训练和推理。
在训练过程中,PyTorch 通过编排一个持续的循环,辅助神经网络的学习过程:存储为张量 (Tensor) 的输入数据流经模型,从而生成预测结果。该框架的定义特征在于其自动微分引擎 autograd,该引擎会自动准确计算模型参数需要如何变化以降低错误 (反向传播)。该过程涉及在 NVIDIA GPU 上并行处理的海量计算工作负载,使模型能够在大规模时迭代并提高其模型精度。
模型经过训练后,即进入 AI 推理阶段,在该阶段中,它会应用于其所学到的模式来处理新的现实世界数据。在这个阶段,PyTorch 将重点从学习转向执行,让模型能够生成内容或进行预测。这种精简性能对于部署至关重要,使开发者能够采用复杂的生成式 AI 模型,并在加速计算硬件上高效运行,以提供实时、低延迟的响应。
张量是 PyTorch 的基本构建块。与多维数组或 NumPy 的 ndarray 类似,张量会存储和操纵模型输入、输出和参数。至关重要的是,PyTorch 张量专为在 NVIDIA GPU 上运行而设计,可实现大规模并行计算以加速训练和推理。
神经网络将嵌套函数应用于参数 (权重和偏差) 来转换输入数据。目标是通过反向传播计算梯度 (偏导数) 来优化这些参数。
从历史上看,PyTorch 使用动态计算图形 (Eager Mode),在执行代码时即时构建图形。这种“运行时定义”的方法让调试更轻松,并能很好地处理动态结构。
随着 PyTorch 2.0 的发布,该框架引入了 torch.compile。这使得用户可以在需要时将图捕获为静态结构,并在不更改底层模型代码的情况下针对 NVIDIA GPU 上的性能进行优化。这兼顾了两者的优点:兼具开发阶段动态图的灵活性,以及生产环境中静态图的速度优势。
快速链接
PyTorch 已从专注于研究的框架,发展为生成式 AI 和生产部署领域的行业标准。它最初由 Meta AI 开发,目前由独立的 PyTorch 基金会 (隶属于 Linux 基金会) 负责管理,NVIDIA 则是其创始首席成员。
灵活高速:PyTorch 基于直观的 Python 前端构建,注重可读性和快速迭代。然而,现代 PyTorch 也专为速度而设计。借助 TorchScript 和 TorchDynamo 等特性,开发者可以从即时模式实验无缝过渡到高性能生产部署。
生成式 AI 标准:PyTorch 是生成式 AI 革命的原生语言。凭借对使用 FSDP (全分片数据并行) 在数千个 GPU 上进行分布式训练的支持,PyTorch 成为构建大语言模型 (LLM) 和扩散模型的首选框架。
开发者生态系统:PyTorch API 一直保持一致且用户友好,既便于初学者使用,同时为专家提供深度控制能力。 其庞大的生态系统包括面向计算机视觉、自然语言处理 (NLP) 和强化学习的库。因此,只要出现新的 AI 技术,往往也会有相应的 PyTorch 实现。
PyTorch 框架是全球最先进 AI 背后的引擎,涵盖计算机视觉、强化学习,尤其支撑了生成式 AI 的蓬勃发展。
从实时翻译到智能编码助手,PyTorch 为现代 NLP 提供支持。虽然早期的方法依赖于循环神经网络 (RNN),但该领域已转向 Transformer 架构。PyTorch 为训练 Transformer 提供必要的基础构件,从而能够创建像 gpt-oss 和 NVIDIA Nemotron™ 这样的大规模基础模型。这些模型能够理解上下文、生成文本,并在复杂领域中进行推理。
PyTorch 仍然是目标检测和分割等基础计算机视觉任务的标准,但它也是多模态 AI (能够处理并理解文本、图像、音频和视频等多种数据类型的 AI) 兴起背后的引擎。借助 PyTorch,模型能够同时处理文本、图像和音频,为最新的视觉语言模型 (VLM) 和扩散模型提供支持。这助力开发者构建可以根据文本描述生成逼真图像或实时对视觉数据进行推理的系统,并利用 NVIDIA GPU 来处理这些混合模态工作负载的巨大计算要求。
前沿 AI 模型拥有数十亿至数万亿个参数,在密集线性代数、注意力机制和通信原语中呈现出大规模并行性。要高效实现这种并行性,需要在框架、编译器、内核和硬件中实现极致的协同设计。PyTorch 通过其调度器、自动微分引擎和编译器栈,将高级模型图形映射到融合 GPU 内核上,而 GPU 则提供这些内核所面向的执行模型、存储层级和互连。PyTorch 与 GPU 的紧密协同演进,为现代 AI 模型实现可扩展、高吞吐量的训练和推理。
NVIDIA GPU 利用数千个核心,同时处理海量并行工作负载。大语言模型 (LLM) 由数十亿或数万亿个参数组成,可自然地映射到该架构,从而让训练和推理速度实现指数级提升。
结合 NVIDIA GPU 使用时,PyTorch 提供了业界一流的 NVIDIA 硬件支持。它提供原生 CUDA® 支持 (通过 torch.cuda),使开发者只需一行代码即可将张量迁移到 GPU。
经优化的生态系统:
虽然 PyTorch 是现代开源软件 (OSS) 深度学习生态系统的基石,但要充分释放其在可扩展科学计算中的潜力,需要与优化的硬件加速建立共生关系。以下示意图描绘了这一关键协同效应,并说明 NVIDIA 的全面栈如何在 PyTorch 灵活接口之下充当底层引擎。
软件栈中的这些层对于高性能 AI 至关重要:从 CUDA-X™ 的基础数学原语 (如 cuBLAS 和 cuDNN) 开始,向上延伸至 TensorRT、NeMo 和 Transformer Engine 等专用 NVIDIA 深度学习框架,最终构建于覆盖数据中心 GPU 到云环境的加速基础设施之上。