기술 요약: torch-sla

문제 정의

미분 가능한 희소 선형 대수는 그래프 신경망 (GNN), 비정형 메시 위의 신경 연산자, 그리고 미분 가능한 시뮬레이션 파이프라인을 포함한 과학적 머신러닝의 기초적인 요구사항입니다. 그러나 PyTorch 생태계는 이러한 워크로드를 지원할 통합 라이브러리가 부족합니다. 네이티브 PyTorch 지원은 분산되어 있습니다: torch.sparse 는 저수준 커널과 미분 불가능하며 CPU 전용인 spsolve(SuperLU 래퍼) 를 제공하지만, torch.linalg 는 밀집 행렬로 제한됩니다.

기존의 우회 방법들은 대규모 또는 복잡한 학습 시나리오에는 충분하지 않습니다:

1. 그래프 단절: 외부 솔버 (예: cuDSS, PETSc) 를 호출하면 autograd 그래프가 단절되어, 사용자는 모든 문제에 대해 기울기를 포기하거나 커스텀 torch.autograd.Function 래퍼를 작성해야 합니다.
2. 비효율적인 미분: $k$번의 반복 솔버를 통해 순전파를 수행하면 $O(k)$ 노드 크기의 계산 그래프가 생성됩니다. $10^6$개의 자유도 (DOF) 와 1,000 번의 반복을 가진 문제의 경우, 중간 벡터를 저장하는 데 약 80GB 의 메모리가 필요하여 대규모 학습이 불가능해집니다.
3. 확장성 한계: $10^7$–$10^9$개의 미지수를 가진 문제들은 분산 도메인 분할이 필요합니다. 기존 PyTorch 분산 원시 연산들은 역전파 단계에서 희소 연산자의 수반 (adjoint) 에 필요한 전치 헤일로 교환을 기본적으로 지원하지 않습니다.
4. 백엔드 분열: 직접 솔버 (예: cuDSS) 는 작은 문제 ($<10^5$ DOF) 에서는 빠르지만, 더 큰 시스템에서는 $O(n^{1.5})$의 채움 (fill-in) 과 메모리 고갈을 겪습니다. 반복 솔버는 확장성이 더 좋지만 작은 문제에서는 커널 시작 오버헤드가 발생합니다. 장치와 문제 크기에 따라 이질적인 백엔드 간에 디스패치할 통합 인터페이스는 존재하지 않습니다.

방법론

저자들은 이 격차를 메우기 위해 세 가지 핵심 구성 요소를 통해 오픈소스 라이브러리인 torch-sla를 제시합니다:

1. 통합 백엔드 추상화

torch-sla 는 장치와 문제 크기에 따라 다섯 가지 교체 가능한 백엔드 간에 디스패치하는 단일 autograd 인식 API (.solve, .matvec, .eigsh, .det) 를 노출합니다:

• CPU: SciPy(SuperLU, UMFPACK, CG, BiCGStab) 및 Eigen(CG, BiCGStab).
• GPU: NVIDIA cuDSS(LU, Cholesky, LDLT), CuPy(반복 및 직접), 그리고 PyTorch 네이티브 반복 솔버.
• 자동 디스패치: 시스템은 백엔드를 자동으로 선택합니다: 작은~중간 크기 문제 (예: GPU 에서 $<2 \times 10^6$ DOF) 에는 직접 솔버를, 메모리 오버플로우를 피하기 위해 더 큰 문제에는 반복 솔버를 사용합니다. 대칭성과 양의 정부호성은 인수분해를 최적화하기 위해 감지됩니다 (예: LU 를 Cholesky 로 업그레이드).

이 라이브러리는 국소성과 배치에 따라 조직된 네 가지 희소 텐서 유형을 지원합니다: SparseTensor(단일 행렬/공유 패턴), SparseTensorList(배치별 서로 다른 패턴), 그리고 이들의 분산 변형인 DSparseTensor 및 DSparseTensorList.

2. 수반 미분 프레임워크

순전파의 $O(k)$ 메모리 및 그래프 복잡성을 피하기 위해 torch-sla 는 암시적 함수 정리를 사용합니다.

• 선형 시스템: $Ax=b$의 경우, 역전파는 단일 수반 시스템$A^\top \lambda = \partial L / \partial x$를 풉니다. 행렬에 대한 기울기는 희소 패턴에서만$-\lambda_i x_j$로 계산됩니다.
• 비선형 시스템: $F(u, \theta)=0$의 경우, 역전파는 비선형 반복 횟수와 무관하게 하나의 수반 선형 풀이 ($J^\top \lambda = \partial L / \partial u$) 와 벡터 - 야코비 곱을 필요로 합니다.
• 고유값 문제: 고유값에 대한 기울기는 Hellmann–Feynman 정리 ($\partial \lambda / \partial A_{ij} = v_i v_j$) 를 활용하며, 희소 패턴에서의 외적만 필요합니다.
• 복잡도: 이 프레임워크는 autograd 그래프를 $O(1)$ 노드로 줄이고 메모리를 솔버 반복 횟수 ($k$) 와 무관하게 $O(n + \text{nnz})$로 줄입니다.

3. 헤일로 교환이 포함된 분산 레이어

멀티 GPU 실행을 위해 torch-sla 는 각 프로세스가 행 블록과 헤일로 메타데이터를 소유하는 도메인 분할을 구현합니다.

• 순전파: 프로세스들은 로컬 희소 행렬 - 벡터 (SpMV) 연산 전에 헤일로 값을 교환합니다.
• 역전파: 라이브러리는 전치 헤일로 교환을 구현합니다. 순전파는 소유된 값을 이웃 헤일로에 분산시키는 반면, 역전파는 들어오는 헤일로 기울기를 소유 프로세스의 경계 항목으로 다시 합산합니다.
• 구현: 순전파 및 역전파 통신 모두 torch.distributed를 사용하여 torch.autograd.Function 원시 연산으로 래핑되어 분산 풀이 전반에 걸친 엔드투엔드 미분 가능성을 보장합니다.

주요 기여

1. 통합 API: 다섯 가지 백엔드 (SciPy, Eigen, cuDSS, CuPy, PyTorch 네이티브) 에 걸쳐 직접, 반복, 비선형, 고유값 솔버를 위한 autograd 인식 인터페이스를 제공하는 단일 라이브러리.
2. 확장 가능한 수반 프레임워크: 반복 솔버 역전파의 메모리 병목 현상을 제거하여 수백만 개의 미지수를 가진 문제의 학습을 가능하게 하는 $O(1)$-그래프 미분 방법.
3. 분산 Autograd: 전치 헤일로 교환을 통해 엔드투엔드 미분을 기본적으로 지원하는 PyTorch 내의 첫 번째 분산 희소 솔버 구현.
4. 확장성: 공통 인터페이스를 구현함으로써 외부 라이브러리 (예: PETSc, Trilinos, hypre) 또는 학습된 전구조건자를 통합할 수 있는 백엔드 추상화.

실험 결과

저자들은 NVIDIA H200 GPU 를 사용하여 2D 푸아송 벤치마크에서 torch-sla 를 평가했습니다:

• 확장성: 이 라이브러리는 단일 GPU 에서 169M DOF까지, 그리고 세 개의 GPU 에 걸쳐 400M DOF까지 확장됩니다.
• 성능: 1M DOF 문제의 경우, torch-sla CG 백엔드 (190ms) 는 cuDSS(7.3 초) 보다 38 배 빠른 성능을 보입니다. 2M DOF 이상에서는 직접 솔버가 메모리 부족으로 실행되지만, 반복 솔버는 거의 선형적인 확장성을 유지합니다.
• 메모리 효율성: 1,000 번의 CG 반복에 대한 순전파와 수반 방법 비교에서, 순전파 방식은 약 64GB 를 필요로 했으며 (2,000 번 반복 시 OOM 발생), 수반 방법은 일정한 약 328MB 사용량을 유지하여 195 배의 메모리 감소를 달성했습니다.
• 기울기 검증: 비선형 및 고유값 문제에 대해 유한 차분과 비교하여 기울기를 검증한 결과, 상대 오차가 $10^{-5}$ 미만으로 나타났습니다.
• 엔드투엔드 애플리케이션: 푸아송 방정식에서 공간적으로 변하는 전도도를 복원하는 역계수 학습 작업은 표준 PyTorch Adam 루프를 사용하여 1 분 이내에 0.23% 의 상대 L2 오차로 성공적으로 수렴했습니다.

의의

이 논문은 torch-sla 가 PyTorch 생태계의 오랜 격차를 해소하여 JAX 의 jax.lax.custom_linear_solve 및 jax.scipy.sparse.linalg 와 유사한 기능을 통합된 PyTorch 라이브러리 내에서 제공한다고 주장합니다. 이를 통해 연구자들은 다음과 같은 작업을 수행할 수 있습니다:

• 기울기를 포기하거나 코드를 JAX 로 포팅하지 않고 희소 선형 풀이가 포함된 모델을 학습합니다.
• 분산 GPU 클러스터에서 수백만 개의 자유도를 가진 과학적 머신러닝 워크로드를 확장합니다.
• 희소 솔버를 미분 가능한 레이어로 취급하여 학습된 전구조건자 연구 및 물리 시뮬레이션의 엔드투엔드 최적화를 용이하게 합니다.

저자들은 현재 PyTorch 네이티브 백엔드에 대한 야코비 전구조건화에 의존한다는 점 (이는 매우 큰 시스템의 수렴 속도를 제한함) 과 2D 푸아송 벤치마크에 초점을 맞췄다는 한계를 지적하며, 3D 편미분방정식 (PDE) 과 부정정 시스템은 향후 과제로 식별했습니다.