Getting JAX to perform on AMD: Optimising hybrid MD simulations on LUMI

 

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

 

Qubit Pharmaceuticals needed its hybrid molecular dynamics (MD) workload, which combines deep learning models with physics-based simulation, to run efficiently on LUMI, the Cray-built EuroHPC supercomputer powered by AMD MI250X GPUs. Getting there meant confronting a software ecosystem still catching up to NVIDIA on several fronts and building the tooling to measure exactly how much it mattered.

 

 

Where the AMD software stack fell short

 

At the start of the project, Qubit faced software ecosystem disparities between NVIDIA and AMD hardware:

 

• – Library availability: the pre-built JAX libraries compatible with LUMI (v0.4.35) lagged behind the state of the art (v0.8.2 and above);
• – Version-specific bugs: the highest available version at the time, v0.5, contained computational errors that made it unusable for high-precision scientific calculations;
• – Build complexity: compiling JAX from source for AMD hardware requires complex dependency resolution and configuration;
• – Driver compatibility: ROCm driver incompatibilities with LUMI, combined with a fast-moving JAX codebase, resulted in subpar performance.

 

Beyond the AMD/JAX gap itself, two further factors made the optimisation harder. Profiling was complicated by the sheer mix of languages and compilation models involved: Fortran, C, HIP, Python, and Just-In-Time compiled Python code all coexisting in the same execution path, a combination most profiling tools aren’t built to handle. And the MD simulation itself is complex, which made testing, profiling, and implementing optimisations slower at every step.

 

 

Where EPICURE focused its HPC support

 

EPICURE’s technical support targeted three areas: gradually updating the AI libraries to improve performance, advancing code profiling to pinpoint AMD-specific bottlenecks, and exploring alternative optimisation routes, including the low-level C PJRT API and converting JAX models to ONNX for inference.

 

1. Containerization. The team built custom Docker build scripts tailored specifically to the LUMI environment, ensuring the build process was both reproducible and isolated from the underlying system.
2. Iterative JAX upgrades. The EPICURE team systematically compiled and tested successive JAX versions, working around the errors present in v0.5 until reaching a stable, recent release.
3. Profiling tool selection. After testing multiple profiling tools, the team settled on the Tau profiler for its ability to simultaneously trace Python and HIP/ROCm kernels inside a Fortran-compiled executable.
4. Micro-benchmarking. MD-specific micro-benchmarks were developed to directly compare the impact of different optimisation strategies.
5. Final deployment. The result was a finalised container image running JAX v0.8.3, a significant upgrade from the initial v0.4.35.

 

 

Measuring the impact

 

 

Metric	Initial State	Final State	Improvement
Profiling	Rocprofv1	Tau profiler	Trace capability, including Python methods.
JAX Version	v0.4.35	v0.8.3	Major version upgrade resulting in significant performance improvements.
Performance	Baseline	2x – 8x faster	Significant reduction in simulation runtime for multiple internal MD benchmarks ranging from 2000 to 96000 atoms, where systems with more atoms see a larger speedup.

 

Beyond the headline numbers, the profiling and micro-benchmarking work also yielded practical recommendations for Qubit’s future hybrid MD development.

 

 

What this means going forward

 

• – For AMD supercomputing: Qubit now has a verified, high-performance JAX environment for AMD-based systems. This is directly relevant to their upcoming work on Alice Recoque, the next-generation French supercomputer that is set to become a primary target for their operations;
• – For the profiling ecosystem: This project reinforced a broader lesson that there’s no one-size-fits-all tool for profiling complex hybrid codebases. Choosing the right tool still depends on the specific mix of languages and execution models involved.

 

 

Sharing this work with the community

 

The EPICURE team at the Danish e-infrastructure consortium (DeiC) is turning this project into resources the wider HPC community can use directly:

 

• – A how-to guide for building newer versions of JAX on AMD systems. This has already been shared with the LUMI AI Factory (LAIF) to eventually incorporate JAX into standard LAIF containers;
• – A technical how-to blog post on using the low-level JAX runtime (PJRT) to launch AI inference from C;
• – A technical report on converting JAX models to ONNX for inference.