Creating a PyTorch/TensorFlow code environment on AMD GPUs — ROCm Blogs
注意: 本博客之前是 AMD lab notes 系列博客的一部分。
目标: 机器学习生态系统正在迅速扩展,我们旨在通过这一系列的机器学习博客,简化移植到 AMD GPU 的过程。
受众: 数据科学家和机器学习从业者,以及在 AMD GPU 上使用 PyTorch/TensorFlow 的软件工程师。无论你是机器学习新手,还是有使用 Nvidia GPU 经验的专家,这篇博客都会对你有帮助。
动机: 当启动一个新的机器学习项目时,你可能会发现 GitHub 上的许多现有代码几乎总是基于 CUDA。如果你有 AMD GPU 并按照他们的说明运行代码,通常是行不通的。我们基于自身经验提供了一些步骤,可以帮助你为实验创建一个工作代码环境,并管理在 AMD GPU 上运行 CUDA 基础代码库的过程。
与现有在线资源的 区别:
- 这是从机器学习从业者的角度出发,指导你避免由于习惯和偏好(如使用 Jupyter Notebooks 和 pip install)而掉入的陷阱。
- 这篇博客并不是教你如何在 ROCm 上安装 PyTorch/TensorFlow,因为仅这一步骤往往不足以成功运行机器学习代码。
- 这篇博客并不是教你如何 HIPify 代码,而是告诉你有时甚至不需要这一步骤。
- 截至目前,这是唯一拥有从头到尾的 AMD GPU 上创建 PyTorch/TensorFlow 代码环境的文档。
前提是已安装 ROCm,请按照 此处 的说明进行安装。
在 ROCm 上安装 PyTorch 或 TensorFlow
本节非常简要地介绍了如何安装 PyTorch 或 TensorFlow:
选项 1. PyTorch
我们建议按照官方 ROCm PyTorch 网站上的说明进行操作。
选项 2. TensorFlow
我们建议按照官方 ROCm TensorFlow 网站上的说明进行操作。
注意:我们还强烈推荐使用预装了PyTorch或TensorFlow的 Docker 镜像。原因是,如果你创建一个虚拟环境或 conda 环境,某些 ROCm 依赖项可能无法正确安装。安装这些依赖项可能并非易事。
注意:你不需要使用 --gpus all
标志来在 AMD GPU 上运行 Docker。
Git 克隆你想要运行的源码
git clone –-recursive <https://github.com/project/repo.git>
根据 GitHub 仓库安装所需的库
-
跳过创建虚拟环境或 conda 环境的命令。它们通常在
machine_install.sh
或setup.sh
文件中。 -
直接查看库列表并移除
torch
和tensorflow
,因为这些库默认是基于 CUDA 的。Docker 容器应该已经预装了适用于 ROCm 的这些库。您通常可以在 requirements.txt 文件中找到库列表。 -
运行
pip3 install –r requirements.txt
,其中requirements.txt
包含以单行格式列出的包名(可能还有包的版本号)。
运行你的代码
如果你能运行你的代码而没有问题,那么你已经在AMD GPU上成功创建了代码环境!
如果不行,那么可能是由于`requirements.txt`中的额外包依赖于CUDA,需要对其进行HIP化以便在AMD GPU上运行。
获取HIP化的库源码
以下是两种HIP化你的代码的选项:
选项1. 查找现有的HIP化库源码
你可以简单地在线或在GitHub上搜索“库名” + “ROCm”。如果存在HIP化代码,它会显示出来。
由于这一步并不容易,这里有一个例子:
如果你试图运行与大型语言模型相关的代码,你可能需要库`bitsandbytes`(见链接)。
在线搜索“bitsandbytes ROCm”你会找到这个fork,它增加了ROCm支持,并有一个HIP编译目标。
git clone https://github.com/agrocylo/bitsandbytes-rocm
cd bitsandbytes-rocm
export ROCM_HOME=/opt/rocm/
make hip -j
python3 setup.py install
注意:安装位置可能带有版本号,例如/opt/rocm-5.5.0。
选项2. 如有必要,对代码进行HIP化
我们建议参考以下教程进行这个选项。
将更改提交到 Docker 映像
一旦按照第一步(“在 ROCm 上安装 PyTorch 或 TensorFlow”)修改完新的 Docker 容器,退出:
exit
提示系统显示已启动的容器列表并找到 Docker 容器 ID:
docker ps -a
通过提交更改创建一个新的映像:
docker commit [CONTAINER_ID] [new_image_name]
总之,这篇文章介绍了如何在 AMD GPU 上创建 PyTorch/TensorFlow 代码环境的关键步骤。ROCm 是一个不断成熟的生态系统,更多的 GitHub 代码最终会包含 ROCm/HIPified 端口。未来的 AMD 实验室笔记将讨论从 CUDA 迁移到 HIP 的具体细节,以及运行流行社区模型的指南(如 HuggingFace)。
作者要感谢 Justin Chang、Rajat Arora 和 Austin Ellis 的有益评论和建议。如果你有任何问题或意见,请访问我们在 GitHub 上的 [讨论区](https://github.com/ROCm/rocm-blogs/discussions) 联系我们。