기술 요약: 최대 수천 개의 GPU 를 활용한 스파이킹 신경망의 확장 가능한 구축

문제 제기

인간 대뇌 피질 규모로 대규모 스파이킹 신경망 (SNN) 을 시뮬레이션하는 것은 두 가지 주요 과제를 제시합니다: 개별 뉴런과 시냅스에 대한 막대한 메모리 요구사항과, 밀리초 미만의 정밀도로 역학을 해결하기 위한 높은 처리 속도의 필요성입니다. 수천 개의 GPU 를 장착한 고성능 컴퓨팅 (HPC) 시스템은 필요한 계산 밀도를 제공하지만, 기존 GPU 기반 시뮬레이션 소프트웨어는 계산 신경과학의 인프라 및 정확도 요구사항을 충족하면서 전체 컴퓨팅 클러스터로 확장되는 능력을 아직 입증하지 못했습니다.

대규모 포인트 뉴런 네트워크의 분산 시뮬레이션에서 특정 병목 현상은 컴퓨팅 클러스터의 서로 다른 노드 간에 스파이크를 통신하는 데 있습니다. 디지털 브레인 (Digital Brain) 이나 GeNN 과 같은 이전 접근 방식들은 개별 시냅스 정보를 생략하거나 단일 GPU 실행으로 제한되었습니다. 또한, NEST 와 같은 전통적인 CPU 기반 시뮬레이터는 라운드 로빈 뉴런 분산과 집단 통신에 의존하는데, 이는 동질적인 네트워크 구조를 가정하며 생물학적 뇌의 위상적 및 공간적 이질성을 활용하지 못합니다. NEST GPU 가 이러한 문제 중 일부를 해결했음에도 불구하고, 초기 네트워크 구축은 CPU 에서 GPU 메모리로 데이터를 전송하는 데 의존했으며, 동적 구축 방법은 이전에는 단일 GPU 시뮬레이션으로 제한되었습니다.

방법론

이 연구는 메시지 전달 인터페이스 (MPI) 를 사용하여 멀티 GPU 시스템에서 대규모 SNN 을 직접 구축하고 시뮬레이션하기 위한 새로운 메모리 효율적 방법을 제시합니다. 핵심 혁신은 구축 단계 동안 프로세스 간 통신 없이 네트워크 구축을 GPU 메모리 내에서 ("온보드") 완전히 수행하는 데 있습니다.

핵심 알고리즘

이 방법은 로컬 연결(동일한 MPI 프로세스 내의 뉴런) 과 원격 연결(서로 다른 프로세스에 걸친 뉴런) 을 구분합니다.

1. 독립적 구축: 각 MPI 프로세스는 네트워크의 해당 부분을 독립적으로 구축합니다. 다른 프로세스와 통신하지 않고 로컬 연결성을 생성하고 원격 연결을 위한 데이터 구조를 준비합니다.
2. 프록시 표현: 원격 연결의 경우, 이 방법은 대상 프로세스에서 "이미지 뉴런"(프록시) 을 사용합니다. 이들은 다른 MPI 랭크에 위치한 소스 뉴런의 가상 표현입니다.
3. 통신 맵: 알고리즘은 스파이크를 효율적으로 라우팅하기 위해 GPU 메모리에 연속적인 통신 맵을 인스턴스화합니다. 이러한 맵은 소스 랭크의 소스 뉴런 인덱스를 대상 랭크의 이미지 뉴런 인덱스와 연관시킵니다.
4. 통신 방식: 이 프레임워크는 네트워크 아키텍처에 따라 사용자가 선택할 수 있는 두 가지 MPI 통신 모드를 지원합니다.
   • 포인트 투 포인트: 두 프로세스 간의 직접 통신을 사용합니다. 뉴런이나 시냅스의 불균등한 분포를 가진 네트워크 (예: 멀티 에어리어 모델) 에 최적화되어 있습니다. 이는 스파이크를 라우팅하기 위해 특정 매핑 구조 $(R_{\tau,\sigma}, L_{\tau,\sigma})$ 와 시퀀스 $(T, P)$ 를 활용합니다.
   • 집단: 그룹 기반 통신 (예: MPI_Allgather) 을 사용합니다. 이는 균일한 통신 페이로드를 가진 균형 잡힌 네트워크에 유리합니다. 이는 여러 프로세스에 걸쳐 스파이크 라우팅을 관리하기 위해 그룹별 인덱싱 배열 및 호스트 배열을 사용합니다.

GPU 메모리 최적화

GPU 메모리 소비와 시뮬레이션 속도 사이의 균형을 맞추기 위해, 저자들은 네 가지 GPU 메모리 레벨 (GML) 을 구현했습니다:

• 레벨 0: 원격 연결 맵과 연결 수는 CPU 메모리에 저장됩니다.
• 레벨 1: 레벨 0 과 유사하지만 모든 소스 뉴런이 대상 프로세스에 이미지를 가진다고 가정하여 실제 사용 여부를 확인하는 것을 생략합니다 (더 빠른 구축, 잠재적으로 더 높은 메모리 낭비).
• 레벨 2: 맵과 연결 인덱스는 GPU 메모리에 저장되며, 연결 수는 실시간으로 계산됩니다. 이는 기본 레벨입니다.
• 레벨 3: 연결 수를 포함한 모든 데이터 구조가 GPU 메모리에 저장되어, 더 높은 GPU 메모리 사용량을 대가로 CPU-GPU 데이터 전송을 최소화합니다.

평가된 모델

1. 멀티 에어리어 모델 (MAM): 마카크 원숭이 피질의 32 개 시각 관련 영역에 대한 생물학적 세부 모델 ($4.13 \times 10^6$개 뉴런, $24.2 \times 10^9$개 시냅스). 이 모델은 복잡한 계층적 연결성을 특징으로 하며 포인트 투 포인트 통신을 사용하여 시뮬레이션되었습니다.
2. 확장 가능한 균형 네트워크: 고정된 인-디그리 연결성을 가진 흥분성 및 억제성 뉴런의 무작위 네트워크로, 약한 확장성 성능을 평가하도록 설계되었습니다. 이 모델은 최대 1,024 개의 GPU 에서 집단 통신을 사용하여 시뮬레이션되었습니다.

주요 결과

네트워크 구축 성능

"온보드"GPU 구축 방법은 이전 "오프보드"(CPU 기반) 접근 방식에 비해 상당한 속도 향상을 보여주었습니다:

• MAM 시뮬레이션: 네트워크 구축 시간이 686.0 초(오프보드) 에서 55.5 초(온보드) 로 감소하여 12.4 배의 속도 향상을 보였습니다.
  • 로컬 연결 생성은 20 배의 속도 향상을 보였습니다.
  • 원격 연결 생성은 9 배의 속도 향상을 보였습니다.
  • 뉴런/장치 생성 및 시뮬레이션 준비는 각각 350 배와 50 배의 속도 향상을 보였습니다.
• 확장 가능한 균형 네트워크: 이 방법은 1,024 개의 GPU(256 개 노드) 에서 최대 2 억 3,040 만 개 뉴런과 2.59 조 개 시냅스로 구성된 네트워크를 성공적으로 구축했습니다.

상태 전파 및 확장성

• MAM: 상태 전파 시간 (실시간 인자로 측정) 은 오프보드 및 온보드 버전 간에 유사하게 유지되었습니다 (약 15~16). 이는 구축 최적화가 시뮬레이션 역학에 부정적인 영향을 미치지 않았음을 나타냅니다.
• 균형 네트워크: 시스템은 1,024 개의 GPU 까지 약한 확장성을 보여주었습니다.
  • 메모리 효율성: GPU 메모리 레벨 0 은 NVIDIA A100 GPU(64 GB) 의 메모리 한계를 초과하지 않고 4,096 개 노드까지 시뮬레이션을 가능하게 했습니다. 더 높은 메모리 레벨 (2 및 3) 은 더 빠른 구축 및 시뮬레이션 속도를 제공했지만 더 낮은 노드 수 (레벨 3 의 경우 약 3,072 개 노드) 에서 메모리 한계에 도달했습니다.
  • 성능: 균형 네트워크에서 스파이크 기록을 비활성화하면 상태 전파 시간이 약 20% 감소했습니다.

검증

새로운 구축 방법은 이전 오프보드 버전 및 CPU 기반 NEST 시뮬레이터와 비교하여 검증되었습니다. 새로운 알고리즘으로 인한 난수 생성 시퀀스의 변화에도 불구하고, 스파이킹 활동의 통계적 특성 (방전율, 스파이크 간격의 변동 계수, 쌍별 피어슨 상관관계) 이 보존되어 시뮬레이션의 생물학적 타당성을 확인했습니다.

중요성 및 주장

이 논문은 개별 시냅스 정보를 저장하면서 수천 개의 GPU 로 구성된 전체 컴퓨팅 클러스터로 확장 가능한 최초의 GPU 기반 SNN 시뮬레이션 소프트웨어를 제공한다고 주장합니다. 주요 기여점은 다음과 같습니다:

1. 확장 가능한 구축: CPU-GPU 전송 병목 현상을 제거하고 구축 단계 동안 MPI 통신을 피하면서 GPU 메모리 내에서 직접 네트워크 연결성을 구축하는 새로운 알고리즘.
2. 유연성: 계층적 대 무작위/균형과 같은 다양한 네트워크 위상에 적응할 수 있도록 포인트 투 포인트 및 집단 MPI 통신 모두 지원.
3. 엑사스케일 준비성: 저자들은 향후 엑사스케일 슈퍼컴퓨터인 JUPITER 에서 $2 \times 10^{10}$개 뉴런과 $10^{14}$개 시냅스로 구성된 네트워크를 시뮬레이션할 수 있다고 외삽했습니다. 이 규모는 개별 시냅스 분해능을 유지하면서 인간 대뇌 피질의 연결성에 근접합니다.
4. 효율성: GML 시스템을 통해 메모리 사용을 최적화함으로써 기존 하드웨어에서 더 큰 네트워크의 시뮬레이션을 가능하게 합니다 (예: 32 개 대신 8 개의 GPU 에서 MAM 적합) 그리고 미래 엑사스케일 시스템의 전체 용량을 활용할 수 있는 경로를 제공합니다.

저자들은 이 접근 방식이 분산 시뮬레이션에서 스파이크 통신의 중요한 병목 현상을 해결하고, 현대 HPC 아키텍처에서 대규모 생물학적 세부 신경 시뮬레이션을 위한 참조 플랫폼으로 NEST GPU 를 확립한다고 결론지었습니다.