@@ -0,0 +1,951 @@
.. meta ::
:description: PyTorch compatibility
:keywords: GPU, PyTorch compatibility
********************************************************************************
PyTorch compatibility
********************************************************************************
`PyTorch <https://pytorch.org/> `_ is an open-source tensor library designed for
`MIOpen <https://github.com/ROCm/MIOpen> `_ and
`RCCL <https://github.com/ROCm/rccl> `_ libraries.
- ROCm PyTorch release:
- Offers :ref: `Docker images <pytorch-docker-compat>` with ROCm and PyTorch
- ROCm PyTorch repository: `<https://github.com/rocm/pytorch>`__
- See the :doc: `ROCm PyTorch installation guide <rocm-install-on-linux:install/3rd-party/pytorch-install>` to get started.
- Official PyTorch release:
`<https://github.com/pytorch/pytorch>`__
`Nightly and latest stable version installation guide <https://pytorch.org/get-started/locally/> `_
`Previous versions <https://pytorch.org/get-started/previous-versions/> `_ to get started.
.. _pytorch-docker-compat:
Docker image compatibility
================================================================================
`PyTorch <https://hub.docker.com/r/rocm/pytorch> `_
`ROCm 6.3.0 <https://repo.radeon.com/rocm/apt/6.3/> `_ .
.. list-table :: PyTorch Docker image components
:header-rows: 1
:class: docker-image-compatibility
* - Docker
- PyTorch
- Ubuntu
- Python
- Apex
- torchvision
- TensorBoard
- MAGMA
- UCX
- OMPI
- OFED
* - .. raw:: html
`2.4.0 <https://github.com/ROCm/pytorch/tree/release/2.4> `_
- 24.04
- `3.12 <https://www.python.org/downloads/release/python-3128/> `_
- `1.4.0 <https://github.com/ROCm/apex/tree/release/1.4.0> `_
- `0.19.0 <https://github.com/pytorch/vision/tree/v0.19.0> `_
- `2.13.0 <https://github.com/tensorflow/tensorboard/tree/2.13> `_
- `master <https://bitbucket.org/icl/magma/src/master/> `_
- `1.10.0 <https://github.com/openucx/ucx/tree/v1.10.0> `_
- `4.0.7 <https://github.com/open-mpi/ompi/tree/v4.0.7> `_
- `5.3-1.0.5.0 <https://content.mellanox.com/ofed/MLNX_OFED-5.3-1.0.5.0/MLNX_OFED_LINUX-5.3-1.0.5.0-ubuntu20.04-x86_64.tgz> `_
* - .. raw:: html
`2.4.0 <https://github.com/ROCm/pytorch/tree/release/2.4> `_
- 22.04
- `3.10 <https://www.python.org/downloads/release/python-31016/> `_
- `1.4.0 <https://github.com/ROCm/apex/tree/release/1.4.0> `_
- `0.19.0 <https://github.com/pytorch/vision/tree/v0.19.0> `_
- `2.13.0 <https://github.com/tensorflow/tensorboard/tree/2.13> `_
- `master <https://bitbucket.org/icl/magma/src/master/> `_
- `1.10.0 <https://github.com/openucx/ucx/tree/v1.10.0> `_
- `4.0.7 <https://github.com/open-mpi/ompi/tree/v4.0.7> `_
- `5.3-1.0.5.0 <https://content.mellanox.com/ofed/MLNX_OFED-5.3-1.0.5.0/MLNX_OFED_LINUX-5.3-1.0.5.0-ubuntu20.04-x86_64.tgz> `_
* - .. raw:: html
`2.4.0 <https://github.com/ROCm/pytorch/tree/release/2.4> `_
- 22.04
- `3.9 <https://www.python.org/downloads/release/python-3918/> `_
- `1.4.0 <https://github.com/ROCm/apex/tree/release/1.4.0> `_
- `0.19.0 <https://github.com/pytorch/vision/tree/v0.19.0> `_
- `2.13.0 <https://github.com/tensorflow/tensorboard/tree/2.13> `_
- `master <https://bitbucket.org/icl/magma/src/master/> `_
- `1.10.0 <https://github.com/openucx/ucx/tree/v1.10.0> `_
- `4.0.7 <https://github.com/open-mpi/ompi/tree/v4.0.7> `_
- `5.3-1.0.5.0 <https://content.mellanox.com/ofed/MLNX_OFED-5.3-1.0.5.0/MLNX_OFED_LINUX-5.3-1.0.5.0-ubuntu20.04-x86_64.tgz> `_
* - .. raw:: html
`2.3.0 <https://github.com/ROCm/pytorch/tree/release/2.3> `_
- 22.04
- `3.10 <https://www.python.org/downloads/release/python-31016/> `_
- `1.3.0 <https://github.com/ROCm/apex/tree/release/1.3.0> `_
- `0.18.0 <https://github.com/pytorch/vision/tree/v0.18.0> `_
- `2.13.0 <https://github.com/tensorflow/tensorboard/tree/2.13> `_
- `master <https://bitbucket.org/icl/magma/src/master/> `_
- `1.14.1 <https://github.com/openucx/ucx/tree/v1.14.1> `_
- `4.1.5 <https://github.com/open-mpi/ompi/tree/v4.1.5> `_
- `5.3-1.0.5.0 <https://content.mellanox.com/ofed/MLNX_OFED-5.3-1.0.5.0/MLNX_OFED_LINUX-5.3-1.0.5.0-ubuntu20.04-x86_64.tgz> `_
* - .. raw:: html
`2.2.1 <https://github.com/ROCm/pytorch/tree/release/2.2> `_
- 22.04
- `3.10 <https://www.python.org/downloads/release/python-31016/> `_
- `1.2.0 <https://github.com/ROCm/apex/tree/release/1.2.0> `_
- `0.17.1 <https://github.com/pytorch/vision/tree/v0.17.1> `_
- `2.13.0 <https://github.com/tensorflow/tensorboard/tree/2.13> `_
- `master <https://bitbucket.org/icl/magma/src/master/> `_
- `1.14.1 <https://github.com/openucx/ucx/tree/v1.14.1> `_
- `4.1.5 <https://github.com/open-mpi/ompi/tree/v4.1.5> `_
- `5.3-1.0.5.0 <https://content.mellanox.com/ofed/MLNX_OFED-5.3-1.0.5.0/MLNX_OFED_LINUX-5.3-1.0.5.0-ubuntu20.04-x86_64.tgz> `_
* - .. raw:: html
`2.2.1 <https://github.com/ROCm/pytorch/tree/release/2.2> `_
- 20.04
- `3.9 <https://www.python.org/downloads/release/python-3921/> `_
- `1.2.0 <https://github.com/ROCm/apex/tree/release/1.2.0> `_
- `0.17.1 <https://github.com/pytorch/vision/tree/v0.17.1> `_
- `2.13.0 <https://github.com/tensorflow/tensorboard/tree/2.13.0> `_
- `master <https://bitbucket.org/icl/magma/src/master/> `_
- `1.10.0 <https://github.com/openucx/ucx/tree/v1.10.0> `_
- `4.0.3 <https://github.com/open-mpi/ompi/tree/v4.0.3> `_
- `5.3-1.0.5.0 <https://content.mellanox.com/ofed/MLNX_OFED-5.3-1.0.5.0/MLNX_OFED_LINUX-5.3-1.0.5.0-ubuntu20.04-x86_64.tgz> `_
* - .. raw:: html
`1.13.1 <https://github.com/ROCm/pytorch/tree/release/1.13> `_
- 22.04
- `3.9 <https://www.python.org/downloads/release/python-3921/> `_
- `1.0.0 <https://github.com/ROCm/apex/tree/release/1.0.0> `_
- `0.14.0 <https://github.com/pytorch/vision/tree/v0.14.0> `_
- `2.18.0 <https://github.com/tensorflow/tensorboard/tree/2.18> `_
- `master <https://bitbucket.org/icl/magma/src/master/> `_
- `1.14.1 <https://github.com/openucx/ucx/tree/v1.14.1> `_
- `4.1.5 <https://github.com/open-mpi/ompi/tree/v4.1.5> `_
- `5.3-1.0.5.0 <https://content.mellanox.com/ofed/MLNX_OFED-5.3-1.0.5.0/MLNX_OFED_LINUX-5.3-1.0.5.0-ubuntu20.04-x86_64.tgz> `_
* - .. raw:: html
`1.13.1 <https://github.com/ROCm/pytorch/tree/release/1.13> `_
- 20.04
- `3.9 <https://www.python.org/downloads/release/python-3921/> `_
- `1.0.0 <https://github.com/ROCm/apex/tree/release/1.0.0> `_
- `0.14.0 <https://github.com/pytorch/vision/tree/v0.14.0> `_
- `2.18.0 <https://github.com/tensorflow/tensorboard/tree/2.18> `_
- `master <https://bitbucket.org/icl/magma/src/master/> `_
- `1.10.0 <https://github.com/openucx/ucx/tree/v1.10.0> `_
- `4.0.3 <https://github.com/open-mpi/ompi/tree/v4.0.3> `_
- `5.3-1.0.5.0 <https://content.mellanox.com/ofed/MLNX_OFED-5.3-1.0.5.0/MLNX_OFED_LINUX-5.3-1.0.5.0-ubuntu20.04-x86_64.tgz> `_
Critical ROCm libraries for PyTorch
================================================================================
.. list-table ::
:header-rows: 1
* - ROCm library
- Version
- Purpose
- Used in
* - `Composable Kernel <https://github.com/ROCm/composable_kernel> `_
`` torch.permute `` , `` torch.view `` , `` torch.matmul `` ,
`` torch.mm `` , `` torch.bmm `` , `` torch.nn.Conv2d `` , `` torch.nn.Conv3d ``
`` torch.nn.MultiheadAttention `` .
* - `hipBLAS <https://github.com/ROCm/hipBLAS> `_
* - `hipBLASLt <https://github.com/ROCm/hipBLASLt> `_
`` torch.matmul `` , `` torch.mm `` , and the
* - `hipCUB <https://github.com/ROCm/hipCUB> `_
`` torch.sum `` , `` torch.cumsum `` , `` torch.sort ``
`` torch.topk `` . Operations on sparse tensors or tensors with
* - `hipFFT <https://github.com/ROCm/hipFFT> `_
`` torch.fft `` module.
* - `hipRAND <https://github.com/ROCm/hipRAND> `_
`` torch.rand `` , `` torch.randn `` and stochastic layers like
`` torch.nn.Dropout `` .
* - `hipSOLVER <https://github.com/ROCm/hipSOLVER> `_
`` torch.linalg.solve `` ,
`` torch.linalg.eig `` , and `` torch.linalg.svd `` .
* - `hipSPARSE <https://github.com/ROCm/hipSPARSE> `_
`` torch.sparse `` .
* - `hipSPARSELt <https://github.com/ROCm/hipSPARSELt> `_
`` torch.sparse `` .
* - `hipTensor <https://github.com/ROCm/hipTensor> `_
* - `MIOpen <https://github.com/ROCm/MIOpen> `_
`` torch.nn.Conv2d `` , `` torch.nn.ReLU `` , and `` torch.nn.LSTM `` .
* - `MIGraphX <https://github.com/ROCm/AMDMIGraphX> `_
`` torch.nn.Conv2d `` , `` torch.nn.ReLU `` , and `` torch.nn.LSTM `` .
* - `MIVisionX <https://github.com/ROCm/MIVisionX> `_
`` torch.utils.data ``
`` torchvision `` workflows.
* - `rocAL <https://github.com/ROCm/rocAL> `_
`` torch.utils.data `` and
`` torchvision `` data load workloads.
* - `RCCL <https://github.com/ROCm/rccl> `_
`` torch.nn.parallel.DistributedDataParallel `` ).
* - `rocDecode <https://github.com/ROCm/rocDecode> `_
`` torch.utils.data `` , `` torchvision.transforms ``
`` torch.distributed `` .
* - `rocJPEG <https://github.com/ROCm/rocJPEG> `_
`` torchvision.io.decode_jpeg `` and
`` torchvision.io.encode_jpeg `` and can be integrated in
`` torch.utils.data `` and `` torchvision `` .
* - `RPP <https://github.com/ROCm/RPP> `_
`` torch.utils.data `` and
`` torchvision `` data load workloads.
* - `rocThrust <https://github.com/ROCm/rocThrust> `_
* - `rocWMMA <https://github.com/ROCm/rocWMMA> `_
`` torch.nn.Linear `` ), convolutional layers
`` torch.nn.Conv2d `` ), attention layers, general tensor operations that
`` torch.matmul `` , `` torch.bmm `` , and
Supported and unsupported features
================================================================================
torch
--------------------------------------------------------------------------------
`torch <https://pytorch.org/docs/stable/index.html> `_ is the central module of
Tensor data types
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
`` dtype `` attribute or argument, and PyTorch supports a wide range of data types for different use cases.
`torch.Tensor <https://pytorch.org/docs/stable/tensors.html> `_ 's single data types:
.. list-table ::
:header-rows: 1
* - Data type
- Description
- Since PyTorch
- Since ROCm
* - `` torch.float8_e4m3fn ``
* - `` torch.float8_e5m2 ``
* - `` torch.float16 `` or `` torch.half ``
* - `` torch.bfloat16 ``
* - `` torch.float32 `` or `` torch.float ``
* - `` torch.float64 `` or `` torch.double ``
* - `` torch.complex32 `` or `` torch.chalf ``
* - `` torch.complex64 `` or `` torch.cfloat ``
* - `` torch.complex128 `` or `` torch.cdouble ``
* - `` torch.uint8 ``
* - `` torch.uint16 ``
* - `` torch.uint32 ``
* - `` torch.uint64 ``
* - `` torch.int8 ``
* - `` torch.int16 `` or `` torch.short ``
* - `` torch.int32 `` or `` torch.int ``
* - `` torch.int64 `` or `` torch.long ``
* - `` torch.bool ``
* - `` torch.quint8 ``
* - `` torch.qint8 ``
* - `` torch.qint32 ``
* - `` torch.quint4x2 ``
.. note ::
`` uint8 `` are currently only have limited support in
`` torch.compile `` ).
:doc: `ROCm precision support page <rocm:reference/precision-support>`
torch.cuda
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
`` torch.cuda `` in PyTorch is a module that provides utilities and functions for
.. list-table ::
:header-rows: 1
* - Data type
- Description
- Since PyTorch
- Since ROCm
* - Device management
* - Tensor operations on GPU
* - Streams and events
* - Memory management
`` torch.cuda.memory_allocated() `` , `` torch.cuda.max_memory_allocated() `` ,
`` torch.cuda.memory_reserved() `` and `` torch.cuda.empty_cache() `` .
* - Running process lists of memory management
`` torch.cuda.memory_stats() `` and `` torch.cuda.memory_summary() `` .
* - Communication collectives
* - `` torch.cuda.CUDAGraph ``
* - TunableOp
* - NVIDIA Tools Extension (NVTX)
* - Lazy loading NVRTC
* - Jiterator (beta)
.. Need to validate and extend.
torch.backends.cuda
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
`` torch.backends.cuda `` is a PyTorch module that provides configuration options
.. list-table ::
:header-rows: 1
* - Data type
- Description
- Since PyTorch
- Since ROCm
* - `` cufft_plan_cache ``
* - `` matmul.allow_tf32 ``
* - `` matmul.allow_fp16_reduced_precision_reduction ``
* - `` matmul.allow_bf16_reduced_precision_reduction ``
* - `` enable_cudnn_sdp ``
* - `` enable_flash_sdp ``
* - `` enable_mem_efficient_sdp ``
* - `` enable_math_sdp ``
* - `` allow_fp16_bf16_reduction_math_sdp ``
..
FIXME:
- Partial?
.. Need to validate and extend.
torch.backends.cudnn
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
`` torch `` options:
.. list-table ::
:header-rows: 1
* - Data type
- Description
- Since PyTorch
- Since ROCm
* - `` allow_tf32 ``
* - `` deterministic ``
Automatic mixed precision: torch.amp
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.. list-table ::
:header-rows: 1
* - Data type
- Description
- Since PyTorch
- Since ROCm
* - Autocasting
* - Gradient scaling
’
’
* - CUDA op-specific behavior
`` torch.nn.Module `` , as a function, or as a `` torch.Tensor `` method. If
Distributed library features
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
:ref: `rocm-for-ai-pytorch-distributed` for more information.
.. list-table ::
:header-rows: 1
* - Features
- Description
- Since PyTorch
- Since ROCm
* - TensorPipe
* - RPC Device Map Passing
* - Gloo
* - MPI
`` torch.distributed `` module.
* - TorchElastic
..
FIXME: RPC Device Map Passing "Since ROCm version"
torch.compiler
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.. list-table ::
:header-rows: 1
* - Features
- Description
- Since PyTorch
- Since ROCm
* - `` torch.compiler `` (AOT Autograd)
`` TorchInductor `` .
* - `` torch.compiler `` (TorchInductor)
`` torch.compile `` deep learning compiler that generates fast
`` torch.compile `` possible. For AMD,
torchaudio
--------------------------------------------------------------------------------
`torchaudio <https://pytorch.org/audio/stable/index.html> `_ library provides
`` torchaudio.transforms `` , you need to move audio
`` .to('cuda') `` .
`` torchaudio `` features are GPU-accelerated.
.. list-table ::
:header-rows: 1
* - Features
- Description
- Since torchaudio version
- Since ROCm
* - `` torchaudio.transforms.Spectrogram ``
* - `` torchaudio.transforms.MelSpectrogram ``
* - `` torchaudio.transforms.MFCC ``
* - `` torchaudio.transforms.Resample ``
torchvision
--------------------------------------------------------------------------------
`torchvision <https://pytorch.org/vision/stable/index.html> `_ library
`` torchvision `` features are GPU-accelerated.
.. list-table ::
:header-rows: 1
* - Features
- Description
- Since torchvision version
- Since ROCm
* - `` torchvision.transforms.functional ``
* - `` torchvision.ops ``
`` torchvision.ops.roi_align `` , `` torchvision.ops.nms `` and
`` box_convert `` .
* - `` torchvision.models `` with `` .to('cuda') ``
`` torchvision `` provides several pre-trained models (ResNet, Faster
* - `` torchvision.io ``
’
torchtext
--------------------------------------------------------------------------------
`torchtext <https://pytorch.org/text/stable/index.html> `_ library provides
`` .to('cuda') `` .
* torchtext does not implement its own kernels. ROCm support is enabled by linking against ROCm libraries.
* Only official release exists.
torchtune
--------------------------------------------------------------------------------
`torchtune <https://pytorch.org/torchtune/stable/index.html> `_ library for
* Usage: It works out-of-the-box, enabling developers to fine-tune ROCm PyTorch solutions.
* Only official release exists.
torchserve
--------------------------------------------------------------------------------
`torchserve <https://pytorch.org/torchserve/> `_ is a PyTorch domain library
* torchtext does not implement its own kernels. ROCm support is enabled by linking against ROCm libraries.
* Only official release exists.
torchrec
--------------------------------------------------------------------------------
`torchrec <https://pytorch.org/torchrec/> `_ is a PyTorch domain library for
* torchrec does not implement its own kernels. ROCm support is enabled by linking against ROCm libraries.
* Only official release exists.
Unsupported PyTorch features
----------------------------
.. list-table ::
:widths: 30, 60, 10
:header-rows: 1
* - Data type
- Description
- Since PyTorch
* - APEX batch norm
* - `` torch.backends.cuda `` / `` matmul.allow_tf32 ``
* - `` torch.cuda `` / NVIDIA Tools Extension (NVTX)
* - `` torch.cuda `` / Lazy loading NVRTC
* - `` torch-tensorrt ``
* - `` torch.backends `` / `` cudnn.allow_tf32 ``
* - `` torch.backends.cuda `` / `` matmul.allow_fp16_reduced_precision_reduction ``
* - `` torch.backends.cuda `` / `` matmul.allow_bf16_reduced_precision_reduction ``
* - `` torch.nn.functional `` / `` scaled_dot_product_attention ``
* - `` torch.backends.cuda `` / `` enable_cudnn_sdp ``
* - `` torch.backends.cuda `` / `` enable_flash_sdp ``
* - `` torch.backends.cuda `` / `` enable_mem_efficient_sdp ``
* - `` torch.backends.cuda `` / `` enable_math_sdp ``
* - Dynamic parallelism
`` torch.utils.cpp_extension `` module.
* - Unified memory support in PyTorch
Use cases and recommendations
================================================================================
* :doc: `Using ROCm for AI: training a model </how-to/rocm-for-ai/train-a-model>` provides
* :doc: `Single-GPU fine-tuning and inference </how-to/llm-fine-tuning-optimization/single-gpu-fine-tuning-and-inference>`
* :doc: `Multi-GPU fine-tuning and inference optimization </how-to/llm-fine-tuning-optimization/multi-gpu-fine-tuning-and-inference>`
* The :doc: `Instinct MI300X workload optimization guide </how-to/tuning-guides/mi300x/workload>` provides detailed
* The :doc: `Inception with PyTorch documentation </conceptual/ai-pytorch-inception>`
`ROCm PyTorch blog posts <https://rocm.blogs.amd.com/blog/tag/pytorch.html> `_
Peter Park