ROCm PyTorch训练容器 (v25.7)

针对ROCm优化的统一PyTorch基础训练容器。

概述

PyTorch是一款开源机器学***框架，广泛用于模型训练，其GPU优化组件适用于基于Transformer的模型。

ROCm PyTorch训练Docker镜像rocm/pytorch-training:v25.5（可通过AMD Infinity Hub获取）提供预构建的优化环境，用于在AMD Instinct™ MI300X和MI325X加速器上进行模型微调与预训练。详细文档请参考ROCm文档页面：[***]

Docker部署示例

拉取镜像

bash
docker pull rocm/pytorch-training:v25.5

运行容器（模型微调示例）

以下命令展示如何运行容器进行Qwen2模型微调：

bash
docker run -it --rm \
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  --ipc=host \
  -v /本地数据路径:/data \
  -v /模型 checkpoint 路径:/checkpoints \
  rocm/pytorch-training:v25.5 \
  bash -c "cd /workspace && \
           torchtune finetune --model qwen2-7b --data /data/train_data.json --checkpoint /checkpoints/qwen2-7b-base --output_dir /data/output"

说明：

--device：挂载ROCm设备以访问GPU
--group-add video：添加视频组权限
--ipc=host：使用主机IPC命名空间，避免共享内存限制
-v：挂载本地数据和模型文件到容器内

v25.7 发布说明

更新内容

库更新：
- ROCm: 6.4.2
- Python: 3.10.18
- PyTorch: 2.8.0a0+gitd06a406
- Transformer Engine: 2.2.0.dev0+94e53dd8
- Flash Attention: 3.0.0.post1
- hipBLASLt: 1.1.0-4b9a52edfc
- Triton: 3.3.0
增强FP8支持覆盖范围
- Torchtune全权重微调
- Torchtitan预训练
新增Qwen模型微调支持
- Qwen2 - 1.5B、7B
- Qwen2.5 - 32B、72B
- Qwen3 - 8B、32B
支持GPT-OSS模型微调
修复因PyTorch更新导致的SemiAnalysis模型问题

v25.6 发布说明

更新内容

库更新：
- PyTorch: 2.8.0a0+git7d205b2
- Triton: 3.3.0
- Hipblaslt: 0.15.0-8c69191d
- Transformer Engine: 1.14.0+2f85f5f2
开箱即用地全面支持上游pytorch/TorchTune
新增与最新Torchtune匹配的支持矩阵
为所有支持的模型和用法添加基准测试示例
新增统一的Torchtune_Tester.sh脚本，用于Torchtune微调基准测试
新增模型支持及基准测试示例：
- Llama 4 17B_16E (scout)：全权重SFT、LoRA
- Llama 3.2 Vision 11B：仅全权重SFT
- Llama 3.2 Vision 90B：仅全权重SFT
开箱即支持Torchtitan：
- 更新Torchtitan以匹配上游版本
- 基于CK的Flash Attention实现开箱即用地全面支持BF16

已知问题

由于ROCm 6.4运行时问题，仍使用ROCm 6.3.4，已确定解决方法，将在下一版本应用
上游Torchtune中Llama 3.2视觉模型LoRA微调存在已知问题
ROCm 6.3存在内存泄漏问题（将在ROCm 6.4中修复）
移除2024年12月使用的Semi Analysis模型的基准测试支持，因其实现不再支持PyTorch 2.8+

v25.5 发布说明

更新内容

库更新：
- ROCm: 6.3.4
- Triton: 3.2.0
- Hipblaslt: 0.13.0-98e224a3
- Transformer Engine: 1.12.0.dev0+25a33da
全面支持TorchTune：
- 模型支持：Llama-3.3、Llama-3.2、Llama 3.1、Llama 2
- 功能支持：全权重微调、LoRA、qLoRA
Torchtune打包输入设置下性能提升约38%