Installing deep learning frameworks for ROCm

2025-09-11

3 min read time

Applies to Linux

ROCm provides a comprehensive ecosystem for deep learning development, including libraries for optimized deep learning operations and ROCm-aware versions of popular deep learning frameworks and libraries such as PyTorch, TensorFlow, and JAX. ROCm works closely with these frameworks to ensure that framework-specific optimizations take advantage of AMD accelerator and GPU architectures.

The following guides provide information on compatibility and supported features for these ROCm-enabled deep learning frameworks.

• PyTorch compatibility

• TensorFlow compatibility

• JAX compatibility

• verl compatibility

• Stanford Megatron-LM compatibility

• DGL compatibility

• Megablocks compatibility

• Taichi compatibility

This chart steps through typical installation workflows for installing deep learning frameworks for ROCm.

See the installation instructions to get started.

• PyTorch for ROCm

• TensorFlow for ROCm

• JAX for ROCm

• verl for ROCm

• Stanford Megatron-LM for ROCm

• DGL for ROCm

• Megablocks for ROCm

• Taichi for ROCm

Note

For guidance on installing ROCm itself, refer to ROCm installation for Linux.

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

• Training

• Fine-tuning LLMs

• Inference

• Inference optimization