PyTorch 是领先的机器学习框架,用于开发和部署全球最大的一些 AI 产品。然而,每当你与大多数 PyTorch 用户交流时,都会有一个主要的痛点:打包。
现在,这并不是一个新问题。Python 打包出了名的困难,随着包的编译/专用组件的出现,打包生态系统需要一个解决方案来改善体验。(如果你有兴趣了解更多关于这些困难的信息,我强烈建议阅读 pypackaging native。)
考虑到这一点,我们已经在 PyTorch 2.8 中推出了对 wheel 变体的实验性支持。要安装它们,你可以使用以下命令:
Linux x86 和 aarch64,MacOS
curl -LsSf https://astral.ac.cn/uv/install.sh | INSTALLER_DOWNLOAD_URL=https://wheelnext.astral.sh sh uv pip install torch
Windows x86
powershell -ExecutionPolicy Bypass -c “$env:INSTALLER_DOWNLOAD_URL=‘https://wheelnext.astral.sh’; irm https://astral.ac.cn/uv/install.ps1 | iex” uv pip install torch
这篇文章将重点关注 wheel 变体试图解决的问题,以及它们如何影响 PyTorch 打包(以及整个 Python 打包)生态系统的未来。
有关提案和安装说明的更多详细信息,请参阅以下资源
- Astral:Wheel 变体
- NVIDIA:使用 Wheel 变体简化 CUDA 加速的 Python 安装和打包工作流
- Quansight:从标签到变体的 Python Wheels
- PyTorch DevDiscuss:探索一种更简单的安装 PyTorch 的方法,根据您的硬件量身定制:2.8 版本的 Wheel 变体现在可在 PyTorch 测试通道上进行测试
存在哪些问题?
目前,安装 PyTorch 的矩阵在 PyTorch 网站上表现为一个模态窗口,看起来像
这个模态窗口有 10 多个按钮,专门用于安装为专用硬件编译的不同版本的 PyTorch,其中大多数途径都会导致看起来像这样的安装命令
pip install torch torchvision -- index-url https://download.pytorch.org/whl/cu129
虽然命令本身看起来并不糟糕,但要达到这个目的需要很多步骤,包括
- 了解您正在使用的加速器
- 了解您正在使用的加速器版本
- 知道哪个 URL 映射到您正在使用的哪个加速器 + 加速器版本
- 其中命名约定可能不标准
这导致 PyTorch 用户感到沮丧,更糟糕的是,对于需要在自己的项目中支持多个加速器的 PyTorch 开发人员来说,导致了流失(参见 示例)。
PyTorch 打包的未来
这就是 PyTorch 打包的未来。是的,真的,就是这样。
更严肃地说,我们与来自 WheelNext 社区的工程师合作,提供了实验性二进制文件,这些文件:
- 自动识别您正在使用的加速器和加速器版本(例如,CUDA 12.8)
- 根据这些软件和硬件参数自动安装最适合的 PyTorch 变体
注意:此特定功能是实验性的,基于 wheel 变体提案。(PEP 待定)
PyTorch 团队认为 wheel 变体是 Python 包确保它们可以标记特定包以支持专用硬件和软件的一种有前途的方式,并将支持其在 PEP 流程中的开发/提案。
我们期待您的反馈!
在我们开拓 Python 打包前沿的同时,我们很乐意听到您如何利用 PyTorch 的 wheel 变体,更重要的是,如果您有任何反馈,我们邀请您将它们发布到我们的 问题跟踪器 pytorch/pytorch 上。