Copyright © 2023-2026 Advanced Micro Devices, Inc. All rights reserved.

zentorch: A PyTorch extension for AMD EPYC™ CPUs.

Table of Contents

• About zentorch
  • Overview
  • Structure
  • Third Party Libraries
  • Supported OS
• Installation
  • From PyPI
  • From Source
• Unit Tests
  • Install Unit tests Dependencies
  • Run All Unit Tests
  • Run All Tests
  • Run Individual Tests
• Usage
  • General Usage
  • CNN Models
  • HuggingFace NLP models
  • HuggingFace Generative LLM models
  • Quantized Models
    • Quantizing Models
      • Weight Only Quantization (WOQ)
      • Dynamic Quantization
    • Running Quantized Models
  • vLLM Zentorch Plugin
• Logging and Debugging
  • ZenDNN logs
  • zentorch logs
    • Understanding Profiler Output with Operator Fusion
  • Support for TORCH_COMPILE_DEBUG
• Performance tuning and Benchmarking
• Additional Utilities
  • zentorch attributes

1. About zentorch

1.1. Overview

The latest ZenDNN Plugin for PyTorch* (zentorch) [2.14.1.0] is here

The ZenDNN plugin for PyTorch is called zentorch.

This upgrade continues the focus on optimizing inference with Recommender Systems and Large Language Models on AMD EPYC™ CPUs. It includes AMD EPYC™ enhancements for bfloat16 performance, expanded support for cutting-edge models like Llama 3.2 and 3.3, Microsoft Phi, and more as well as support for a wide-variety of quantization configurations. The quantization support included 4-bit weight-only quantization, along with support for INT8 dynamic activation and INT8 weight quantization, and quantized support for the DLRM-v2 model with a mix of 8-bit and 4-bit quantization. This also includes support for running generative models with vLLM. This release introduces functional support for running LLMs using float16 precision with vLLM on 6th Gen AMD EPYC™ processors.

Under the hood, ZenDNN’s enhanced AMD-specific optimizations operate at every level. In addition to highly optimized operator microkernels, these include comprehensive graph optimizations including pattern identification, graph reordering, and fusions. They also incorporate optimized embedding bag kernels and enhanced zenMatMul matrix splitting strategies which leverage the AMD EPYC™ microarchitecture to deliver enhanced throughput and latency.

Combined with PyTorch's torch.compile, zentorch transforms deep learning pipelines into finely-tuned, AMD-specific engines, delivering unparalleled efficiency and speed for large-scale inference workloads

zentorch supports PyTorch v2.14.1 (recommended) and v2.13.0. Install the zentorch version matching your PyTorch — zentorch 2.14.1.0 for PyTorch v2.14.1 or zentorch 2.13.0.3 for PyTorch v2.13.0, offering a high-performance experience for deep learning on AMD EPYC™ platforms.

Support

We welcome feedback, suggestions, and bug reports. Should you have any of the these, please kindly file an issue on the ZenDNN Plugin for PyTorch GitHub page here

License

AMD copyrighted code in ZenDNN is subject to the Apache-2.0, MIT, or BSD-3-Clause licenses; consult the source code file headers for the applicable license. Third party copyrighted code in ZenDNN is subject to the licenses set forth in the source code file headers of such code.

1.2. Structure

zentorch consists of three parts. They are

• ZenDNN Integration Code
• zentorch backend
• Build System

1.2.1. ZenDNN Integration Code

ZenDNN is integrated into zentorch using CPP code which interfaces ATen API to ZenDNN's API. This code exports torch compatible API similar to ATen IR of PyTorch. The exported API is made available for usage from python code using TORCH_LIBRARY, TORCH_LIBRARY_FRAGMENT and TORCH_LIBRARY_IMPL. Integration code is linked and compiled into zentorch using CppExtension provided by PyTorch.

The following ops are integrated as of now:

• Embedding bag op
• Embedding op
• Matmul ops
• Custom Fusion ops
• Rope op
• MHA op

1.2.2. The zentorch custom backend to torch.compile

We have registered a custom backend to torch.compile called zentorch. This backend integrates ZenDNN optimizations after AOTAutograd through a function called optimize. This function operates on the FX based graph at the ATEN IR to produce an optimized FX based graph as the output.

1.2.3. Build System

The static library for ZenDNN and the cpp Extension modules that bind the ZenDNN operators with Python are built using setup.py script.

1.2.3.1. CMake Based Build

CMake downloads the ZenDNN during configure stage. It generates a config.h with GIT hashes of ZenDNN. It downloads ZenDNN as a static library.

1.2.3.2. Packaging into a Wheel File

The CPP code, being an extension module, is built through CppExtension. It takes various attributes to zentorch. setup.py also adds in various attributes to the zentorch for debugging and providing additional information.

1.3. Third Party Libraries

zentorch uses following library for its functionality.

• ZenDNN

1.4. Supported OS

Refer to the support matrix for the list of supported operating systems.

2. Installation

zentorch can be installed using binary wheel file or can be built from source itself. zentorch supports PyTorch v2.14.1 (recommended) and v2.13.0. Install the zentorch version matching your PyTorch — zentorch v2.14.1.0 for PyTorch v2.14.1 or zentorch v2.13.0.3 for PyTorch v2.13.0.

2.1. From PyPI

• Create conda or python environment and activate it.
• Uninstall any existing zentorch installations.

pip uninstall zentorch

• Install PyTorch v2.14.0 or v2.13.0

pip install torch==2.14.0 --index-url https://download.pytorch.org/whl/cpu

pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cpu

• Install latest stable zentorch v2.14.0.1 or v2.13.0.2

pip install zentorch==2.14.0.1

pip install zentorch==2.13.0.2

Notes:

• Dependent packages 'numpy' and 'torch' will be installed by 'zentorch' if not already present.
• If you get the error: ImportError: /lib64/libstdc++.so.6: version `GLIBCXX_.a.b.cc' not found, export LD_PRELOAD as: export LD_PRELOAD=<path_to_conda>/envs/<env_name>/lib/libstdc++.so.6:$LD_PRELOAD

2.2. From Source

Run the following commands:

git clone https://github.com/amd/ZenDNN-pytorch-plugin.git
cd ZenDNN-pytorch-plugin

Notes:

• export ZENDNNL_MANYLINUX_BUILD=1 is needed for build from source for RHEL/FEDORA/Almalinux/CentOS OS families

2.2.1. Preparing third party repositories

Build setup downloads the ZenDNN repo into third_party folder.

2.2.2. Linux build

2.2.2.1. Create conda environment for the build

conda create -n pt-zentorch python=3.12 -y
conda activate pt-zentorch

2.2.2.2. Install PyTorch

Install PyTorch v2.14.1 (recommended):

pip install torch==2.14.1 --index-url https://download.pytorch.org/whl/cpu

Alternatively, PyTorch v2.13.0 is also supported:

pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cpu

Notes:

• This README uses Python 3.12.
• Zentorch follows PyTorch’s Python version compatibility. For PyTorch 2.14.1 and 2.13.0 Zentorch supports Python versions 3.10 through 3.14. For other PyTorch releases, refer to the PyTorch Release Compatibility Matrix.
• Zentorch does not support experimental versions of Python (3.13T/3.14T)

2.2.2.3. Install Dependencies

pip install -r requirements.txt

2.2.2.4. To build & generate wheel file of zentorch

python setup.py bdist_wheel

Note: The wheel file will be generated in dist folder in ZenDNN-pytorch-plugin directory

2.2.2.5. To install the wheel file of zentorch

pip install dist/zentorch-2.14.1.0-cp312-cp312-linux_x86_64.whl

or

pip install dist/zentorch-2.13.0.3-cp312-cp312-linux_x86_64.whl

Note: After installation, run the following script for recommended environment settings:

source scripts/zentorch_env_setup.sh

Refer to zentorch_env_setup.sh for details.

2.2.2.6. Build Cleanup

python setup.py clean --all

3. Unit Tests

3.1 Install Unit tests Dependencies

python test/install_requirements.py

Note: Before running any unit tests, export the following environment variables to disable ZenDNN caching:

export ZENDNNL_ZP_COMP_CACHE=0
export ZENDNNL_ENABLE_POSTOP_CACHE=0

3.2 Run All Unit Tests

python -m unittest discover -s ./test/unittests

3.3 Run All Tests

python -m unittest discover -s ./test

3.4 Run Individual Tests

python -m unittest test/unittests/op_tests/test_bmm.py

4. Usage

4.1 General Usage

# Using torch.compile with 'zentorch' as backend
import torch
import zentorch
--snip--
compiled_model = torch.compile(model, backend='zentorch')
with torch.no_grad():
    output = compiled_model(input)

4.1.1 Using freeze path with zentorch

Additionally, zentorch supports freezing the model. Enable it before running:

export TORCHINDUCTOR_FREEZING=1

import torch
import zentorch
--snip--
compiled_model = torch.compile(model, backend='zentorch')
with torch.no_grad():
    output = compiled_model(input)

Note:

• zentorch is able to do the zentorch op replacements in both non-inference and inference modes. But some of the zentorch optimizations are only supported for the inference mode, so it is recommended to use torch.no_grad() if you are running the model for inference only.

4.2 CNN Models

For CNN models, set dynamic=False when calling for torch.compile as below:

model = torch.compile(model, backend='zentorch', dynamic=False)
with torch.no_grad():
    output = model(input)

4.3 HuggingFace NLP models

For HuggingFace NLP models, optimize them as below:

model = torch.compile(model, backend='zentorch')
with torch.no_grad():
    output = model(input)

4.4 HuggingFace Generative LLM models

HuggingFace generative LLM models can be run with torch.compile(model, backend='zentorch'). For optimal performance, use the vLLM zentorch plugin. Refer to vLLM Zentorch Plugin.

4.5 Quantized Models

zentorch supports quantized inference for large language models on AMD EPYC™ CPUs. The following quantization methods are supported:

• Weight Only Quantization (WOQ): INT4 quantized weights with both symmetric and asymmetric configurations, providing significant memory savings while maintaining model accuracy.
• Dynamic Quantization: INT8 weights with INT8 dynamic activations, quantizing activations dynamically at runtime for minimal accuracy loss.

Models can be quantized with LLM Compressor. TorchAO quantization support is deprecated and will be removed in the next release.

4.5.1 Quantizing Models

Models quantized with LLM-Compressor for this release are published in the AMD EPYC compatible models collection from Red Hat AI on the Hugging Face Hub. See the ZenDNN v6.1 Consolidated Model Support Matrix for the supported models and their quantized checkpoints. Each AMD model ID follows a consistent pattern: base model, quantization scheme, optional symmetric or asymmetric grouping, quantizer, and optional tool version. For example, w4a16-llmcompressor denotes 4-bit weight-only quantization produced with LLM-Compressor, and w4a16-asym-torchao-v0.17.0 denotes 4-bit weight-only quantization with asymmetric per-group weights, produced with TorchAO v0.17.0. Red Hat AI checkpoints end in quantized.w8a8 or quantized.w4a16. The checkpoints below cover the dynamic and weight-only formats supported in section 4.5:

• amd/gpt-oss-20b-BF16-w8a8-llmcompressor — DA8W8 (w8a8) dynamic quantization with LLM-Compressor: INT8 per-channel weights and dynamic INT8 per-token activations.
• amd/Llama-3.1-8B-Instruct-w4a16-llmcompressor — symmetric per-group w4a16 (group size 128) weight-only configuration with LLM-Compressor. The matrix lists it under DA8W4 and W4A16 (WOQ-Sym-PG).
• amd/Llama-3.1-8B-Instruct-w4a16-asym-torchao-v0.17.0 — asymmetric per-group w4a16 (W4A16 (WOQ-Asym)) with TorchAO v0.17.0. TorchAO support is deprecated.

Note: The quantization recipe and evaluation details for each checkpoint are documented on its Hugging Face model card. TorchAO checkpoints remain in the zentorch TorchAO collection for PyTorch 2.11 until that support is removed in the next release.

4.5.2 Running Quantized Models

Quantized models can be run using vLLM with the zentorch backend. Refer to Quantization examples for detailed instructions on running models with vLLM.

Note: When running quantized models, use the path to your quantized model directory (e.g., ./quantized_model) as the model path instead of the original HuggingFace model name.

4.6 vLLM Zentorch Plugin

The vLLM-ZenTorch plugin enhances the capabilities of the vLLM inference engine, enabling plug-and-play acceleration of large language model inference on AMD EPYC™ CPUs. By incorporating ZenTorch with vLLM, users can experience substantial throughput enhancements for LLM workloads without requiring any modifications to their existing code.

For more details regarding vLLM-ZenTorch Plugin refer to this Readme.

5. Logging and Debugging

5.1 ZenDNN logs

Logging for ZenDNN is disabled by default but can be enabled by using the environment variable ZENDNN_API_LOG_LEVEL before running any tests. An example to turn on info logging is given below.

export ZENDNNL_API_LOG_LEVEL=4

To enable the profiling logs, you can use:

export ZENDNNL_ENABLE_PROFILER=1

For further details on ZenDNN logging mechanism, refer to ZenDNN user-guide from this page.

5.2 zentorch logs

For zentorch, CPP specific logging can be enabled by setting the environment variable TORCH_CPP_LOG_LEVEL. This has four levels: INFO, WARNING, ERROR and FATAL in decreasing order of verbosity. Similarly, python logging can be enabled by setting the environment variable ZENTORCH_PY_LOG_LEVEL, this has five levels: DEBUG, INFO, WARNING, ERROR and CRITICAL, again in decreasing order of verbosity. An example to enable INFO level logs for cpp and DEBUG level for python (most verbose) is given below:

export TORCH_CPP_LOG_LEVEL=INFO
export ZENTORCH_PY_LOG_LEVEL=DEBUG

The default level of logs is WARNING for both cpp and python sources but can be overridden as discussed above.

Note: The log levels are the same as those provided by the python logging module.

Info: Since all OPs implemented in zentorch are registered with torch using the TORCH_LIBRARY(), TORCH_LIBRARY_FRAGMENT() and TORCH_LIBRARY_IMPL() macros in bindings, the PyTorch profiler can be used without any modifications to measure the op level performance.

5.2.1 Understanding Profiler Output with Operator Fusion

When profiling models optimized with zentorch, the profiler output reflects how zentorch tracks operator fusion using a hierarchical structure.

The hierarchical structure of an umbrella function like zentorch_linear_unary looks like:

zentorch::zentorch_linear_unary
  ├── zentorch::zentorch_linear
  ├── zentorch::zentorch_linear_relu
  ├── zentorch::zentorch_linear_gelu_erf
  ├── zentorch::zentorch_linear_tanh
  ├── ...

zentorch uses umbrella functions to track fused operations, with internal record_function calls for each fusion pattern. In the profiler output:

• Outer/umbrella function — records the overall fused operation execution time
• Inner record functions — track individual fusion components for debugging

As a result, "Self CPU %" is typically small for umbrella functions (time in the function itself), while "CPU total %" is larger (includes all child operations within the fusion).

For example, in the snippet below, zentorch::zentorch_linear_unary is the umbrella function with a low "Self CPU" of 134.041us but a high "CPU total" of 22.196ms, while zentorch::zentorch_linear_tanh is one of its inner fused operations:

Name                                 Self CPU %    Self CPU    CPU total %    CPU total
---------------------------------------------------------------------------------------
zentorch::zentorch_linear_unary         0.29%      134.041us      47.74%      22.196ms
zentorch::zentorch_linear_tanh          0.07%       33.491us       0.20%      92.271us

5.3 Support for TORCH_COMPILE_DEBUG

PyTorch offers a debugging toolbox that comprises a built-in stats and trace function. This functionality facilitates the display of the time spent by each compilation phase, output code, output graph visualization, and IR dump. TORCH_COMPILE_DEBUG invokes this debugging tool that allows for better problem-solving while troubleshooting the internal issues of TorchDynamo and TorchInductor. This functionality works for the models optimized using zentorch, so it can be leveraged to debug these models as well. To enable this functionality, users can either set the environment variable TORCH_COMPILE_DEBUG=1 or specify the environment variable with the runnable file (e.g., test.py) as input.

# test.py contains model optimized by torch.compile with 'zentorch' as backend
TORCH_COMPILE_DEBUG=1 python test.py

For more information about TORCH_COMPILE_DEBUG refer to the official PyTorch documentation available.

6. Performance tuning and Benchmarking

For zentorch v2.14.1.0 and zentorch 2.13.0.3 plugin. Please see the Tuning Guidelines section of ZenDNN User Guide for performance tuning. ZenDNN User Guide can be downloaded from here

7. Additional Utilities:

7.1 zentorch attributes:

To check the version of zentorch use the following command:

python -c 'import zentorch; print(zentorch.__version__)'

To check the build config of zentorch use the following command:

python -c 'import zentorch; print(*zentorch.__config__.split("\n"), sep="\n")'