Building from source¶

This page describes how to install AIMET from source in a uv environment and within docker container.

You can also use a virtual environment (venv), provided your system has the required Python version and necessary dependencies that aren’t available via pip, such as CUDA and cuDNN.

UV environment¶

Install uv¶

Following https://docs.astral.sh/uv/getting-started/installation/ to isntall UV

On Linux/MacOS, you can run the following command to install UV:

curl -LsSf https://astral.sh/uv/install.sh | sh

Create a new uv environment with Python 3.10¶

An example of conda environment setup is shown below:

# Create new uv environment with Python 3.10
uv venv --python=3.10 aimet-dev

# Activate the environment
. aimet-dev/bin/activate

NVIDIA CUDA support¶

Skip the following step, if you don’t want to compile with CUDA support or already have CUDA installed.

Here, we show how to install CUDA Toolkit 12.1 on Ubuntu 22.04. You can find instructions for other versions and platforms in NVIDIA’s documentation: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html

Set environment variables to build desired AIMET wheel¶

General Toggles

GPU build: -DENABLE_CUDA=ON
CPU-only build: -DENABLE_CUDA=OFF
Build C++ tests: -DENABLE_TESTS=ON
Skip building C++ tests: -DENABLE_TESTS=OFF

Variant-specific Toggles

Variant	CMake flags
aimet-onnx	-DENABLE_ONNX=ON -DENABLE_TORCH=OFF
aimet-torch	-DENABLE_TORCH=ON -DENABLE_ONNX=OFF
Docs	-DENABLE_ONNX=ON -DENABLE_TORCH=ON -DENABLE_CUDA=OFF

# Example: Build for aimet-onnx with GPU
export 'CMAKE_ARGS=-DENABLE_CUDA=ON -DENABLE_ONNX=ON -DENABLE_TORCH=OFF -DENABLE_TESTS=OFF'
export 'SKBUILD_BUILD_TARGETS=all'

Compile and install pip package dependencies¶

# cd to AIMET root directory
cd aimet/

# Compile requirements from pyproject.toml with constraints
uv pip compile pyproject.toml --extra=dev --extra=test --output-file=/tmp/requirements.txt

# Install the compiled dependencies
uv pip install -r /tmp/requirements.txt

Build AIMET wheel and run unit tests¶

# Build AIMET wheel
python3 -m build --wheel --no-isolation .

# Install the built wheel
pip install dist/aimet*.whl

# Run unit tests (ONNX)
cd TrainingExtensions/onnx/test/python
pytest

Build AIMET documentation¶

# cd to AIMET root directory
cd aimet/

# Example: Build for Documentation Only
export 'CMAKE_ARGS=-DENABLE_ONNX=ON -DENABLE_TORCH=ON -DENABLE_CUDA=OFF -DENABLE_TESTS=OFF'
export 'SKBUILD_BUILD_TARGETS=all;doc'

# Pin torch, onnxruntime versions
echo "onnxruntime==1.22.0" >> /tmp/constraints.txt
echo "torch==2.1.2" >> /tmp/constraints.txt

# Compile requirements from pyproject.toml with constraints
uv pip compile pyproject.toml -v --constraint=/tmp/constraints.txt --extra=dev --extra=test,docs --output-file=/tmp/requirements.txt

# Install the compiled dependencies
python3 -m pip install -r /tmp/requirements.txt

# Build AIMET docs (aimet/build/Docs/index.html)
python3 -m build --wheel --no-isolation .

Docker environment¶

Build and run docker container locally¶

Docker build argument examples for AIMET Variants.

Variant	Build args
aimet-onnx	VER_PYTHON=3.10 VER_ONNXRUNTIME=1.22.0 VER_CUDA=12.1.0
aimet-torch	VER_PYTHON=3.10 VER_TORCH=2.1.2 VER_CUDA=12.1.1

# cd to AIMET root directory
cd aimet

# Example: Build docker image for aimet-onnx with GPU
docker buildx build --build-arg VER_PYTHON=3.10 --build-arg VER_ONNXRUNTIME=1.22.0 --build-arg VER_CUDA=12.1.0 -t onnx-gpu:1.0 -f Jenkins/fast-release/Dockerfile.ci .

# Run the container
docker run -it -v /local/mnt/workspace:/local/mnt/workspace/ --gpus all --user root onnx-gpu:1.0

# Set up the conda environment inside the container
. ${VIRTUAL_ENV}/bin/activate

Set environment variables to build desired AIMET wheel¶

General Toggles

GPU build: -DENABLE_CUDA=ON
CPU-only build: -DENABLE_CUDA=OFF
Build C++ tests: -DENABLE_TESTS=ON
Skip build C++ tests: -DENABLE_TESTS=OFF

Variant-specific Toggles

Variant	CMake flags
aimet-onnx	-DENABLE_ONNX=ON -DENABLE_TORCH=OFF
aimet-torch	-DENABLE_TORCH=ON -DENABLE_ONNX=OFF

# Example: Build for aimet-onnx with GPU
export 'CMAKE_ARGS=-DENABLE_CUDA=ON -DENABLE_ONNX=ON -DENABLE_TORCH=OFF -DENABLE_TESTS=OFF'
export 'SKBUILD_BUILD_TARGETS=all'

Build AIMET wheel and run unit tests¶

# Build AIMET wheel
python3 -m build --wheel --no-isolation .

# Install the built wheel
pip install dist/aimet*.whl

# Run unit tests (ONNX)
cd TrainingExtensions/onnx/test/python/
pytest