使用 AMD ROCm 容器入门：从基础镜像到定制解决方案

最新推荐文章于 2025-04-01 11:15:42 发布

109702008

最新推荐文章于 2025-04-01 11:15:42 发布

阅读量1.2k

点赞数 18

分类专栏： # ROCm 文章标签：人工智能

本文链接：https://blog.csdn.net/eidolon_foot/article/details/145377233

版权

ROCm 专栏收录该内容

122 篇文章

订阅专栏

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进行推理和基准测试的高级示例进行讲解。