什么是 PyTorch?

PyTorch 是一款开放源代码机器学习框架,在现代 AI 应用中以其灵活性、易用性和性能而闻名。这在一定程度上是因为其与机器学习开发者和 AI 研究者所青睐的热门 Python 高级编程语言相兼容。

PyTorch 的工作原理?

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 中的大规模专家混合模型训练

NeMo™ Automodel 基于经 NVIDIA 性能优化加速的 PyTorch 分布式能力构建,让大规模 MoE 训练更易普及。

快速链接

张量

张量是 PyTorch 的基本构建块。 与多维数组或 NumPy 的 ndarray 类似,张量会存储和操纵模型输入、输出和参数。至关重要的是,PyTorch 张量专为在 NVIDIA GPU 上运行而设计,可实现大规模并行计算以加速训练和推理。

图形

神经网络将嵌套函数应用于参数 (权重和偏差) 来转换输入数据。 目标是通过反向传播计算梯度 (偏导数) 来优化这些参数。

从历史上看,PyTorch 使用动态计算图形 (Eager Mode),在执行代码时即时构建图形。 这种“运行时定义”的方法让调试更轻松,并能很好地处理动态结构。

随着 PyTorch 2.0 的发布,该框架引入了 torch.compile。这使得用户可以在需要时将图捕获为静态结构,并在不更改底层模型代码的情况下针对 NVIDIA GPU 上的性能进行优化。这兼顾了两者的优点:兼具开发阶段动态图的灵活性,以及生产环境中静态图的速度优势。

为何选择 PyTorch?

PyTorch 已从专注于研究的框架,发展为生成式 AI 和生产部署领域的行业标准。它最初由 Meta AI 开发,目前由独立的 PyTorch 基金会 (隶属于 Linux 基金会) 负责管理,NVIDIA 则是其创始首席成员。

灵活高速:PyTorch 基于直观的 Python 前端构建,注重可读性和快速迭代。然而,现代 PyTorch 也专为速度而设计。借助 TorchScriptTorchDynamo 等特性,开发者可以从即时模式实验无缝过渡到高性能生产部署。

生成式 AI 标准:PyTorch 是生成式 AI 革命的原生语言。凭借对使用 FSDP (全分片数据并行) 在数千个 GPU 上进行分布式训练的支持,PyTorch 成为构建大语言模型 (LLM) 和扩散模型的首选框架。

开发者生态系统:PyTorch API 一直保持一致且用户友好,既便于初学者使用,同时为专家提供深度控制能力。 其庞大的生态系统包括面向计算机视觉、自然语言处理 (NLP) 和强化学习的库。因此,只要出现新的 AI 技术,往往也会有相应的 PyTorch 实现。

PyTorch 的用途

PyTorch 框架是全球最先进 AI 背后的引擎,涵盖计算机视觉、强化学习,尤其支撑了生成式 AI 的蓬勃发展。

生成式 AI 和自然语言处理

从实时翻译到智能编码助手,PyTorch 为现代 NLP 提供支持。虽然早期的方法依赖于循环神经网络 (RNN),但该领域已转向 Transformer 架构。PyTorch 为训练 Transformer 提供必要的基础构件,从而能够创建像 gpt-oss 和 NVIDIA Nemotron™ 这样的大规模基础模型。这些模型能够理解上下文、生成文本,并在复杂领域中进行推理。

计算机视觉和多模态 AI

PyTorch 仍然是目标检测和分割等基础计算机视觉任务的标准,但它也是多模态 AI (能够处理并理解文本、图像、音频和视频等多种数据类型的 AI) 兴起背后的引擎。借助 PyTorch,模型能够同时处理文本、图像和音频,为最新的视觉语言模型 (VLM) 和扩散模型提供支持。这助力开发者构建可以根据文本描述生成逼真图像或实时对视觉数据进行推理的系统,并利用 NVIDIA GPU 来处理这些混合模态工作负载的巨大计算要求。

PyTorch 的优势是什么?

规模化性能

与 NVIDIA GPU 结合时,PyTorch 可实现巨大的吞吐量提升。 FlashAttention 和 Transformer 引擎 (利用 NVIDIA Blackwell GPU 上的 NVFP4 精度) 等特性可显著降低训练时间和显存用量。

生产效率

在推理方面,NVIDIA TensorRT™ LLM 采用 PyTorch 作为默认后端最新基准测试显示,结合 NVIDIA Dynamo 和 wide-expert 并行等其他兼容优化后,性能提升超过 10 倍

可扩展性

PyTorch 的分布式训练能力使其能够在数千个 NVIDIA GPU 上自动扩展,这对于训练具有数十亿参数的基础模型至关重要。

GPU 上的 PyTorch

前沿 AI 模型拥有数十亿至数万亿个参数,在密集线性代数、注意力机制和通信原语中呈现出大规模并行性。要高效实现这种并行性,需要在框架、编译器、内核和硬件中实现极致的协同设计。PyTorch 通过其调度器、自动微分引擎和编译器栈,将高级模型图形映射到融合 GPU 内核上,而 GPU 则提供这些内核所面向的执行模型、存储层级和互连。PyTorch 与 GPU 的紧密协同演进,为现代 AI 模型实现可扩展、高吞吐量的训练和推理。

搭载 NVIDIA GPU 的 PyTorch

NVIDIA GPU 利用数千个核心,同时处理海量并行工作负载。大语言模型 (LLM) 由数十亿或数万亿个参数组成,可自然地映射到该架构,从而让训练和推理速度实现指数级提升。

结合 NVIDIA GPU 使用时,PyTorch 提供了业界一流的 NVIDIA 硬件支持。它提供原生 CUDA® 支持 (通过 torch.cuda),使开发者只需一行代码即可将张量迁移到 GPU。

经优化的生态系统:

  • NVIDIA NGC (NVIDIA GPU CLOUD):开发者可以从 NGC 获取经过优化的 PyTorch 容器,其中预加载了最新驱动程序、库和依赖项。
  • 基于 TensorRT-LLM 的原生推理:PyTorch 现已成为 TensorRT-LLM 的默认后端,使开发者能够直接从 PyTorch 检查点无缝编译和优化 LLM,以实现更高吞吐量和更低延迟。
  • 现代化可扩展训练:PyTorch 2.0 现在通过 NVIDIA 的 Transformer 引擎和图形编译 (torch.compile) 原生集成 FP4 执行,以实现最大吞吐量,而全分片数据并行 (FSDP) 可实现在数千个 GPU 上高效扩展超大规模模型。

NVIDIA 深度学习框架和 PyTorch

虽然 PyTorch 是现代开源软件 (OSS) 深度学习生态系统的基石,但要充分释放其在可扩展科学计算中的潜力,需要与优化的硬件加速建立共生关系。以下示意图描绘了这一关键协同效应,并说明 NVIDIA 的全面栈如何在 PyTorch 灵活接口之下充当底层引擎。

软件栈中的这些层对于高性能 AI 至关重要:从 CUDA-X™ 的基础数学原语 (如 cuBLAS 和 cuDNN) 开始,向上延伸至 TensorRT、NeMo 和 Transformer Engine 等专用 NVIDIA 深度学习框架,最终构建于覆盖数据中心 GPU 到云环境的加速基础设施之上。

后续步骤

探索 PyTorch 技术博客

探索 NVIDIA 持续增长的 PyTorch 技术博客库。

NVIDIA PyTorch NGC 目录

PyTorch NGC 容器针对 GPU 加速进行了优化,并包含一组经过验证的库,可实现和优化 GPU 性能。

在 NVIDIA 优化框架上运行 PyTorch

了解主要特性、软件增强与改进、已知问题,以及如何运行 PyTorch 容器。