magnum.np.distributed: accelerating finite difference micromagnetic simulations with multiple GPUs

Tsz Chung Cheng, Yuichiro Kurokawa, Hiromi Yuasa

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, to our knowledge, 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 on a synthetic benchmark with 42.8M cells across 8 GPUs connected via NVLink, whereas the Halo exchange required for Heisenberg exchange shows limited scalability due to kernel dispatch latency. We also demonstrated the framework’s versatility by achieving a 6.8x speedup in demagnetisation field computation on the CPU with NUMA pinning via the MPI backend of PyTorch Distributed, using a 1M-cell synthetic benchmark. Faster turnaround times will enable researchers to explore larger, more complex systems and accelerate the design cycle for novel spintronic devices. Source code: https://gitlab.com/jedcheng/magnum-np-distributed .