Deep learning frameworks for ROCm

2026-01-16

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.

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 guide	Installation options	GitHub
PyTorch	link	Docker image Wheels package ROCm Base Docker image Upstream Docker file
TensorFlow	link	Docker image Wheels package
JAX	link	Docker image
verl	link	Docker image
Stanford Megatron-LM	link	Docker image
DGL	link	Docker image
Megablocks	link	Docker image
Ray	link	Docker image Wheels package ROCm Base Docker image
llama.cpp	link	Docker image ROCm Base Docker image
FlashInfer	link	Docker image ROCm Base Docker image

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