2x to 8x Faster: Helping Qubit Pharmaceuticals run hybrid molecular dynamics on LUMI

By Julius Philipp Roede, Danish e-infrastructure consortium (DeiC)

 

Qubit Pharmaceuticals runs a hybrid workload that blends deep learning models with traditional physics-based molecular dynamics (MD) simulations, a combination that speeds up the iteration rate of molecular design, a critical factor in modern drug discovery. To scale this work, the team needed their software stack running efficiently on LUMI, the AMD-based EuroHPC supercomputer, and a clear picture of how its performance compared to the NVIDIA systems they were used to.

 

 

The challenge: an AMD-shaped gap

 

Qubit’s workload depends on JAX, and that’s where the trouble started. The pre-built JAX libraries compatible with LUMI were several versions behind the state of the art, and the most recent version available at the time contained computational errors, rendering it unusable for high-precision scientific calculatiosimns. Building newer JAX versions from source for AMD hardware isn’t straightforward either, it requires complex dependency resolution; on top of ROCm driver compatibility issues with the LUMI system that were causing subpar performance.

 

On top of the AMD/JAX gap, Qubit faced two further challenges: their MD simulations mix Fortran, C, HIP, Python, and Just-In-Time compiled code, a combination that most profiling tools struggle to handle; and the simulations themselves are complex enough to make testing and optimisation slow going.

 

 

How EPICURE helped to upgrade JAX and profile for AMD GPUs

 

EPICURE’s support focused on three fronts: gradually updating the AI libraries to improve performance, advancing code profiling to pinpoint bottlenecks specific to AMD systems, and exploring optimisation routes such as the low-level C PJRT API and converting JAX models to ONNX.

 

In practice, this meant:

 

• – Building custom Docker build scripts tailored to the LUMI system to ensure reproducible and isolated build processes;
• – Systematically testing JAX versions, until the team could move past the broken v0.5 release toward a stable release;
• – Settling on the Tau profiler for its ability to trace both Python and HIP/ROCm kernels inside the same Fortran-compiled executable;
• – Implementing MD micro-benchmarks to compare different optimisation strategies;
• – Finalising with a container image running JAX v0.8.3, a significant upgrade from the initial v0.4.35.

 

 

The results: 2x–8x faster MD simulations

 

The work paid off on both fronts:

 

• – JAX was upgraded from v0.4.35 to v0.8.3, a major version jump that had been holding the team back until now.
• – Simulation runtime improved by 2× to 8× across internal MD benchmarks ranging from 2,000 to 96,000 atoms, with larger systems seeing the biggest gains.

 

Beyond these numbers, the profiling and benchmarking work also yielded concrete recommendations for Qubit’s future hybrid MD work on AMD systems.

 

What’s next: scaling to Alice Recoque and the wider HPC community

 

A verified, high-performance JAX environment on AMD architecture is part of Qubit’s roadmap for Alice Recoque, the upcoming French supercomputer that is set to become a key system for the company going forward.

 

EPICURE is also turning this work into resources for the wider community, including a how-to guide for building newer JAX versions on AMD systems, already shared with the LUMI AI Factory, with more technical write-ups on the way.