Running the OpenADMET stack on AMD GPUs with ROCm

OpenADMET

— 5 min read
Share

By Sean Colby

This guide walks through setting up an AMD GPU workstation to run openadmet-models, the open-source ADMET property prediction platform. It covers OS installation, AMD ROCm driver setup (including Secure Boot compatibility), and the full openadmet-models environment. We include dual-boot instructions since practitioners with a consumer AMD GPU are likely starting from a Windows machine. A Linux-only setup is easy to derive from the steps below.

Tested on: Xubuntu 24.04, ROCm 6.4.x, AMD Radeon RX 7900 XTX

System setup

Creating installation media

Download the latest Xubuntu 24.04 LTS ISO or other distribution of choice.
Use Rufus to write it to a USB drive with the following settings:

  • Partition scheme: GPT
  • Target system: UEFI

Installing Xubuntu

The Xubuntu 24.04 installer (Subiquity backend) is fragile when modifying partition tables interactively. The workaround is to partition the drive first using GParted, then point the installer at the existing layout.

Step 1: Partition with GParted

  1. Boot from the USB drive and open GParted from the application menu.

  2. Select your target drive from the dropdown in the top-right corner (e.g., /dev/nvme1n1). Confirm it is not your Windows drive.

  3. Go to Device > Create Partition Table, select GPT, and click Apply. This removes any metadata that might crash the installer.

  4. Create two partitions:

    Partition 1 -- EFI Boot Loader:

    • Right-click unallocated space > New
    • Size: 1024 MiB
    • File system: fat32
    • Label: EFI
    • Click Add

    Partition 2 -- Root Filesystem:

    • Right-click remaining space > New
    • Size: remaining
    • File system: ext4
    • Label: ROOT
    • Click Add

  5. Click the checkmark icon (Apply All Operations) and wait for it to finish.

  6. Right-click the 1 GB EFI partition > Manage Flags. Check boot and esp. Close GParted.

Step 2: Run the installer

  1. Open the Install Xubuntu icon on the desktop.

  2. Proceed to the Installation Type screen and select "Something Else" (manual partitioning).

  3. Configure the two partitions you just created:

    EFI partition (p1):

    • Use as: FAT32 file system
    • Mount point: /boot/efi

    Root partition (p2):

    • Use as: Ext4 journaling file system
    • Mount point: /
    • Check Format

  4. Verify the bootloader device dropdown at the bottom shows your Linux drive (e.g., /dev/nvme1n1), not the Windows drive.

  5. Click Install Now.

Note: This guide skips drive encryption. Consider enabling it if the machine is an on-the-go laptop or stores regulated data such as HIPAA patient records or proprietary molecules.

First-boot setup

When the installer finishes, click Restart Now. Remove the USB when prompted and press Enter. At the GRUB menu, select Ubuntu (or Xubuntu), then log in and verify network connectivity. Open a terminal (Ctrl+Alt+T) and run:

# Update the system
sudo apt update && sudo apt upgrade -y

# Build tools and Python support
# Required for driver compilation and Python C extensions
sudo apt install -y build-essential libstdc++-12-dev python3-venv python3-pip \
  python3-dev git wget libsuitesparse-dev ninja-build

# Install a browser
sudo apt install -y firefox

Set up direnv to manage per-project environment variables automatically:

sudo apt install -y direnv
echo 'eval "$(direnv hook bash)"' >> ~/.bashrc
source ~/.bashrc

Installing AMD ROCm drivers

First, identify which ROCm version the latest PyTorch wheels target by checking pytorch.org/get-started/locally. As of this writing that is ROCm 6.2, but AMD driver releases 6.4.x are backward-compatible and recommended.

Find the corresponding driver package from AMD's release history. The following installs 6.4.3:

cd /tmp/

# Download the installer package
wget https://repo.radeon.com/amdgpu-install/6.4.3/ubuntu/noble/amdgpu-install_6.4.60403-1_all.deb

# Install it
sudo apt install ./amdgpu-install_6.4.60403-1_all.deb

# Run the driver installation
sudo amdgpu-install -y --usecase=graphics,rocm

Note: apt install ./file.deb runs under a restricted user (_apt) that cannot access your home directory. Running from /tmp/ avoids permission errors.

Enrolling the MOK key (Secure Boot)

Since Secure Boot is enabled, the AMD kernel module must be signed with a trusted key. During the install, if prompted for a password, enter one (e.g., 12345678) and note it down.

After installation, reboot:

sudo reboot

Watch the monitor before the OS loads. You should see a blue screen titled "Shim UEFI key management":

  1. Select Enroll MOK
  2. Select Continue
  3. Select Yes
  4. Enter the password you set (nothing appears on screen while typing)
  5. Select Reboot

If you miss this screen, re-trigger MOK enrollment manually:

# Use whichever file exists on your system
sudo mokutil --import /var/lib/dkms/mok.pub
# or
sudo mokutil --import /var/lib/shim-signed/mok/MOK.der

Enter and confirm a password, reboot, and follow the same UEFI key management steps above.

Add user to GPU groups

sudo usermod -a -G render,video $LOGNAME

Log out and back in for the group changes to take effect.

Verify the driver install

# Show GPU status
rocm-smi

# Confirm the GPU is detected correctly
/opt/rocm/bin/rocminfo | grep -E 'Name:|Marketing'

You should see your GPU listed (e.g., Radeon RX 7900 XTX).

Installing the OpenADMET stack

openadmet-models provides open-source ADMET property prediction models and tooling, built on PyTorch and a broad scientific Python stack.

Set up environment

ROCm-compatible PyTorch wheels are not available on conda-forge, so we use pip with venv for environment management. direnv handles automatic activation when you enter the project directory.

Create a base directory for your OpenADMET work and initialize a virtual environment:

mkdir -p ~/projects/openadmet/
cd ~/projects/openadmet/

python3 -m venv .venv

# Create a direnv config to auto-activate it
echo 'source .venv/bin/activate' > .envrc
direnv allow

Your prompt should reflect the activated environment whenever you are inside this directory.

Install dependencies

# Ensure pip is up to date
pip install --upgrade pip

# Build dependencies (required for compiling PyG extensions)
pip install ninja setuptools wheel

# Install PyTorch 2.5.1 with ROCm 6.2 wheels
# This release has the broadest compatibility with dependencies
pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 \
  --index-url https://download.pytorch.org/whl/rocm6.2

Before compiling GPU extensions, set environment variables to target your GPU architecture. The values below are for consumer RDNA3 hardware (RX 7000 series):

export HSA_OVERRIDE_GFX_VERSION=11.0.0
export PYTORCH_ROCM_ARCH=gfx1100
export ROCM_HOME=/opt/rocm
export FORCE_CUDA=1
export MAX_JOBS=4

Note: gfx1100 is the architecture identifier for RDNA3 (RX 7000 series). If you are using a different GPU, look up the correct gfx code for your architecture.

# Core ML libraries
pip install pytorch-lightning "chemprop>=2.1.1"

# Scientific and cheminformatics stack
pip install ase catboost codecov datamol h5py hyperopt \
  intake intake-parquet joblib jupyterlab lightgbm loguru matplotlib \
  MDAnalysis modal mordredcommunity numpy pandas pytest pytest-cov pytest-xdist \
  rdkit rich scikit-learn seaborn splito statsmodels tabulate umap-learn \
  wandb xgboost "zarr>=3.0.0" email-validator reportlab uncertainty-toolbox

# boto3/botocore must be pinned to avoid version conflicts with molfeat
pip install "botocore==1.42.19" "boto3==1.42.19"

# OpenADMET and upstream git dependencies
# Installed last so they build against the complete environment
pip install phx-class-registry
pip install tabpfn-extensions
pip install git+https://github.com/PriorLabs/TabPFN.git@main
pip install git+https://github.com/JacksonBurns/neural-pairwise-regression.git
pip install git+https://github.com/OpenADMET/molfeat.git@main
pip install git+https://github.com/PatWalters/useful_rdkit_utils.git@master

For users running models:

pip install git+https://github.com/OpenADMET/openadmet-models.git

For contributors developing the library:

git clone git@github.com:OpenADMET/openadmet-models.git
cd openadmet-models
pip install -e .

Verify

Check for dependency conflicts:

pip check

Confirm the full stack is functional:

python3 -c "
import torch
import pytorch_lightning
import rdkit
import chemprop

print(f'PyTorch Version: {torch.__version__}')
print(f'ROCm Available:  {torch.cuda.is_available()}')
print(f'GPU Name:        {torch.cuda.get_device_name(0)}')
print('ALL SYSTEMS GO')
"

Save the working environment:

pip freeze > requirements_openadmet_rocm_stable.txt


Read more