Getting started with AMD ROCm containers: from base images to custom solutions — ROCm Blogs
2025年1月16日,作者: Matt Elliott
在从事技术工作超过二十年后,我亲身见证了容器化如何改变我们开发和部署应用程序的方式。容器将应用程序及其依赖项打包成标准单元,使软件在不同环境中具有良好的可移植性和一致性。当我们将这种容器化的力量与 AMD Instinct™ 加速器结合起来时,我们就能快速部署 AI 和机器学习工作负载。这篇博客文章是探索 ROCm 容器化系列文章的第一篇,我想分享一下我对 AMD ROCm™ 容器的一些见解,并向你展示如何构建和定制你自己的 GPU 加速工作负载。你将学习如何选择适当的基础镜像,为你的具体需求修改容器,以及实现 GPU 启用容器化的最佳实践,并尝试亲手操作的例子。
为什么使用容器来处理 AI 工作负载?
最近,我遇到了一个许多人可能会遇到的场景:一个复杂的机器学习应用程序在开发环境中运行良好,但在生产环境中却失败了。根本原因?框架版本、系统库和 Python 依赖项在不同环境中存在微妙差异。以下是容器在现代开发环境中变得如此重要的几个关键优势:
-
一致性: 开发环境与你的生产环境完全匹配
-
可移植性: 应用程序可以一致地在不同系统和 ROCm 版本上运行
-
隔离性: 多个应用程序可以在同一系统上使用不同的 ROCm 版本
ROCm 容器生态系统
AMD 通过其 Docker Hub 仓库 提供几种类型的官方容器。以下是一些最受欢迎的容器列表。
开发容器
这些容器包括 ROCm 工具和库。
-
rocm/dev-ubuntu-22.04: 基于 Ubuntu 22.04 的开发镜像
-
rocm/dev-ubuntu-24.04: 基于 Ubuntu 24.04 的开发镜像
-
rocm/dev-centos-7: 基于 CentOS 7 的开发镜像
特定框架的容器
-
rocm/pytorch: 预构建的 PyTorch 环境
-
rocm/pytorch-nightly: PyTorch 的每日构建
-
rocm/tensorflow: 预构建的 TensorFlow 环境
-
rocm/tensorflow-build: TensorFlow 的每日构建
-
rocm/jax: 预构建的 JAX 环境
-
rocm/jax-build: JAX 的每日构建
-
rocm/vllm: 针对 LLM 推理和基准测试的 ROCm 优化 vLLM 构建
-
rocm/vllm-dev: vLLM 的开发构建
关于标签的重要说明
从 ROCm 6.2.1 开始,我们对标签策略进行了重要的变化:
-
ROCm 6.2.1 and later:
rocm/pytorch:latest
指向最新稳定版的 PyTorch -
Pre-ROCm 6.2.1:
rocm/pytorch:latest
指向开发版本
例如:
# 最新稳定版 PyTorch 环境 docker pull rocm/pytorch:latest # 开发版本(如果需要) docker pull rocm/pytorch-nightly:latest
容器要求和设置#
为了使容器能够访问 GPU,你需要特定的 Docker 选项。这里是命令结构及其说明:
docker run \ --device=/dev/kfd \ # 所有 GPU 共享的主要计算接口 --device=/dev/dri \ # 每个 GPU 的直接渲染接口 --security-opt seccomp=unconfined \ # 启用 HPC 环境的内存映射 --group-add video \ --ipc=host \ -e HIP_VISIBLE_DEVICES=0,1 \ # 可选:限制使用特定的 GPU <image>
查看ROCm安装先决条件页面获取有关配置用户和组权限的更多信息。如果你想限制容器使用特定的GPU,可以有两种方法:
你可以通过两种方式限制GPU访问:
-
使用
HIP_VISIBLE_DEVICES
环境变量(推荐) -
显式映射特定的GPU设备(例如,第一个GPU的`/dev/dri/renderD128`)
实用示例
让我们探索两个常见的容器配置,这展示了两种可用的ROCm容器。我们将从一个基础的开发容器开始,它对于GPU加速应用程序开发非常合适,然后转向一个用于机器学习工作负载的专用容器。
基础开发容器
这个Dockerfile创建了一个我喜欢称为“ROCm开发者工作室”的环境——它是一个最小但完备的GPU开发环境。它适用于构建GPU加速应用程序的团队,提供了基本的ROCm工具,同时易于定制。
FROM rocm/dev-ubuntu-22.04:latest # 安装基本开发工具 RUN apt-get update && apt-get install -y \ build-essential \ git \ cmake \ python3-pip \ && rm -rf /var/lib/apt/lists/* # 设置工作目录 WORKDIR /workspace # 添加你的项目文件 COPY . .
我会逐节解释:
FROM rocm/dev-ubuntu-22.04:latest
我们从AMD的官方ROCm容器镜像开始,该镜像预安装了所有必需的ROCm工具和库。它就像一个工作台,上面已经安装了所有专用的GPU工具。
RUN apt-get update && apt-get install -y \ build-essential \ git \ cmake \ python3-pip \ && rm -rf /var/lib/apt/lists/*
这里是我们添加日常开发工具的地方。可以想象成在你的基础工具集中添加基本的工具:
-
build-essential
: 你的基本编译器和构建工具 -
git
: 用于版本管理 -
cmake
: 用于项目构建和配置 -
python3-pip
: 用于Python包管理
在末尾的 rm -rf /var/lib/apt/lists/* 是一个最佳实践,我总是包括在内——它在安装后清理包列表,以保持我们的镜像精简。
WORKDIR /workspace COPY . .
这设置了我们的活动工作空间并复制了我们的项目文件。我发现 /workspace 是一个干净、逻辑的位置用于开发工作。
使用案例示例
这个容器非常适合一个开发GPU加速图像处理库的团队。ROCm基本镜像提供了所需的框架、编译C++代码的构建工具以及创建方便的绑定的Python支持。使用以下命令构建容器:
docker buildx build -t rocm-dev-envionment .
源代码可以在启动容器时挂载为卷:
docker run -it --device=/dev/kfd --device=/dev/dri \ --security-opt seccomp=unconfined --group-add video \ --ipc=host -v /path/to/code:/workspace \ rocm-dev-envionment
这个设置创建了一个理想的开发循环——你可以在主机上编辑代码,并在容器内编译/测试,所有依赖项都已就位。
我建议将其作为任何ROCm开发项目的起点——它足够简洁以便于理解,但又足够完整,可以立即使用。你可以根据项目需求增加更多工具。
ML培训容器
这是一个用于ML培训的更复杂的示例:
# 从最新的稳定版PyTorch ROCm容器开始 FROM rocm/pytorch:latest # 安装额外的Python包 RUN pip3 install \ pandas \ scikit-learn \ matplotlib \ wandb # 设置培训目录结构 WORKDIR /opt/training RUN mkdir -p /opt/training/data /opt/training/models # 复制培训代码 COPY train.py . COPY requirements.txt . # 默认命令 CMD ["python3", "train.py"]
这个Dockerfile创建了一个基于AMD官方PyTorch容器的随时可用机器学习培训环境。我特别喜欢这种设置,因为它结合了PyTorch的GPU加速和基本的数据科学工具,如
pandas
和scikit-learn
, 同时还集成了wandb用于实验跟踪。容器结构中设置了数据和模型文件的专用目录 (/opt/training/data
和/opt/training/models
), 你可以在运行时轻松挂载这些目录作为卷。当你启动此容器时,它会自动运行你的培训脚本(train.py
), 这使得它非常适合交互式开发和自动化培训管道。
使用这种设置的一个实际技巧: 虽然容器准备好了进行培训, 但你需要在运行时将数据集目录和模型输出目录挂载为卷。类似 docker run -v $(pwd)/data:/opt/training/data -v $(pwd)/models:/opt/training/models your-image-name
的操作可以完成这项工作。这样可以将您的宝贵数据和训练模型安全地存储在主机系统上,同时享受容器化培训环境的所有好处。 Docker Compose 文件可以用于简化启动和停止容器的流程,带有所需的文件夹挂载和其他选项。构建和运行此容器的过程类似于前一个示例:
docker buildx build -t rocm-training . docker run -it --device=/dev/kfd --device=/dev/dri \ --security-opt seccomp=unconfined --group-add video \ --ipc=host \ -v /home/user/models:/opt/training/models \ -v /home/user/data:/opt/training/data \ rocm-training
总结
容器已经改变了我们开发和部署GPU加速应用程序的方式。通过这篇博客文章,我们探索了如何从基本镜像构建自定义ROCm容器,设置GPU加速应用程序的开发环境,以及为ML培训工作负载创建专用容器。你学习了挂载卷、管理依赖项和在容器中实现GPU支持的实用方法。这些示例展示了ROCm容器生态系统如何为开发和生产部署提供坚实基础,使在AMD GPU上工作变得前所未有的简单。在下一篇博客文章中,我将通过使用vLLM进行推理和基准测试的高级示例进行讲解。