rocm/jax-training

JAX训练Docker镜像文档

镜像概述

JAX训练Docker镜像专为机器学习研究与开发设计，提供预配置的JAX运行环境，集成必要的科学计算与深度学习依赖库。该镜像旨在简化JAX开发环境搭建流程，减少环境配置时间，让用户专注于模型开发与训练任务。

核心功能与特性

• 预配置环境：内置JAX及其核心依赖（如NumPy、SciPy、Optax、Flax），无需手动安装
• 硬件加速支持：原生支持GPU（CUDA）与TPU加速，自动适配硬件环境
• 多版本兼容性：提供不同JAX版本标签（如latest、0.4.20），满足不同项目依赖需求
• 轻量级设计：基于Ubuntu或Alpine基础镜像构建，平衡功能完整性与镜像体积
• 灵活配置：支持通过环境变量自定义训练参数、日志路径及硬件资源分配

使用场景与适用范围

• 机器学习研究：学术或工业界研究人员进行JAX相关算法开发与验证
• 深度学习模型训练：大规模神经网络模型（如Transformer、CNN）的训练与调优
• 科学计算任务：结合JAX的自动微分与向量化能力，执行高性能科学计算
• 教学与实验：高校或培训机构用于JAX框架教学，统一实验环境

使用方法与配置说明

镜像获取

从Docker仓库拉取镜像（以最新版为例）：

bash
docker pull jax/training:latest

基础运行命令

启动交互式终端，验证JAX环境：

bash
docker run -it --rm jax/training:latest python -c "import jax; print('JAX version:', jax.__version__); print('GPU available:', jax.devices())"

高级配置示例

1. 挂载数据卷（持久化存储）

将本地数据集与训练脚本挂载至容器：

bash
docker run -it --rm \
  -v /path/to/local/data:/app/data \
  -v /path/to/local/scripts:/app/scripts \
  jax/training:latest \
  python /app/scripts/train.py --data-path /app/data

2. GPU加速配置

启用GPU支持（需主机已安装NVIDIA Docker运行时）：

bash
docker run -it --rm \
  --gpus all \
  -e JAX_PLATFORM_NAME=gpu \
  jax/training:latest \
  python -c "import jax; print('GPU devices:', jax.devices('gpu'))"

3. 环境变量配置

通过环境变量调整JAX运行参数：

bash
docker run -it --rm \
  -e JAX_ENABLE_X64=True \  # 启用64位浮点数支持
  -e JAX_LOG_LEVEL=INFO \    # 设置日志级别
  -e TOKENIZERS_PARALLELISM=False \  # 禁用分词器并行（避免警告）
  jax/training:latest

Docker Compose配置示例

创建docker-compose.yml文件管理服务：

yaml
version: '3.8'
services:
  jax-training:
    image: jax/training:latest
    volumes:
      - ./data:/app/data
      - ./models:/app/models
    environment:
      - JAX_PLATFORM_NAME=gpu
      - JAX_ENABLE_X64=True
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    command: python /app/scripts/train.py --output-path /app/models

启动服务：

bash
docker-compose up -d

注意事项

• 确保主机已安装Docker Engine（20.10+推荐）及相应硬件驱动（GPU需NVIDIA Driver 450+）
• TPU加速需在Google Cloud等支持TPU的环境中使用，并配置相应权限
• 不同版本镜像依赖库版本可能存在差异，建议根据项目需求指定具体标签（如jax/training:0.4.20）
• 大规模训练任务建议配合Kubernetes等编排工具进行资源管理与扩展