Magnum.np.distributed: Accelerating Finite Difference Micromagnetic Simulations with Multiple GPUs
This work addresses the single-device limitation of existing GPU-accelerated micromagnetic solvers, enabling faster simulations of larger systems for researchers in nanomagnetism and spintronics.
The authors present the first Python-native multi-GPU micromagnetic simulation framework by extending magnum.np with PyTorch Distributed, achieving a 7.0x speedup across 8 GPUs for demagnetisation field calculations and a 6.8x speedup on CPU with NUMA pinning.
Micromagnetic simulations are essential tools in nanomagnetism and spintronics research. Although widely adopted solvers like Mumax3 and the Python-native magnum.np use GPU acceleration to improve performance, these tools are limited to single-device computation. In this work, we present the first Python-native multi-GPU micromagnetic framework by extending magnum.np with PyTorch Distributed. This leverages high-speed communication and computation across multiple GPUs while retaining the benefits of ease of installation, platform-agnostic design, and compatibility with Python. For computationally intensive demagnetisation effective-field calculations, we achieve a 7.0x speedup across 8 GPUs connected via NVLink, whereas Halo exchange required for Heisenberg exchange shows limited scaling due to kernel dispatch latency. We also demonstrated the framework's versatility by achieving a 6.8x speedup in demagnetisation field computation on CPU with NUMA pinning via the MPI backend of PyTorch Distributed. Faster turnaround times will enable researchers to explore larger, more complex systems and accelerate the design cycle for novel spintronic devices.