很多初学者在安装 PyTorch 时都会遇到诸如CUDA、ROCm、MPS、cuDNN、显卡驱动、Python 版本等各种专业术语。它们之间到底是什么关系为什么有的人安装了 CUDA有的人却不用安装为什么同样都是 GPU不同品牌却需要不同版本的 PyTorch本文将从整体架构出发系统梳理 PyTorch 的 GPU 加速生态帮助你建立完整的知识体系。一、PyTorch深度学习框架的核心PyTorch是目前最流行的开源深度学习框架之一由 Facebook AI ResearchFAIR现 Meta AI开发目前由 PyTorch Foundation 维护。很多资料会说“PyTorch 基于 Torch 开发。”这种说法并不准确。更准确地说PyTorch 借鉴了早期 Lua TorchTorch7的设计思想但底层已经完全重构核心采用 CLibTorch/ATen实现而 Python 只是最常用的接口。PyTorch 提供了Tensor张量计算自动求导Autograd神经网络模块torch.nn数据加载DataLoader模型训练与推理分布式训练GPU 加速其中Tensor张量是 PyTorch 的核心数据结构。它与 NumPy ndarray 非常相似但最大的区别在于可以进行自动求导可以直接运行在 GPU 上支持 CPU、CUDA、ROCm、MPS 等多种计算后端例如importtorch xtorch.randn(3,3)print(x.device)# cpuxx.to(cuda)print(x.device)# cuda:0这里真正发生的是Python 代码 → PyTorch → GPU 后端 → GPU 硬件。二、PyTorch 内部架构很多人认为 PyTorch 的执行流程非常简单Python ↓ PyTorch ↓ CUDA ↓ GPU实际上PyTorch 内部还有多个重要模块Python API │ ┌──────────────┴──────────────┐ │ │ torch.nn torch.optim │ Autograd自动求导 │ Dispatcher统一调度器 │ ├──────── CPU Backend ├──────── CUDA Backend ├──────── ROCm Backend └──────── MPS Backend │ 各种计算库 │ GPU / CPU其中1ATenATen 是 PyTorch 的底层张量计算库。几乎所有 Tensor 运算都会最终调用 ATen。例如ab a b torch.matmul()torch.conv2d()最终都会转换成底层计算接口。2AutogradAutograd 是 PyTorch 的自动求导系统。例如loss.backward()梯度就是由 Autograd 自动计算。这也是 PyTorch 能够训练神经网络的重要基础。3DispatcherDispatcher 可以理解成PyTorch 的交通指挥中心。例如xyDispatcher 会自动判断如果数据在 CPUCPU Backend如果数据在 NVIDIA GPUCUDA Backend如果数据在 AMD GPUROCm Backend如果数据在 Apple SiliconMPS Backend因此大多数情况下开发者无需修改算法代码只需要切换设备即可。三、四种计算后端BackendBackend后端负责将 PyTorch 的运算转换成对应硬件能够执行的指令。后端对应硬件主要生态CPU所有 CPUoneDNN、MKL、OpenMPCUDANVIDIA GPUCUDAROCmAMD GPUHIP、rocBLAS、MIOpenMPSApple SiliconMetal Performance Shaders可以理解为PyTorch │ ├── CPU ├── CUDA ├── ROCm └── MPS四、CUDANVIDIA GPU 的计算平台CUDACompute Unified Device Architecture是 NVIDIA 推出的 GPU 并行计算平台和编程模型也是目前最成熟、生态最完善的 GPU 计算方案。需要注意的是CUDA 并不仅仅是一个软件而是一整套 GPU 软件生态。它主要包括CUDA Toolkit │ ├── CUDA Runtime ├── CUDA Driver API ├── cuBLAS ├── cuDNN ├── NCCL ├── nvcc └── 其它开发工具其中CUDA Runtime程序运行所需的运行时库CUDA Driver API负责与显卡驱动通信cuBLAS高性能矩阵运算库cuDNN深度学习加速库NCCL多 GPU 通信库nvccCUDA 编译器因此PyTorch 并不是直接与 GPU 通信而是通过 CUDA Runtime、cuBLAS、cuDNN、NCCL 等组件完成 GPU 加速。五、CUDA Runtime 与 CUDA Toolkit 的区别这是初学者最容易混淆的地方。很多人认为使用 PyTorch GPU 必须安装 CUDA Toolkit。事实上并非如此。需要区分三个概念组件是否必须作用NVIDIA Driver✅ 必须驱动 GPUCUDA Runtime✅ 必须通常随 PyTorch 提供GPU 运行环境CUDA Toolkit❌ 大多数用户无需安装CUDA 开发工具对于绝大多数用户pipinstalltorch torchvision安装的实际上是torch ├── CUDA Runtime ├── cuBLAS ├── cuDNN ├── NCCL因此即使系统没有安装 CUDA ToolkitPyTorch 依然可以正常调用 GPU。只有在以下场景才需要安装 CUDA Toolkit编写 CUDA 程序使用 nvcc 编译 CUDA 代码编译 PyTorch CUDA Extension开发 CUDA 算法六、cuDNN、cuBLAS 与 NCCL很多人误以为 CUDA 完成了所有 GPU 计算。实际上CUDA 更像一个平台真正承担大量工作的是多个专用计算库。1cuBLAS负责矩阵乘法向量计算GEMM 运算Transformer、MLP、Linear 等大量运算都依赖 cuBLAS。2cuDNN负责深度学习相关算子例如卷积Conv池化PoolingBatchNormRNNLSTMTransformer 部分算子训练 CNN 或 Transformer 时大量计算都会调用 cuDNN。3NCCLNCCLNVIDIA Collective Communications Library负责多 GPU 通信。例如DistributedDataParallelDDPFSDP多 GPU 训练多机训练都会大量使用 NCCL 完成AllReduceBroadcastGatherReduceScatter因此对于大模型训练来说NCCL 是不可或缺的重要组件。七、ROCmAMD GPU 的官方计算平台ROCmRadeon Open Compute是 AMD 推出的 GPU 计算平台。可以把它理解成CUDA 在 AMD 世界中的对应方案。其生态包括ROCm │ ├── HIP ├── rocBLAS ├── MIOpen ├── RCCL分别对应 CUDA 世界中的CUDAROCmCUDA RuntimeHIPcuBLASrocBLAScuDNNMIOpenNCCLRCCL目前ROCm 已支持越来越多的 Radeon 与 Instinct 系列 GPU但整体生态成熟度、第三方软件支持以及模型适配速度仍与 CUDA 存在一定差距。八、MPSApple Silicon 的 GPU 后端对于搭载 Apple SiliconM 系列芯片的 MacPyTorch 使用MPSMetal Performance Shaders作为 GPU 后端。可以通过以下方式检测是否可用torch.backends.mps.is_available()使用方式devicetorch.device(mps)MPS 已能够支持大多数常见模型训练与推理但相比 CUDA 仍存在一些限制部分算子尚未实现某些操作可能自动回退到 CPU多 GPU 支持有限生态成熟度仍低于 CUDA因此在苹果平台上MPS 非常适合开发、学习和中小规模模型实验而在大规模训练场景中CUDA 仍是主流选择。九、CPU 后端并非没有优化没有 GPU 时PyTorch 依然可以充分利用现代 CPU 的计算能力。不同平台通常会调用不同的优化库例如平台常见优化库Intel CPUoneDNN、MKLLinuxOpenMPApple CPUAccelerate Framework因此即使在 CPU 环境下PyTorch 也能获得较好的计算性能。十、显卡驱动整个 GPU 生态的基础显卡驱动是操作系统与 GPU 硬件之间的底层通信接口。它主要负责识别 GPU管理显存调度 GPU为 CUDA Runtime 提供底层接口对于 NVIDIA GPU 而言驱动版本决定了系统能够支持的最高 CUDA Runtime 版本。需要注意的是驱动版本号如 551.xx、560.xx与 CUDA 版本号如 12.4、12.6属于两套不同的版本体系不能直接对应。例如CUDA 12.4 需要 Driver ≥ 550.xx不同 CUDA 版本都有对应的最低驱动要求应以 NVIDIA 官方兼容矩阵为准。十一、理解 nvidia-smi 与 nvcc 的区别这是环境配置中最容易产生误解的问题之一。假设执行nvidia-smi输出CUDA Version: 12.6很多人会误认为我安装了 CUDA 12.6。实际上这表示的是当前 NVIDIA 驱动最高支持 CUDA 12.6 Runtime。如果想查看系统是否安装了 CUDA Toolkit应执行nvcc--version因此命令查看内容nvidia-smi驱动版本、GPU 信息、驱动支持的最高 CUDA Runtimenvcc --version已安装的 CUDA Toolkit 版本两者代表的含义完全不同。十二、Python 与 PyTorch 的版本兼容性PyTorch 是 Python 库因此必须与 Python 版本兼容。通常不同版本的 PyTorch 支持不同的 Python 版本范围。新版本 PyTorch 一般要求 Python 3.10 或更高版本具体以官方发布说明为准。建议使用官方支持范围内的 Python 版本并优先参考安装页面的推荐组合。为了避免不同项目之间发生依赖冲突建议使用venvCondaMiniforge等虚拟环境进行隔离。十三、PyTorch 的设备Device抽象PyTorch 提供了统一的设备接口使同一套代码可以运行在不同硬件上。例如devicetorch.device(cpu)devicetorch.device(cuda)devicetorch.device(cuda:0)devicetorch.device(mps)模型model.to(device)数据data.to(device)因此大多数情况下只需修改设备即可在 CPU、NVIDIA GPU、AMD GPU 或 Apple Silicon GPU 之间切换而无需改动模型逻辑。十四、PyTorch 安装包命名说明安装 PyTorch 时经常会看到如下版本torch2.x.xcpu torch2.x.xcu124 torch2.x.xrocm6.x其中后缀含义cpuCPU 版本cu124内置 CUDA 12.4 Runtime 的版本rocm6.xROCm 版本这些后缀表示PyTorch 所携带的运行时后端而不是要求系统必须安装对应版本的 CUDA Toolkit。十五、推荐的安装方式由于 PyTorch、Python、CUDA、ROCm 等版本之间存在兼容关系最推荐的方式是使用 PyTorch 官方安装页面自动生成安装命令。用户只需选择操作系统包管理器pip 或 CondaPython 环境计算平台CPU、CUDA、ROCm即可获得对应的安装命令避免因版本选择错误导致环境配置失败。十六、完整的软件栈与兼容关系从底层硬件到上层应用一个完整的 PyTorch GPU 计算栈如下Python │ PyTorch API │ ┌──────── Dispatcher ────────┐ │ │ │ CPU CUDA ROCm / MPS │ │ │ oneDNN/MKL cuBLAS、cuDNN、NCCL HIP、MIOpen、Metal └─────────────┴──────────────┘ │ 显卡驱动Driver │ GPUNVIDIA / AMD / Apple整个计算链路可以概括为硬件层CPU 或 GPUNVIDIA、AMD、Apple Silicon。操作系统层Windows、Linux 或 macOS。驱动层显卡驱动为操作系统与硬件提供通信能力。计算平台层CUDA、ROCm、MPS 等后端及其运行时与数学库。深度学习框架层PyTorch通过 Dispatcher 调用不同后端完成计算。编程语言层Python作为主要开发接口。只有当这些环节的版本相互兼容时PyTorch 才能顺利调用 GPU 完成训练与推理。总结PyTorch 并不是一个直接操作 GPU 的程序而是构建在多个底层计算平台之上的统一深度学习框架。CUDA是 NVIDIA GPU 的计算生态ROCm是 AMD GPU 的计算生态MPS是 Apple Silicon 的 GPU 后端cuBLAS、cuDNN、NCCL等库负责高性能数学计算与多 GPU 通信显卡驱动提供硬件通信能力PyTorch Dispatcher负责在 CPU、CUDA、ROCm、MPS 等后端之间自动调度Python则为开发者提供统一、易用的编程接口。理解了这些组件之间的关系就能够从会安装 PyTorch进一步提升到理解 PyTorch 为什么能够利用 GPU 加速从而在环境配置、性能优化和问题排查时更加得心应手。