Deep learning frameworks for ROCm

2025-08-22

3 min read time

Applies to Linux

Deep learning frameworks provide environments for machine learning, training, fine-tuning, inference, and performance optimization.

ROCm offers a complete ecosystem for developing and running deep learning applications efficiently. It also provides ROCm-compatible versions of popular frameworks and libraries, such as PyTorch, TensorFlow, JAX, and others.

The AMD ROCm organization actively contributes to open-source development and collaborates closely with framework organizations. This collaboration ensures that framework-specific optimizations effectively leverage AMD GPUs and accelerators.

The table below summarizes information about ROCm-enabled deep learning frameworks. It includes details on ROCm compatibility and third-party tool support, installation steps and options, and links to GitHub resources. For a complete list of supported framework versions on ROCm, see the Compatibility matrix topic.

Framework	Installation	Installation options	GitHub
PyTorch		Docker image Wheels package ROCm Base Docker image Upstream Docker file
TensorFlow		Docker image Wheels package
JAX		Docker image
verl		Docker image
Stanford Megatron-LM		Docker image
DGL		Docker image
Megablocks		Docker image
Taichi		Docker image Wheels package

Learn how to use your ROCm deep learning environment for training, fine-tuning, inference, and performance optimization through the following guides.

• Use ROCm for AI

• Use ROCm for training

• Use ROCm for fine-tuning LLMs

• Use ROCm for AI inference

• Use ROCm for AI inference optimization