[Paper] FastMPS: 대규모 Matrix Product State 샘플링에서 데이터 병렬 재검토

Source: arXiv - 2512.20064v1

Overview

Fast‑MPS는 새로운 “tensor‑parallel” 레이어와 결합하여 결합 차원(bond dimension) 따라 작업을 분할함으로써 대규모 매트릭스 프로덕트 스테이트(Matrix Product States, MPS) 샘플링을 위한 데이터 병렬성을 되살립니다. 저자들은 이 이중 레벨 병렬성이 이전 MPS 시뮬레이터를 제한해 온 메모리 및 I/O 병목 현상을 깨뜨려, 8 k+ 사이트, χ = 10⁴와 같은 전례 없는 규모의 양자 샘플링 실행을 10배 이상의 속도 향상으로 가능하게 함을 보여줍니다.

주요 기여

• Multi‑level parallel framework: 독립적인 샘플들에 대한 고전적인 데이터 병렬성과 무거운 텐서 수축을 MPS 결합 차원에 따라 분산시키는 새로운 텐서 병렬성을 결합합니다.
• Memory‑ and I/O‑aware design: 실시간 압축과 겹치는 통신/연산을 도입하여 프로세스당 메모리 사용량을 낮게 유지하고 데이터 이동 지연을 숨깁니다.
• Scalable implementation: MPI + NCCL(또는 유사한 고성능 집합 연산) 위에 구축되었으며 수천 개의 CPU/GPU 프로세스에서 효율적으로 실행됩니다.
• Benchmark on Gaussian Boson Sampling (GBS): 최첨단 시뮬레이터 대비 10배 이상의 속도 향상을 달성했으며 χ = 10⁴인 8,176 격자 사이트까지 확장했습니다.
• Open‑source reference: 저자들은 기존 텐서 네트워크 툴킷에 플러그인할 수 있는 프로토타입 라이브러리를 공개했습니다.

Methodology

1. Problem decomposition

   • Data parallelism: 각 MPI 랭크는 전체 MPS 샘플 중 일부를 독립적으로 생성합니다.
   • Tensor parallelism: 하나의 랭크 안에서, MPS 텐서는 두 번째 차원의 프로세스들에 걸쳐 분할되어 결합 인덱스(χ 차원)의 서로 다른 슬라이스를 각각 보유합니다.

2. Compressed tensor storage

   • 분배하기 전에 각 텐서는 양자화/압축(예: 저‑랭크 근사 또는 블록‑스파스 인코딩)되어 메모리 사용량을 줄이면서 샘플링 정확도는 유지합니다.

3. Overlap of communication & computation

   • 텐서의 한 슬라이스가 수축되는 동안, 다음 슬라이스는 원격 랭크에서 스트리밍됩니다. 이 파이프라인은 논블로킹 MPI 호출과 CUDA 스트림(GPU 사용 시)으로 조정됩니다.

4. Collective contraction engine

   • 커스텀 all‑reduce / all‑gather 패턴이 텐서‑패럴렐 그룹으로부터 부분 결과를 집계한 뒤, MPS 체인의 다음 사이트로 진행합니다.

5. Sampling loop

   • 알고리즘은 MPS 사이트를 순차적으로 탐색하면서 일련의 조건부 확률 계산(텐서 수축을 통해)을 수행해 각 물리적 인덱스를 추출합니다. 이는 표준 MPS 샘플링과 동일하지만 완전히 병렬화된 형태입니다.

전체 워크플로는 두 차원 프로세스 격자로 시각화할 수 있습니다: 행은 서로 다른 샘플을 처리하고, 열은 동일 샘플의 서로 다른 χ‑슬라이스를 처리합니다.

결과 및 발견

저자들은 또한 샘플링된 분포의 통계적 특성(예: GBS에서의 광자 수 히스토그램)이 압축에도 불구하고 정확한 MPS와 거의 차이가 없음을 보여주어, 압축이 과학적 정확성을 저하시키지 않음을 확인하였다.

Practical Implications

• Quantum‑sampling research: 연구자들은 이제 훨씬 더 큰 GBS 인스턴스를 시뮬레이션할 수 있게 되었으며, 이는 근시일 내 양자 광자 장치의 검증과 양자 우위 주장에 대한 벤치마크에 도움이 된다.
• Tensor‑network libraries: Fast‑MPS의 두 단계 병렬 처리를 인기 있는 파이썬 패키지(e.g., tensornetwork, quimb)의 백엔드로 추상화할 수 있어, 개발자에게 MPS 기반 워크플로우에 즉시 적용 가능한 성능 향상을 제공한다.
• High‑performance ML: MPS는 시퀀스 데이터에 대한 컴팩트 모델로 주목받고 있으며, Fast‑MPS는 대규모 데이터셋에 대한 학습/추론을 기존 HPC 클러스터에서도 실현 가능하게 만든다.
• Memory‑constrained environments: 압축 + 겹침 전략은 I/O가 병목인 다른 대규모 텐서 작업(e.g., 딥러닝 모델 병렬화, 과학 시뮬레이션)에도 적용할 수 있다.
• Scalable cloud deployments: 이 접근 방식이 표준 MPI/NCCL 프리미티브에 의존하기 때문에, 맞춤형 하드웨어 없이도 클라우드 기반 GPU 팜(AWS, Azure)으로 이식할 수 있어, 온‑디맨드 대규모 텐서 샘플링을 서비스 형태로 제공한다.

제한 사항 및 향후 작업

• Compression trade‑offs: 저자들은 샘플링 정확도에 최소한의 영향을 보고했지만, 공격적인 압축은 더 섬세한 양자 상태 특성에 영향을 줄 수 있다; 체계적인 오류 분석이 아직 필요하다.
• Hardware heterogeneity: 현재 구현은 비교적 균일한 CPU/GPU 클러스터를 가정한다; 혼합 정밀도 또는 이질적인 노드 구성으로 확장하는 것은 간단하지 않을 수 있다.
• Generality beyond MPS: Fast‑MPS는 MPS의 선형 체인 구조에 맞춰져 있다; 동일한 아이디어를 고차원 텐서 네트워크(PEPS, MERA)에 적용하려면 새로운 통신 패턴이 필요하다.
• Automation of parallel layout: 데이터‑와 텐서‑병렬 그룹 간 최적 분할을 선택하는 현재 방식은 수동 튜닝에 의존한다; 자동 튜너가 프레임워크를 보다 사용자 친화적으로 만들 수 있다.

전반적으로 Fast‑MPS는 이전에 도달할 수 없었던 문제 규모로 MPS 샘플링을 확장할 수 있는 실용적인 경로를 열어 주며, 그 설계 원칙은 보다 넓은 범위의 고성능 텐서 연산에 영향을 미칠 가능성이 있다.

저자

• Yaojian Chen
• Si‑Qiu Gong
• Lin Gan
• Yanfei Liu
• An Yang
• Yinuo Wang
• Chao‑yang Lu
• Guangwen Yang

논문 정보

• arXiv ID: 2512.20064v1
• 카테고리: cs.DC
• 발행일: 2025년 12월 23일
• PDF: Download PDF