rocm/pytorch-training repository overview

⁠ROCm PyTorch Training Container

Unified PyTorch base container optimized for training with ROCm.

⁠[!NOTE]

⁠For a unified training solution on AMD GPUs with ROCm, the rocm/pytorch-training Docker Hub registry will be deprecated soon in favor of rocm/primus. The rocm/primus Docker containers will cover PyTorch training ecosystem frameworks, including torchtitan and Megatron-LM.

⁠See Training a model with Primus and PyTorch⁠ for details.

⁠Overview

PyTorch is an open-source machine learning framework that is widely used for model training with GPU-optimized components for transformer-based models.

The ROCm PyTorch Training Docker rocm/pytorch-training:v25.5 container, available through AMD Infinity Hub, provides prebuilt, optimized environment for fine-tuning, pre-training a model on the AMD Instinct™ MI300X and MI325X accelerator. For detailed documentation, please refer to ROCm documentation page here: https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/training/benchmark-docker/pytorch-training.html⁠

⁠v25.11 Release Notes :

####Updates:

• Primus-LM 0.6.0

  MoE Package full support, for more details please check our blog⁠

  Interactive Auto-benchmarking tool

• MoE performance further improved up to 28% from v25.10 with fully enabled MoE package

⁠v25.10 Release Notes :

####Updates:

• Merge MI300X and MI355X docker

• Primus TorchTitan backend upgrade to recent version

• Primus Megatron backend upgrade to recent version

• Primus v0.5.0

  Support MLFlow

  Improved Megatron-LM profiling

  MoE Package Feature available:MoE CPU Sync-free, 1f1b MoE overlapping, Zero-bubble, Router Fusion

  Tools: Model based memory projection/modeling tool

• Primus-Turbo v0.2.0

  FP8 GEMM

  Grouped GEMM

  DeepEP support

• Added Primus Torchtitan MoE support with matching performance of PyTorch Conference presentation

  DeepSeekV3

• Diffusion Benchmark with Video Gen models

  Hunyuan Video

  Wan2.1-I2V

  Mochi-1

• Support DLRM training with MAD example

⁠v25.9 Release Notes :

####Updates:

• ROCm 7.0 support

• Combined PyTorch Docker supports

  Primus with Megatron-LM and Torchtitan backend - v0.3.0

  Primus Turbo - v0.1.1

  Finetuning: HF PEFT, Torchtune

• Added/updated Diffusion Benchmark support for following models

  Flux with smaller testing dataset

  Support MI300X, MI325X, MI350X, MI355X

• Known Issues:

  Separate images for MI300X/MI325X and MI350X/MI355X due to cross architecture compatibility issues, to be fixed in v25.10 release

⁠v25.8 Release Notes :

####Updates:

• Library updates:

  ROCm: 6.4.3

  PyTorch: 2.8.0a0+gitd06a406

  Transformer Engine: 2.2.0.dev0+a1e66aae

  hipBLASLt 1.1.0-d1b517fc7a

• Primus Front End of Torchtitan (Llama (Dense) only)

• Primus Turbo Support (Partial features)

  Attention

  Grouped GEMM

• Known Issues:

  Memory access fault for torchtune Llama-3.2-vision 90B and Llama4-scott LoRA. Investigating and will update patches on this release given this is the last ROCm 6.4 based release.

⁠v25.7 Release Notes :

####Updates:

• Library updates:

  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

• Increased FP8 support coverage

  Torchtune Full Weight finetuning

  Torchtitan Pre-training

• Added support of Qwen models for finetuning

  Qwen2 - 1.5B, 7B

  Qwen2.5 - 32B, 72B

  Qwen3 - 8B, 32B

• Support of GPT-OSS model for finetuning

• Fixed SemiAnalysis Model issue caused by PyTorch updates

⁠v25.6 Release Note:

⁠Updates:

• Library updates:

  PyTorch 2.8.0a0+git7d205b2

  Triton 3.3.0

  Hipblaslt 0.15.0-8c69191d

  Transformer Engine 1.14.0+2f85f5f2

• Out-of-box full support upstream of pytorch/TorchTune:

• Added Support matrix that match up to date torchtune

• Added benchmarking examples for all the supported models and usage

• Added unified Torchtune_Tester.sh script for torchtune finetuning benchmarking

• New model support with benchmarking example:

  Llama 4 17B_16E (scout): SFT full weight, LoRA

  Llama 3.2 Vision 11B: SFT full weight only

  Llama 3.2 Vision 90B: SFT full weight only

• Out-of-box support of Torchtitan

  Updated torchtitan to match the upstream

  Full BF16 out-of-box support with CK based flash attention

⁠Known Issues:

• Stayed with ROCm 6.3.4 due to ROCm 6.4 runtime issue, workaround identified, will apply to next release

• Known upstream torchtune issue for LoRA finetuning with Llama 3.2 vision models

• Memory leak with ROCm 6.3 (to be fixed in ROCm 6.4)

• Removed benchmarking support of Semi Analysis models used in Dec. 2024 as the implementation is no longer supported with PyTorch 2.8+

⁠v25.5 Release Note:

⁠Updates:

• Library updates:

  ROCm 6.3.4

  Triton 3.2.0

  Hipblaslt 0.13.0-98e224a3

  Transformer Engine 1.12.0.dev0+25a33da

• Full support of TorchTune:

  Model supports: Llama-3.3, Llama-3.2, Llama 3.1, Llama 2

  Feature supports: Full Finetuning, LoRA, qLoRA

• ~38% performance improvement on torchtune packed input setting

⁠Known Issues:

• Numerical Issue reported with SPDA memory efficient backend
• Memory leak reported with ROCm 6.3 (to be fixed in ROCm 6.4)