PyTorch Nightly with ROCm Docker镜像文档

镜像概述和主要用途

本Docker镜像包含最新ROCm（AMD Radeon Open Compute）版本与PyTorch主分支（pytorch/pytorch main branch）尖端代码，通过每日夜间构建（nightly basis）生成。

主要用途：提供PyTorch最新开发版本的运行环境，适用于开发、测试PyTorch新特性，或在AMD GPU上基于最新ROCm和PyTorch代码进行原型验证。

核心功能和特性

• 最新ROCm支持：集成最新版本ROCm，优化AMD GPU硬件加速能力，支持GPU计算任务。
• PyTorch主分支尖端代码：包含PyTorch官方主分支最新提交（未发布）代码，提供尚未正式发布的新功能、性能优化及Bug修复。
• 每日夜间构建：每日自动更新，确保镜像内容与PyTorch主分支和ROCm最新版本同步，时效性强。

使用场景和适用范围

使用场景

• PyTorch开发者测试最新开发特性或修复。
• 研究人员基于PyTorch最新代码验证算法原型。
• 在AMD GPU环境中运行依赖最新PyTorch功能的应用。

适用范围

• 开发与测试环境：适合用于开发、测试阶段，不建议在生产环境中使用（nightly版本稳定性未经长期验证）。
• AMD GPU硬件：需配合支持ROCm的AMD GPU使用（如Radeon RX系列、MI系列等）。

使用方法和配置说明

前提条件

• 已安装Docker Engine（20.10+推荐）。
• 已配置ROCm驱动环境（参考ROCm官方文档）。
• 若使用GPU，需确保Docker支持GPU设备映射（如通过nvidia-container-toolkit类似工具，或ROCm专用容器运行时）。

拉取镜像

假设镜像托管于Docker仓库（如Docker Hub或私有仓库），拉取命令如下（请替换<repo>为实际仓库名称）：

docker pull <repo>/pytorch-nightly-rocm:latest

注：latest标签对应每日夜间构建的最新版本，也可指定具体日期标签（如20240520）获取历史版本。

运行容器

基础运行命令（带GPU支持）

docker run -it --rm \
  --device /dev/kfd \
  --device /dev/dri \
  --group-add video \
  --cap-add=SYS_PTRACE \
  --security-opt seccomp=unconfined \
  <repo>/pytorch-nightly-rocm:latest \
  bash

参数说明：

• --device /dev/kfd --device /dev/dri：挂载ROCm所需的设备文件（AMD GPU访问）。
• --group-add video：添加video用户组，确保容器内访问GPU设备权限。
• --cap-add=SYS_PTRACE --security-opt seccomp=unconfined：提供必要的系统调用权限，避免ROCm运行时受限。

挂载本地目录（开发场景）

挂载本地代码目录至容器内，方便实时开发：

docker run -it --rm \
  --device /dev/kfd --device /dev/dri \
  --group-add video \
  -v /path/to/local/code:/workspace \
  -w /workspace \
  <repo>/pytorch-nightly-rocm:latest \
  bash

• -v /path/to/local/code:/workspace：将本地代码目录挂载到容器内/workspace。
• -w /workspace：设置工作目录为/workspace。

Docker Compose配置示例

创建docker-compose.yml文件，定义容器服务：

version: '3.8'
services:
  pytorch-nightly:
    image: <repo>/pytorch-nightly-rocm:latest
    devices:
      - /dev/kfd:/dev/kfd
      - /dev/dri:/dev/dri
    group_add:
      - video
    cap_add:
      - SYS_PTRACE
    security_opt:
      - seccomp:unconfined
    volumes:
      - ./local_code:/workspace
    working_dir: /workspace
    tty: true
    stdin_open: true

启动服务：

docker-compose up -d

进入容器：

docker-compose exec pytorch-nightly bash

配置参数与环境变量

常用环境变量

镜像支持通过环境变量调整PyTorch及ROCm运行时行为，常见配置如下：

使用示例

运行容器时指定环境变量：

docker run -it --rm \
  --device /dev/kfd --device /dev/dri \
  --group-add video \
  -e HIP_VISIBLE_DEVICES=0 \
  -e PYTORCH_LOG_LEVEL=DEBUG \
  <repo>/pytorch-nightly-rocm:latest \
  bash

注意事项

1. 稳定性：镜像基于PyTorch主分支夜间构建，可能包含未测试的新特性或Bug，不建议用于生产环境。
2. ROCm兼容性：需确保主机ROCm驱动版本与镜像内置ROCm版本匹配（参考镜像标签或文档说明）。
3. GPU支持：仅支持AMD GPU设备，需提前确认硬件兼容性（参考ROCm硬件支持列表）。