ROCm + JAX 镜像技术文档

1. 镜像概述和主要用途

本镜像基于 JAX 官方发布版本构建，集成 AMD ROCm (Radeon Open Compute Platform) 开源 GPU 计算栈，提供开箱即用的 JAX 运行环境，专为在 AMD GPU 硬件上高效运行 JAX 应用程序设计。
核心用途：支持基于 JAX 的机器学研究、深度学训练/推理及高性能计算任务，充分利用 AMD GPU 的计算能力加速数值计算和并行处理。

2. 核心功能和特性

2.1 兼容性

• ROCm 支持：集成最新稳定版 ROCm 驱动及运行时（如 ROCm 5.7+），兼容 AMD 开源计算生态
• JAX 版本同步：与 JAX 官方发布版本（如 jax==0.4.23+）严格对应，包含 jax 及 jaxlib 预编译二进制文件
• GPU 硬件支持：兼容 AMD Radeon RX 6000/7000 系列、Radeon Pro W6000 系列、MI250/MI100/MI50 等数据中心级 GPU

2.2 功能特性

• 完整依赖集成：预装 JAX 核心依赖库（NumPy, SciPy, JAXlib）及常用扩展（如 jax-numpy, jax-scipy）
• GPU 加速优化：内置 ROCm 特化的 JAX 后端，支持自动微分、向量化（vmap）、并行化（pmap）等 JAX 核心功能
• 轻量级设计：基于 Ubuntu 基础镜像，精简系统依赖，最小化镜像体积
• 多平台适配：支持 Linux x86_64 架构，兼容 Docker/Podman 容器运行时

3. 使用场景和适用范围

3.1 典型场景

• 机器学***研究：支持基于 JAX 的神经网络训练（如使用 Flax、Haiku 等框架）
• 深度学***推理：部署经 JAX 优化的模型，利用 AMD GPU 加速推理过程
• 高性能计算：加速科学计算、数值模拟等并行任务（如流体力学、量子化学）

3.2 适用对象

• 需要 AMD GPU 硬件加速的 JAX 用户
• 学术机构、企业研发团队的机器学***/高性能计算项目
• 基于 AMD 数据中心 GPU（如 MI250X）或消费级 GPU（如 RX 7900 XTX）的计算集群

4. 使用方法和配置说明

4.1 前置条件

• 宿主机需安装与镜像 ROCm 版本兼容的内核驱动（如 ROCm 5.7 需 Linux 内核 5.15+）
• 宿主机需启用 IOMMU 并配置 GPU 直通（如用于虚拟化环境）
• Docker/Podman 环境需支持 GPU 设备挂载（推荐 Docker 20.10+ 或 Podman 4.0+）

4.2 Docker 快速启动（docker run）

4.2.1 基础命令（默认配置）

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

4.2.2 挂载本地目录（数据持久化）

docker run -it --rm \
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  --cap-add=SYS_PTRACE \
  --security-opt seccomp=unconfined \
  -v /本地路径:/workspace \  # 挂载宿主机目录到容器内/workspace
  -w /workspace \  # 设置工作目录
  rocm-jax:latest \
  python your_jax_script.py

4.3 Docker Compose 配置示例

创建 docker-compose.yml 文件：

version: '3.8'

services:
  jax-rocm:
    image: rocm-jax:latest
    container_name: jax-rocm-container
    devices:
      - /dev/kfd:/dev/kfd
      - /dev/dri:/dev/dri
    group_add:
      - video
    cap_add:
      - SYS_PTRACE
    security_opt:
      - seccomp:unconfined
    volumes:
      - ./local_data:/workspace  # 挂载本地数据目录
    working_dir: /workspace
    environment:
      - JAX_PLATFORMS=rocm  # 强制使用ROCm后端
      - TF_CPP_MIN_LOG_LEVEL=2  # 抑制TensorFlow日志（JAX依赖）
    command: python train.py  # 启动命令（示例：运行训练脚本）

启动服务：

docker-compose up -d

4.4 环境变量说明

4.5 使用步骤

步骤1：拉取镜像

docker pull rocm-jax:latest  # 或指定版本标签，如 rocm-jax:jax-v0.4.23-rocm5.7

步骤2：启动容器并运行任务

# 示例：运行JAX设备检测脚本
docker run -it --rm \
  --device=/dev/kfd --device=/dev/dri \
  --group-add video \
  rocm-jax:latest \
  python -c "import jax; print('GPU设备数量:', jax.device_count()); print('默认设备:', jax.devices()[0])"

步骤3：验证安装

若输出类似以下内容，表明 ROCm 后端正常加载：

GPU设备数量: 1
默认设备: rocm:0 (AMD Radeon RX 7900 XTX)

5. 注意事项

• 版本兼容性：确保宿主机 ROCm 驱动版本与镜像内置版本一致（如镜像使用 ROCm 5.7，宿主机需安装 5.7 驱动）
• GPU 内存限制：运行大模型时，可通过 --shm-size=16g 调整共享内存大小，避免 OOM 错误
• 权限配置：宿主机需将当前用户添加到 video 和 render 组，以访问 GPU 设备文件

6. 维护与更新

• 镜像版本与 JAX 官方发布同步更新，标签格式为 rocm-jax:jax-<jax版本>-rocm<rocm版本>（如 jax-v0.4.23-rocm5.7）
• 问题反馈：请提交至 JAX 官方 GitHub 仓库（google/jax）或 AMD ROCm 社区（ROCm GitHub）