Magnum.np.distributed: Accelerating Finite Difference Micromagnetic Simulations with Multiple GPUs

이 논문은 더 크고 복잡한 스핀트로닉스 시스템의 시뮬레이션을 가능하게 하기 위해, PyTorch Distributed를 기반으로 구축되어 탈자장(demagnetization field) 계산을 대폭 가속화하는(8개의 GPU에서 최대 7.0배의 속도 향상 달성) 최초의 Python 네이티브 멀티 GPU 마이크로마그네틱 프레임워크인 Magnum.np.distributed를 소개한다.

핵심

당신은 컴퓨터 칩 내부에서 아주 작은 자석이 어떻게 움직이는지 시뮬레이션하려고 한다고 상상해 보세요. 이를 위해 과학자들은 자석을 수백만 개의 작은 레고 블록(이를 '메쉬(mesh)'라고 부릅니다)으로 나누고, 각 블록이 이웃한 블록에 어떻게 밀거나 당기는지를 계산합니다. 이것을 **마이크로마그네틱 시뮬레이션(micromagnetic simulation)**이라고 합니다.

오랫동안 이러한 시뮬레이션은 마치 한 사람이 산을 옮기려는 것과 같았습니다. 강력한 그래픽 카드(GPU)의 도움을 받더라도, 대부분의 소프트웨어는 한 번에 하나의 카드만 사용할 수 있었습니다. 그것도 빠르긴 했지만, 가장 크고 복잡한 자기적 난제를 풀기에는 충분히 빠르지 않았습니다.

이 논문은 Magnum.np.distributed라는 새로운 도구를 소개합니다. 이것은 마치 한 사람이 산을 옮기는 것에서, 모든 작업자가 자신만의 초강력 삽을 가진 채 함께 일하는 건설 팀으로 업그레이드하는 것과 같습니다.

작동 원리는 다음과 같으며, 쉬운 개념들로 나누어 설명합니다:

1. "팀워크" 문제

예전에는 8개의 그래픽 카드를 사용하고 싶다면, 그들이 서로 대화할 수 있도록 복잡하고 어려운 코드(C++나 CUDA 같은)를 작성해야 했습니다. 이는 마치 8명의 사람들이 서로 다른 언어를 말하고 서로 다른 유니폼을 입은 채 집을 짓는 것과 같았습니다.

새로운 Magnum.np.distributed 프레임워크는 현대 데이터 과학의 "영어"인 Python의 언어를 사용합니다. 이 도구는 PyTorch Distributed라는 라이-브러리를 사용하여 여러 개의 GPU가 즉각적으로 서로 통신할 수 있게 해줍니다.

• 비유: 이어달리기를 상상해 보세요. 예전 시스템에서는 바톤(데이터)을 느리고 수동적인 방식으로 전달해야 했습니다. 이 새로운 시스템에서는 주자들(GPU)이 고속 광섬유 터널(NVLink)로 연결되어 있어, 바톤을 거의 즉시 전달할 수 있습니다.

2. 작업을 나누는 방법

시뮬레이션 그리드는 빵 한 덩이를 써는 것처럼 나뉩니다.

• 로컬 작업: 어떤 계산은 특정 빵 조각에만 의존합니다. 이것은 쉽습니다. 각 GPU가 자신의 조각을 독립적으로 수행하면 됩니다. 이는 8명의 사람이 방 안의 서로 다른 벽 8개를 칠하는 것과 같습니다. 그들은 서로 대화할 필요가 없습니다.
• "헤일로(Halo)" 문제: 어떤 계산(예: '교환장(exchange field)')은 이웃이 무엇을 하고 있는지 알아야 합니다. 만약 당신이 벽의 가장자리를 칠하고 있다면, 옆에 있는 벽이 어떤 색인지 알아야 합니다.
  • 해결책: 시스템은 각 조각의 가장자리에 "헤일로"(완충 지대)를 만듭니다. GPU들은 이 가장자리 데이터를 이웃과 교환합니다.
  • 함정: 논문에서는 이 과정이 잘 작동하지만, 컴퓨터의 두뇌(CPU)가 GPU에게 작업을 시작하라고 명령하는 데 걸리는 시간 때문에 이 가장자리 데이터를 주고받는 과정이 때때로 느려진다는 것을 발견했습니다. 이는 마치 달리기 선수는 매우 빠르지만, 코치가 "가!"라고 외치는 데 시간이 너무 오래 걸리는 것과 같습니다.

3. 큰 도전 과제: "탈자기(Demagnetization)" 장

시뮬레이션에서 가장 어려운 부분은 탈자기장을 계산하는 것입니다. 이것은 "글로벌" 계산으로, 자석 전체의 모든 레고 블록이 다른 모든 블록의 영향력을 느끼는 것을 의미합니다.

• 비유: 방 안에 가득 찬 사람들이 동시에 서로에게 메시지를 외쳐야 하는 상황을 상상해 보세요.
• 결과: 새 프레임워크는 FFT(고속 푸리에 변환)라는 수학적 트릭을 사용하여 이 문제를 처리합니다. 데이터를 재배치하여 GPU들이 병렬로 작업할 수 있게 합니다.
  • 속도: 초고속 터널(NVLink)로 연결된 8개의 강력한 GPU 시스템에서, 이 특정 작업은 단 하나의 GPU를 사용할 때보다 7배 더 빨라졌습니다. 이는 거의 완벽한 선형적 속도 향상이었습니다.

4. 일반 컴퓨터(CPU)는 어떤가요?

모두가 8개의 고성능 GPU가 있는 클러스터를 가진 것은 아닙니다. 저자들은 또한 이 작업을 표준 컴퓨터 프로세서(CPU)에서도 테스트했습니다.

• 문제: CPU는 서로 다른 메모리 영역(NUMA)을 가지고 있습니다. 프로그램이 "먼" 메모리 영역에서 데이터를 가져오면 속도가 느려집니다.
• 해결책: 그들은 프로그램이 자신의 "로컬" 메모리 근처에 머물도록 강제하는 NUMA 피닝(pinning) 기술을 사용했습니다.
• 결과: 강력한 듀얼 소켓 CPU에서, 이 기술은 피닝을 하지 않았을 때보다 시뮬레이션을 6.8배 더 빠르게 만들었습니다. 최고급 GPU보다는 여전히 약 10배 정도 느리지만, 이는 값비싼 GPU를 사용할 수 없는 연구자들도 이전보다 훨씬 빠르게 이러한 복잡한 시뮬레이션을 실행할 수 있음을 의미합니다.

5. 실제 테스트: "미로(Maze)" 자석

이것이 실제로 작동하는지 증명하기 위해, 그들은 고급 메모리 장치에 사용되는 실제 자기 스택(Pt/Gd/Co/Ni)을 시뮬레이션했습니다.

• 설정: 2,360만 개의 셀로 구성된 그리드.
• 결과:
  • 1 GPU: 50.6시간 소요.
  • 4 GPU: 8.4시간 소요.
  • 속도 향상: 6배 더 빠름.
• 예상보다 더 좋았던 이유: 작업을 4개의 GPU로 나누었을 때, 데이터가 GPU의 작고 빠른 내부 메모리(캐시)에 더 잘 들어맞아 전체 과정이 훨씬 더 매끄럽게 진행되었습니다.

요약

이 논문은 과학자들이 여러 개의 GPU를 사용하여 작은 자석을 시뮬레이션할 수 있게 해주는 최초의 Python 네이티브 도구를 제시합니다.

• 코딩의 악몽 없음: C++ 전문가가 될 필요가 없습니다. 그냥 Python을 작성하면 됩니다.
• 압도적인 속도: 8개의 GPU에서 시뮬레이션을 7배 더 빠르게 만들 수 있습니다.
• 다재다능함: 고성능 GPU 클러스터와 적절한 설정을 갖춘 일반 컴퓨터 CPU 모두에서 작동합니다.

이를 통해 연구자들은 훨씬 더 짧은 시간 안에 더 크고 복잡한 자기 시스템을 시뮬레이션할 수 있으며, 이는 차세대 컴퓨터 메모리 및 로직 기술인 스핀트로닉스(spintronics) 장치를 더 빠르게 설계하는 데 도움을 줍니다.

기술 요약

기술 요약: Magnum.np.distributed

문제 정의

미로자기(Micromagnetic) 시뮬레이션은 나노자성 및 스핀트로닉스 연구에 필수적이지만, 기존의 고성능 솔버들은 접근성과 확장성 측면에서 상당한 한계에 직면해 있다. Mumax3나 Python 기반의 magnum.np와 같은 도구들은 GPU 가속을 활용하지만, 단일 장치 계산에 국한되어 있다. 이러한 제약은 현대적 소자 설계를 위해 필요한 더 크고 복잡한 시스템의 시뮬레이션을 방해한다. 또한, 널리 사용되는 솔버들은 종종 비-Python 스크립팅 언어(예: Go, TCL)에 의존하거나 C++/CUDA 코드를 소스에서 직접 컴파일해야 하므로, 설치의 어려움, 교차 플랫폼 호환성 문제, 그리고 Python 기반 분석 워크플로와의 원활한 통합을 저해하는 장벽을 만든다. Boris와 같이 C++/CUDA 기반의 시뮬레이터 중 일부는 멀티 GPU 연산을 지원하지만, Python-native 프레임워크가 가진 설치의 용이성과 플랫폼 불가지론적 설계는 갖추지 못했다.

방법론

저자들은 기존의 magnum.np 솔버를 PyTorch Distributed로 확장함으로써, 최초의 Python-native 멀티 GPU 미로자기 프레임워크인 magnum.np.distributed를 제시한다. 이 구현은 vendor 특정 CUDA 최적화를 피하면서 플랫폼 불가지론성을 유지하기 위해 PyTorch의 Just-In-Time (JIT) 컴파일(TorchDynamo 및 TorchInductor를 통한)과 autograd 기능을 활용한다.

핵심 구현 세부 사항

1. 도메인 분할 (Domain Decomposition): 시뮬레이션 메쉬를 x축을 따라 연속적인 슬래브(slab)로 분할하며, 각 프로세스 랭크(rank)에 하나의 슬래브를 할당한다.
2. 로컬 및 글로벌 필드:
   • 로컬 필드 (Local Fields): Slonczewski 스핀 전달 토크(spin-transfer torque) 및 단축 이방성(uniaxial anisotropy)과 같은 항들은 로컬 데이터에 작용하며, 프로세스 간 통신 없이 분산된 슬래브에 직접 적용된다.
   • 헤일로 교환 (Halo Exchange): Heisenberg 교환 및 Dzyaloshinskii-Moriya 상호작용(DMI)과 같은 비-로컬 상호작용을 위해, 프레임워크는 헤일로 영역(슬래브 경계에 두 개의 추가 셀)을 구현한다. 통신은 배치(batched) 처리된 비블로킹(non-blocking) 포인트 투 포인트 전송을 통해 처리된다. 결정적으로, 이러한 헤일로 교환은 TorchDynamo의 최적화 문제를 피하기 위해 JIT 컴파일 범위 외부에 배치되어, 계산 부분만을 컴파일하도록 한다.
3. 감자 필드 (Demagnetization Field): 가장 계산 집약적인 글로벌 상호작용인 감자 필드는 Newell의 공식을 사용하여 이산 컨볼루션(discrete convolution)을 통해 계산된다. 구현에는 분산 3D 고속 푸리에 변환(FFT) 전략이 사용된다:
   • 메쉬를 x축을 따라 분해한다.
   • y- 및 z-방향으로 2D FFT를 수행한다.
   • All-to-all transpose를 통해 데이터를 y-축을 따라 샤드 벡터(shard vectors)로 재배치한다.
   • x-방향으로 제로 패딩(zero-padding) 및 1D FFT를 수행한다.
   • 감자 커널(demagnetization kernel)과의 요소별 곱셈(pointwise multiplication)을 수행한다.
   • 역방향 패스(inverse pass)가 이 단계들을 역순으로 수행한다.
   • 참고: 저자들은 코드 가독성을 유지하고 현재 JIT 컴 compiler가 복소수 연산을 단일 커널 호출로 퓨전(fuse)할 수 없기 때문에, 다른 솔버(Boris 등)에서 사용되는 정밀도를 절반으로 줄이는 기능을 제거하였다.
4. 솔버 통합: 프레임워크는 시간 적분을 위해 Runge-Kutta-Fehlberg 방법(RKF45)을 지원하며, 에너지 최소화를 위해 Barzilai-Borwein 알고리즘을 지원한다. 이러한 솔버들은 모든 랭크가 동일한 타임 스텝으로 전진하거나 동시에 종료되도록 보장하기 위해 글로벌 리덕션(예: 최대 오차 또는 수렴 기준)을 수행하도록 수정되었다.

백엔드 지원

프레임워크는 NVLink/NVSwitch 및 InfiniBand로 연결된 NVIDIA GPU를 위한 NCCL 백엔드와, CPU 기반 계산을 위한 MPI 백엔드를 대상으로 한다. MPI 백엔드는 특히 CPU 성능을 최적화하기 위해 NUMA(Non-Uniform Memory Access) 피닝(pinning)을 지원한다.

주요 기여

• 최초의 Python-Native 멀티 GPU 프레임워크: C++ 또는 CUDA 컴파일 없이도 멀티 GPU를 지원하도록 magnum.np를 확장하여, 설치의 용이성과 Python 생태계 호환성을 보존하였다.
• 분산 FFT 알고리즘: 분산된 슬래브 전반에서 all-to-all transpose와 제로 패딩을 효율적으로 처리하는 확장 가능한 3D FFT 알고리즘을 구현하였다.
• 하이브리드 MPI+OpenMP CPU 지원: PyTorch MPI 백엔드를 통해 NUMA 피닝을 활용하여 CPU 클러스터를 효과적으로 사용할 수 있음을 입증하였으며, 이는 GPU 자원이 제한된 환경을 위한 실행 가능한 대안을 제공한다.
• 최소한의 코드 마이그레이션: 비-분산 버전에서 분산 버전으로 마이그레이션하는 데 필요한 변경 사항이 약 8줄(분산 솔버 임포트, 프로세스 그룹 초기화, 표준 메쉬/상태 객체를 분산형 대응물로 교체)에 불과함을 보여준다.

결과

검증

프레임워크는 muMAG 표준 문제 4(permalloy 자석의 자기 반전) 및 DMI와 도메인 벽 피닝(domain wall pinning)이 포함된 문제에 대해 검증되었다. 분산 시뮬레이션(2, 4, 8개 랭크 사용) 결과는 기존의 비-분산 magnum.np 및 Mumax3(단정밀도) 결과와 일치하였으며, 이를 통해 분산된 LLG-RKF45 솔버, 감자 및 교환 필드 구현의 정확성을 확인하였다.

성능 벤치마크

• 멀티 GPU 확장성 (NVLink/NVSwitch):
  • 감자 필드: NVSwitch로 연결된 8개의 NVIDIA H100 HBM3 GPU를 사용하여 7.0배의 속도 향상을 달al성했다. 확장성은 거의 선형적이었으며 다양한 메모리 대역폭 세대(HBM3 vs HBM2e)에 걸쳐 일관되게 나타났는데, 이는 성능이 원시 대역폭보다는 알고리즘의 통신-계산 비율에 의해 결정됨을 나타낸다.
  • 교환/DMI 필드: 포인트 투 포인트 통신에서의 커널 디스패치 지연 시간으로 인해 작은 문제 크기에서는 확장성이 제한되었다. 의미 있는 속도 향상(예: 8개 GPU에서 4.2배)은 더 큰 메쉬(>8M cells)에서만 관찰되었다.
  • 로컬 필드 (STT): 인터-랭크 통신이 필요하지 않으므로 이상적인 속도 향상을 달성했다.
• 멀티 노드 확장성:
  • 400 Gbps InfiniBand로 연결된 2-노드 클러스터(총 8개 GPU) 테스트 결과, 특히 감자 필드에서 단일 노드 NVSwitch 시스템에 비해 상당한 성능 저하(72% 저하)를 보였다. 이는 인터커넥트 병목 현상(InfiniBand vs NVSwitch)과 all-to-all transpose 중에 발생하는 높은 데이터 전송량 때문으로 분석된다.
• CPU 성능 (NUMA 피ning):
  • 듀얼 소켓 Intel Xeon 노드에서 MPI 백엔드를 통한 NUMA 피닝을 활성화했을 때, $10^6$ 셀 문제에 대한 감자 필드 계산 시간이 6.8배 감소(단계당 204.0 ms에서 29.8 ms로)하였다.
  • 단일 H100 GPU보다 10~15배 느리지만, 이 성능은 GPU 접근 권한이 없는 연구자들에게 CPU 기반 시뮬레이션을 실행 가능한 옵션으로 만들어 준다.
• 실제 응용 사례:
  • 2,360만 개의 셀을 가진 계면-DMI 유도 다층 구조(Pt/Gd/Co/Ni) 시뮬레이션에서, 프레임워크는 단일 GPU 대비 4개의 GPU에서 6.0배의 속도 향상을 달성했다. 이러한 초선형(superlinear) 확장성은 작업 집합(working set)이 여러 장치로 분산될 때 개선된 GPU 캐시 활용 능력 덕분에 발생했다.

의의 및 주장

본 논문은 magnum.np.distributed가 고성능 멀티 GPU 컴퓨팅과 Python-native 과학 소프트웨어 특유의 사용 용이성 사이의 간극을 성공적으로 메웠다고 주장한다. C++/CUDA 컴파일의 필요성을 제거하고, GPU와 CPU(NUMA 최적화 포함) 모두를 위한 통합 인터페이스를 제공함으로써, 고충실도 미로자기 시뮬레이션에 대한 진입 장벽을 낮춘다.

저자들은 빠른 턴어라운드 타임이 연구자들이 더 크고 복잡한 나노자성 시스템을 탐구할 수 있게 하여, 새로운 스핀트로닉 소자의 설계 주기를 가속화할 것이라고 강조한다. 본 연구는 차세대 고대역폭 인터커넥트(예: NVIDIA NVL72)의 혜택을 받을 수 있는 위치에 있으며, 현재로서는 단일 노드 멀티 GPU 및 NUMA 인식 CPU 클러스터를 위한 견고한 솔루션을 제공한다. 코드 소스는 채택과 추가 개발을 촉진하기 위해 공개되어 있다.