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

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 .