magnum.np.distributed：利用多块GPU加速有限差分微磁模拟

《Scientific Reports》：magnum.np.distributed: accelerating finite difference micromagnetic simulations with multiple GPUs

　　

摘要

微磁模拟是纳米磁学和自旋电子学研究中的重要工具。尽管像Mumax3以及基于Python的magnum.np这样的常用求解器采用了GPU加速来提升性能，但这些工具仍仅限于单设备计算。在本研究中，我们通过结合PyTorch Distributed扩展magnum.np，提出了目前所知首个基于Python的多GPU微磁计算框架。该框架能够利用多个GPU之间的高速通信与计算能力，同时保留了安装简便、平台无关以及与Python兼容的优点。在需要大量计算的退磁场有效场计算中，使用通过NVLink连接的8个GPU和4280万个单元的合成测试基准，该框架实现了7.0倍的加速效果；而海森堡交换所需的Halo交换方法由于内核调度延迟的存在，其扩展性较为有限。我们还通过PyTorch Distributed的MPI后端，利用NUMA绑定技术在CPU上对100万个单元的合成测试基准进行退磁场计算，同样实现了6.8倍的加速效果，进一步证明了该框架的通用性。更快的计算速度将帮助研究人员探索更大、更复杂的系统，从而加快新型自旋电子器件的设计进程。源代码地址：https://gitlab.com/jedcheng/magnum-np-distributed。