看起来 PyTorch 正在获得快速 Thunderbolt 通信后端(在 Mac 上分布式模型)
摘要
据报道,PyTorch 正在为 Mac 上的分布式模型训练添加一个快速的 Thunderbolt 通信后端。
查看缓存全文
缓存时间: 2026/07/14 20:30
pytorch/pytorch 源码: https://github.com/pytorch/pytorch
PyTorch 是一个 Python 包,它提供两个高级特性:
- 类似于 NumPy 的张量计算,并具有强大的 GPU 加速能力
- 基于磁带式自动求导系统的深度神经网络
你可以按需复用你喜欢的 Python 包(如 NumPy、SciPy 和 Cython)来扩展 PyTorch。
我们的主干健康度(持续集成信号)可以在 hud.pytorch.org (https://hud.pytorch.org/ci/pytorch/pytorch/main) 上查看。
- 关于 PyTorch 的更多内容
- 一个支持 GPU 的张量库
- 动态神经网络:基于磁带的自动求导
- Python 优先
- 命令式体验
- 快速且精简
- 无痛扩展
- 安装
- Docker 镜像
- 构建文档
- 排查 CI 错误
- 构建 PDF
- 历史版本
- 快速上手
- 资源
- 交流
- 发布与贡献
- 团队
- 许可证
关于 PyTorch 的更多内容
了解 PyTorch 基础 (https://pytorch.org/tutorials/beginner/basics/intro.html)
从更细粒度来看,PyTorch 是一个包含以下组件的库:
| 组件 | 描述 |
|---|---|
| torch (https://pytorch.org/docs/stable/torch.html) | 类似于 NumPy 的张量库,具有强大的 GPU 支持 |
| torch.autograd (https://pytorch.org/docs/stable/autograd.html) | 基于磁带的自动微分库,支持 torch 中所有可微张量操作 |
| torch.jit (https://pytorch.org/docs/stable/jit.html) | 一个编译栈 (TorchScript),用于从 PyTorch 代码创建可序列化、可优化的模型 |
| torch.nn (https://pytorch.org/docs/stable/nn.html) | 一个与 autograd 深度集成、注重最大灵活性的神经网络库 |
| torch.multiprocessing (https://pytorch.org/docs/stable/multiprocessing.html) | Python 多进程,但具有跨进程魔法般的张量内存共享,对数据加载和 Hogwild 训练非常有用 |
| torch.utils (https://pytorch.org/docs/stable/data.html) | DataLoader 和其他便利工具函数 |
通常,PyTorch 被用做以下两种方式之一:
- 用于利用 GPU 能力的 NumPy 替代品
- 提供最大灵活性和速度的深度学习研究平台
进一步阐述:
一个支持 GPU 的张量库
如果你使用过 NumPy,那么你已经用过张量(即 ndarray)。
张量示意图
PyTorch 提供了既可以运行在 CPU 也可以运行在 GPU 上的张量,并极大地加速了计算。我们提供了丰富的张量例程来加速并满足你的科学计算需求,例如切片、索引、数学运算、线性代数、规约等。而且它们快得惊人!
动态神经网络:基于磁带的自动求导
PyTorch 有一种独特的构建神经网络的方式:使用并回放磁带记录器。大多数框架(如 TensorFlow、Theano、Caffe 和 CNTK)采用静态视图。你必须构建一个神经网络并反复重用同一个结构。改变网络的行为意味着需要从头开始。
在 PyTorch 中,我们采用一种称为反向模式自动微分的技巧,允许你任意改变网络行为,而无需额外延迟或开销。我们的灵感来源于多篇相关研究论文,以及当前和过往的工作,如 torch-autograd (https://github.com/twitter/torch-autograd)、autograd (https://github.com/HIPS/autograd)、Chainer (https://chainer.org) 等。虽然这种技巧并非 PyTorch 独有,但它是迄今为止最快的实现之一。为你疯狂的科研带来速度和灵活性的最佳结合。
动态图
Python 优先
PyTorch 并非是一个将 Python 绑定到庞大 C++ 框架上的产物。它被设计为与 Python 深度集成。你可以像使用 NumPy (https://www.numpy.org/) / SciPy (https://www.scipy.org/) / scikit-learn (https://scikit-learn.org) 那样自然地使用它。你可以用 Python 本身编写新的神经网络层,使用你最喜欢的库以及 Cython (https://cython.org/) 和 Numba (http://numba.pydata.org/) 等包。我们的目标是避免在不必要的地方重复造轮子。
命令式体验
PyTorch 被设计为直观、思路线性、易于使用。当你执行一行代码时,它就会被执行。没有异步世界观。当你进入调试器或收到错误信息和堆栈跟踪时,理解起来非常直接。堆栈跟踪精确地指向代码定义的位置。我们希望你永远不必因为糟糕的堆栈跟踪或异步且不透明的执行引擎而花数小时调试代码。
快速且精简
PyTorch 具有极小的框架开销。我们集成了加速库,如 Intel MKL (https://software.intel.com/mkl) 和 NVIDIA(cuDNN (https://developer.nvidia.com/cudnn),NCCL (https://developer.nvidia.com/nccl))以最大化速度。其核心的 CPU 和 GPU 张量及神经网络后端已经成熟并经过了多年的测试。因此,无论运行小型还是大型神经网络,PyTorch 都相当快速。与 Torch 或其他替代方案相比,PyTorch 的内存使用效率极高。我们为 GPU 编写了自定义内存分配器,以确保你的深度学习模型在内存使用上达到最大效率。这使你能够训练比以前更大的深度学习模型。
无痛扩展
编写新的神经网络模块或与 PyTorch 的张量 API 交互被设计得直接且抽象较少。你可以使用 torch API 或你最喜欢的基于 NumPy 的库(如 SciPy)在 Python 中编写新的神经网络层 (https://pytorch.org/tutorials/advanced/numpy_extensions_tutorial.html)。如果你希望用 C/C++ 编写层,我们提供了一个方便且高效的扩展 API,几乎不需要样板代码,也无需编写包装代码。你可以在这里查看教程 (https://pytorch.org/tutorials/advanced/cpp_extension.html) 和示例 (https://github.com/pytorch/extension-cpp)。
安装
二进制包
通过 Conda 或 pip wheels 安装二进制包的命令在我们的网站上:https://pytorch.org/get-started/locally/
NVIDIA Jetson 平台
适用于 NVIDIA Jetson Nano、Jetson TX1/TX2、Jetson Xavier NX/AGX 和 Jetson AGX Orin 的 Python wheels 在此处提供 (https://forums.developer.nvidia.com/t/pytorch-for-jetson-version-1-10-now-available/72048),L4T 容器在此处发布 (https://catalog.ngc.nvidia.com/orgs/nvidia/containers/l4t-pytorch)。它们需要 JetPack 4.2 及以上版本,并由 @dusty-nv (https://github.com/dusty-nv) 和 @ptrblck (https://github.com/ptrblck) 维护。
从源码编译
前置条件
如果你要从源码安装,你需要:
- Python 3.10 或更高版本
- 完全支持 C++20 的编译器,如 clang 或 gcc(在 Linux 上需要 gcc 11.3.0 或更新的版本)
- Visual Studio 或 Visual Studio 生成工具(仅 Windows)
- 至少 10 GB 的可用磁盘空间
- 首次构建需要 30-60 分钟(后续重新构建要快得多)
* PyTorch CI 使用 Visual C++ 生成工具,这些工具随 Visual Studio Enterprise、Professional 或 Community 版本提供。你也可以从 https://visualstudio.microsoft.com/visual-cpp-build-tools/ 安装生成工具。生成工具不随 Visual Studio Code 默认提供。
环境设置示例如下:
-
Linux:
$ source /bin/activate $ conda create -y -n $ conda activate -
Windows:
$ source \Scripts\activate.bat $ conda create -y -n $ conda activate $ call "C:\Program Files\Microsoft Visual Studio\Community\VC\Auxiliary\Build\vcvarsall.bat" x64
conda 环境并非必需。你也可以在标准虚拟环境中进行 PyTorch 构建,例如使用 uv 等工具创建的环境,前提是你的系统已安装了所有必要的、不可作为 pip 包获取的依赖(如 CUDA、MKL)。
NVIDIA CUDA 支持
如果要使用 CUDA 支持编译,从我们的支持矩阵 (https://pytorch.org/get-started/locally/) 中选择一个受支持的 CUDA 版本,然后安装以下内容:
- NVIDIA CUDA (https://developer.nvidia.com/cuda-downloads)
- NVIDIA cuDNN (https://developer.nvidia.com/cudnn) v9.0 或更高版本
- 与 CUDA 兼容的编译器 (https://gist.github.com/ax3l/9489132)
注意:你可以参考 cuDNN 支持矩阵 (https://docs.nvidia.com/deeplearning/cudnn/backend/latest/reference/support-matrix.html) 了解 cuDNN 版本与不同受支持的 CUDA、CUDA 驱动和 NVIDIA 硬件的对应关系。
如果要禁用 CUDA 支持,导出环境变量 USE_CUDA=0。其他可能有用的环境变量可以在 setup.py 中找到。
如果 CUDA 安装在非标准位置,请设置 PATH 以便你想要使用的 nvcc 能被找到(例如 export PATH=/usr/local/cuda-12.8/bin:$PATH)。
如果你是为 NVIDIA 的 Jetson 平台(Jetson Nano、TX1、TX2、AGX Xavier)构建,可在此处找到适用于 Jetson Nano 的 PyTorch 安装说明 (https://devtalk.nvidia.com/default/topic/1049071/jetson-nano/pytorch-for-jetson-nano/)。
AMD ROCm 支持
如果要使用 ROCm 支持编译,安装:
- AMD ROCm (https://rocm.docs.amd.com/en/latest/deploy/linux/quick_start.html) 4.0 及以上版本
- 目前 ROCm 仅支持 Linux 系统。
默认情况下,构建系统期望 ROCm 安装在 /opt/rocm。如果 ROCm 安装在其他目录,必须设置 ROCM_PATH 环境变量指向 ROCm 安装目录。构建系统会自动检测 AMD GPU 架构。也可以使用 PYTORCH_ROCM_ARCH 环境变量显式设置 AMD GPU 架构 (https://rocm.docs.amd.com/projects/install-on-linux/en/latest/reference/system-requirements.html#supported-gpus)。
如果要禁用 ROCm 支持,导出环境变量 USE_ROCM=0。其他可能有用的环境变量可以在 setup.py 中找到。
Intel GPU 支持
如果要使用 Intel GPU 支持编译,请遵循以下步骤:
- PyTorch 针对 Intel GPU 的前置条件 (https://www.intel.com/content/www/us/en/developer/articles/tool/pytorch-prerequisites-for-intel-gpu.html) 说明。
- Intel GPU 支持 Linux 和 Windows。
如果要禁用 Intel GPU 支持,导出环境变量 USE_XPU=0。其他可能有用的环境变量可以在 setup.py 中找到。
获取 PyTorch 源码
git clone https://github.com/pytorch/pytorch
cd pytorch
# 如果你在更新现有仓库
git submodule sync
git submodule update --init --recursive
安装依赖
通用
# 从 PyTorch 目录运行此命令,假设已按照上面的“获取 PyTorch 源码”部分克隆了源代码
pip install --group dev
在 Linux 上
pip install mkl-static mkl-include
# 仅 CUDA:如果需要,为 GPU 添加 LAPACK 支持
# magma 安装:在激活的 conda 环境中运行。指定要安装的 CUDA 版本
.ci/docker/common/install_magma_conda.sh 12.4
# (可选) 如果要使用 torch.compile 搭配 inductor/triton,则安装匹配版本的 triton
# 从克隆后的 pytorch 目录中运行
# 对于 Intel GPU 支持,请在运行命令前显式 `export USE_XPU=1`
make triton
在 Windows 上
pip install mkl-static mkl-include
# 如果需要 torch.distributed,添加这些包
# Windows 上的分布式包支持是原型功能,可能会发生变化。
conda install -c conda-forge libuv=1.51
安装 PyTorch
在 Linux 上
如果你要为 AMD ROCm 编译,请先运行以下命令:
# 仅当你要为 ROCm 编译时才运行此命令
python tools/amd_build/build_amd.py
安装 PyTorch:
# conda 环境的 CMake 前缀
export CMAKE_PREFIX_PATH="${CONDA_PREFIX:-'$(dirname $(which conda))/../'}:${CMAKE_PREFIX_PATH}"
python -m pip install --no-build-isolation -v -e .
# 非 conda 环境(例如 Python venv)的 CMake 前缀
# 激活 venv 后调用以下命令
export CMAKE_PREFIX_PATH="${VIRTUAL_ENV}:${CMAKE_PREFIX_PATH}"
在 macOS 上
python -m pip install --no-build-isolation -v -e .
在 Windows 上
如果要构建旧版 Python 代码,请参考在旧版代码和 CUDA 上构建 (https://github.com/pytorch/pytorch/blob/main/CONTRIBUTING.md#building-on-legacy-code-and-cuda)
仅 CPU 构建
在此模式下,PyTorch 计算将在 CPU 上运行,而不是 GPU。
python -m pip install --no-build-isolation -v -e .
关于 OpenMP 的说明:所需的 OpenMP 实现是 Intel OpenMP (iomp)。要链接 iomp,你需要手动下载该库,并通过调整 CMAKE_INCLUDE_PATH 和 LIB 来设置构建环境。这里的说明 (https://docs.pytorch.org/docs/main/notes/windows.html#building-from-source) 是设置 MKL 和 Intel OpenMP 的示例。如果没有这些 CMake 配置,将使用 Microsoft Visual C OpenMP 运行时 (vcomp)。
基于 CUDA 的构建
在此模式下,PyTorch 计算将通过 CUDA 利用你的 GPU 以实现更快的数值计算。
构建带有 CUDA 的 PyTorch 需要 NVTX (https://docs.nvidia.com/gameworks/content/gameworkslibrary/nvtx/nvidia_tools_extension_library_nvtx.htm)。NVTX 是 CUDA 发行版的一部分,称为“Nsight Compute”。要在已安装的 CUDA 上安装它,请再次运行 CUDA 安装并勾选相应复选框。确保安装了 Nsight Compute 的 CUDA 之后安装 Visual Studio。
目前,VS 2017 / 2019 和 Ninja 被支持作为 CMake 的生成器。如果在 PATH 中检测到 ninja.exe,则 Ninja 将被用作默认生成器,否则将使用 VS 2017 / 2019。如果选择 Ninja 作为生成器,将选择最新的 MSVC 作为底层工具链。
通常还需要其他库,如 Magma (https://developer.nvidia.com/magma)、oneDNN(又称 MKLDNN 或 DNNL)(https://github.com/oneapi-src/oneDNN) 和 Sccache (https://github.com/mozilla/sccache)。请参考安装帮助程序 (https://github.com/pytorch/pytorch/tree/main/.ci/pytorch/win-test-helpers/installation-helpers) 来安装它们。你可以参考 build_pytorch.bat 脚本 (https://github.com/pytorch/pytorch/blob/main/.ci/pytorch/win-test-helpers/build_pytorch.bat) 了解一些其他环境变量配置。
cmd
:: 下载并解压 mkl 包后,设置环境变量,
:: 否则 CMake 会报错 `Could NOT find OpenMP`。
set CMAKE_INCLUDE_PATH={Your directory}\mkl\include
set LIB={Your directory}\mkl\lib;%LIB%
:: 仔细阅读上一节的内容后再继续。
:: [可选] 如果
相似文章
@zcbenz:MLX 在 macOS 上通过 Thunderbolt 实现的 RDMA(远程直接内存访问)现已作为独立库发布……
MLX 面向 macOS 的 RDMA-over-Thunderbolt 实现已独立成库,任何人都能用它组建高速 Mac 集群,在本地跑 AI 负载。
PyTorch分布式:加速数据并行训练的实践经验
本文详细介绍了PyTorch分布式数据并行模块的设计与优化,重点阐述了梯度分桶(gradient bucketing)和计算-通信重叠等技术,这些技术使系统在使用256个GPU时实现了接近线性的可扩展性。
thunderbird/thunderbolt
Thunderbolt 是一个开源的跨平台 AI 客户端,用户可以在本地部署,支持前沿模型、本地模型和本地部署模型,同时保持数据所有权并避免供应商锁定。目前处于早期活跃开发阶段,针对企业用户,支持 Web、iOS、Android、Mac、Linux 和 Windows 平台。
@PyTorch:ExecuTorch 现有一个 MLX 委托,可在 Apple Silicon GPU 上运行 PyTorch 模型。它支持大语言模型、语音转文字、以及……
ExecuTorch 现有一个 MLX 委托,可在 Apple Silicon Mac 上为 PyTorch 模型提供 GPU 加速推理,支持大语言模型、语音转文字以及通过 TorchAO 进行量化的 MoE 模型。
@tom_doerr: 在 16GB 内存 Mac 上运行 35B 模型 https://github.com/walter-grace/mac-code…
该工具支持通过从 SSD 流式加载模型权重,在 16GB Mac 上运行 Qwen3.5-35B 等大型语言模型,经优化配置后最高可达 30 tok/s。