기술 요약: ShardTensor – 과학적 머신러닝을 위한 도메인 병렬화

1. 문제 제기

과학적 머신러닝 (SciML) 은 GPU 메모리 제약으로 인해 고해상도 과학 데이터를 훈련하고 추론할 수 없다는 치명적인 병목 현상에 직면해 있습니다. 대규모 언어 모델 (LLM) 은 메모리가 주로 파라미터 수에 의해 지배되는 반면, SciML 워크로드는 종종 **중간 활성화 (intermediate activations)**에 의해 제약받습니다.

이 논문은 과학 데이터에서 "차원의 저주 (computational curse of dimensionality)"를 식별합니다. LLM 에서 텍스트 길이를 두 배로 늘리면 입력 토큰이 약 두 배로 증가하는 반면, $N$차원 과학 데이터의 해상도를 두 배로 늘리면 데이터 크기는 $2^N$배로 증가합니다. 3 차원 이상의 데이터 (예: 유체 역학, 의료 영상, 기후 모델링) 의 경우, 훈련 중 역전파 단계에서 중간 활성화 상태를 저장하는 데 필요한 메모리는 기하급수적으로 증가합니다.

기존 솔루션들은 확장성이 부족하거나 정확도가 저하되는 경향이 있습니다:

• 정밀도 감소 (Reduced Precision): 메모리를 modest 하게 (2 배) 만 절약할 뿐이며, 과학적 대리 모델 (surrogates) 에서 동적 범위 문제를 겪을 수 있습니다.
• 공간 다운샘플링 (Spatial Downsampling): 메모리는 줄이지만 과학적 돌파구에 필요한 고충실도 해상도를 희생합니다.
• 모델 축소 (Model Reduction): 정확도를 떨어뜨립니다.
• 활성화 체크포인팅/CPU 오프로딩: 효과적이지만 상당한 계산 오버헤드 (재계산 또는 호스트 - 장치 간 전송) 를 초래합니다.
• 표준 병렬화: 데이터 병렬화 (배치 분할) 는 배치 크기가 1 일 때 하드 리미트에 부딪힙니다. 모델 병렬화 (가중치 분할) 는 파라미터를 위해 설계된 것으로, 입력 데이터용이 아니며 동적 입력 데이터 처리 (예: 컨볼루션으로 인한 전역 형태 변경) 에 적용되지 않는 정적 텐서 형태를 가정합니다.

결과적으로, 장치당 배치 크기 1 미만의 입력 데이터를 병렬화할 수 있는 일반화된 프레임워크가 없어, SciML 이 과학자들이 요구하는 네이티브 해상도에 도달하는 것이 불가능합니다.

2. 방법론: ShardTensor

저자들은 입력 데이터의 공간 차원을 하드웨어 제약과 분리하는 새로운 패러다임인 **도메인 병렬화 (domain parallelism)**인 ShardTensor를 소개합니다. 이는 모델 가중치가 아닌 입력 데이터를 처리하도록 특별히 설계된 PyTorch 의 DTensor(Distributed Tensor) 추상화의 확장입니다.

핵심 설계 원칙

• 도메인 병렬화: 배치 차원을 넘어 입력 데이터에 대한 장치 간 병렬화로 정의됩니다. 배치 크기가 1 에 도달하면 ShardTensor 는 여러 GPU 에 걸쳐 높이, 너비, 깊이와 같은 공간 차원을 추가로 세분화합니다.
• 동적 형태 처리: 모델 가중치에 적합한 정적 분포 형태를 가정하는 DTensor 와 달리, ShardTensor 는 동적 형태를 지원합니다. 이는 텐서의 각 로컬 청크 형태를 설명하는 **샤딩 형태 (sharding shapes)**라는 네 번째 메타데이터 구성 요소를 도입합니다. 이를 통해 비구조적이거나 불규칙한 데이터를 임의의 청크로 분할할 수 있습니다.
• 명령형 디스패치 (Imperative Dispatch): 이 프레임워크는 정적이 아닌 명령형입니다. PyTorch 의 __torch_dispatch__ 및 __torch_function__ 메커니즘을 활용하여 연산을 동적으로 라우팅합니다. 이는 사전 컴파일된 계산 그래프나 레이어별 맞춤형 패치가 필요 없이 레이어 단위의 도메인 병렬화를 처리할 수 있게 합니다.
• 확장성: 사용자는 고수준 함수형 인터페이스 (@custom_op) 나 저수준 디스패치 로직 (aten 연산) 을 통해 시스템을 확장할 수 있습니다. 이는 진화하는 모델 아키텍처 및 커스텀 커널과의 호환성을 보장합니다.

구현 세부 사항

ShardTensor 는 PyTorch 의 DTensor 확장으로 구축되었습니다.

• 디스패치 아키텍처: Python 수준에서 연산을 가로챕니다. 특정 연산에 도메인 병렬 핸들러가 있으면 이를 실행하고, 그렇지 않으면 DTensor 시맨틱스나 표준 PyTorch 동작으로 폴백합니다.
• 통신: 이 프레임워크는 컨볼루션이나 정규화와 같은 연산에서 수치적 일관성을 유지하기 위해 필요한 집단 연산 (all-reduce, all-gather 등) 과 "할로 (halo)" 교환 (인접 장치에서 인접 픽셀 가져오기) 을 관리합니다.
• 통합: 기존 병렬화 전략 (데이터 병렬화 및 모델 병렬화) 과 조립 가능하도록 설계되어, 한 축에서는 배치를 분할하고 다른 축에서는 공간 차원을 분할하는 2 차원 이상의 병렬화 전략을 가능하게 합니다.

3. 주요 기여

1. ShardTensor 프레임워크: 배치 크기 1 미만의 입력 데이터에 대한 도메인 병렬화를 가능하게 하는 일반화된 오픈소스 프레임워크 (NVIDIA PhysicsNeMo 를 통해 제공).
2. 새로운 추상화: DTensor 와 같은 가중치 중심 병렬화 도구와 구별되는 동적 입력 텐서 분포를 처리하기 위한 "샤딩 형태" 도입.
3. 비침습적 프로그래밍 모델: 최소한의 코드 변경으로 모델 입력을 래핑할 수 있도록 설계되어, 병렬화를 위해 레이어나 모델을 다시 작성할 필요가 없음.
4. 조립성: 직교하는 메쉬 축에서 도메인 병렬화를 데이터 병렬화 및 모델 병렬화와 결합할 수 있음을 입증.

4. 결과 및 벤치마크

이 논문은 NVIDIA Blackwell (GB200) 및 H100 GPU 에서 합성 벤치마크와 실제 산업 응용을 통해 ShardTensor 를 검증했습니다.

• 링 어텐션 (Ring Attention): 매우 작은 시퀀스 크기에서는 병렬화 오버헤드가 지배적입니다. 그러나 매우 큰 시퀀스 크기에서는 훈련 및 추론 모두에서 GPU 수에 따른 **거의 선형적인 강한 확장 (near-linear strong scaling)**을 달성합니다.
• 비전 트랜스포머 (ViT):
  • 강한 확장: 고해상도 (예: $4096^2$) 에서 추론은 단일 GPU 대비 16 GPU 에서 15 배 더 빠릅니다. $2048^2$에서의 훈련은 8 GPU 로 5 배 더 빠릅니다.
  • 메모리 효율성: 메모리 사용량은 2 차원 데이터의 경우 이차적으로, 3 차원 데이터의 경우 삼차적으로 증가하여 중간 활성화가 지배적임을 확인했습니다. ShardTensor 를 통한 강한 확장은 장치당 메모리를 비례적으로 줄여, 표준 하드웨어에서 10 억 개 이상의 입력 포인트를 가진 3D 모델 훈련을 가능하게 합니다.
• Transolver (자동차 공기역학):
  • DrivaerML 데이터셋을 기반으로 훈련되었으며, 도메인당 최대 120 만 개의 포인트를 처리했습니다.
  • 도메인 크기가 2 배, 4 배, 8 배로 증가함에 따라 시스템은 수치적 안정성을 유지하고 단일 GPU 베이스라인 대비 경쟁력 있는 정확도 (압력 및 속도에 대한 L2 오차) 를 유지했습니다.
  • 전체 전처리 파이프라인도 병렬화되어 엔드투엔드 확장성을 보장했습니다.
• StormScope (기상 예보):
  • 3 km 해상도($1024 \times 1792$ 그리드) 의 대륙 규모 모델 (CONUS).
  • 피크 메모리 사용량 (114 GB) 은 단일 H100 GPU 의 용량 (80 GB) 을 초과했습니다.
  • ShardTensor 는 16 개의 데이터 병렬 그룹 내에서 2 개의 GPU 에 걸쳐 활성화 상태를 분할하여 안정적인 훈련을 가능하게 했습니다. 모델은 안정적으로 수렴하여 저해상도 (6 km) 단일 GPU 베이스라인의 손실 궤적을 추적했으며, 도메인 범위를 희생하지 않고 이전에 접근 불가능했던 개별 대류성 폭풍을 해상도할 수 있게 했습니다.

5. 중요성 및 주장

이 논문은 ShardTensor 가 과학적 머신러닝의 근본적인 병목 현상, 즉 과학자가 요구하는 네이티브 해상도의 데이터에 대한 훈련 및 추론 불가능성을 해결한다고 주장합니다.

• 해상도 제한 워크로드의 해결: 메모리 한계로 인해 이전에 불가능했던 계산을 처리 가능한 것으로 전환하여, 체적 의료 영상, 계산 유체 역학, 기후 모델링 분야에서 고충실도 훈련을 가능하게 합니다.
• 조립성: 데이터 양, 모델 크기, 해상도 간의 트레이드오프를 제거합니다. 과학자들은 이제 여러 병렬화 축을 동시에 확장할 수 있습니다.
• 접근성: 비침습적 프로그래밍 모델을 제공함으로써 분산 시스템 전문가가 아닌 실무자의 진입 장벽을 낮추며, 이전의 맞춤형 아키텍처 특화 병렬화 노력과 대비됩니다.

인정된 한계

저자들은 겸손하게도 명령형, 레이어 단위 디스패치 모델이 Python 수준 오버헤드를 도입하여 작은 연산이나 저해상도 데이터에서는 효율성이 낮아질 수 있다고 지적합니다. 성능은 또한 인터커넥트 대역폭에 의존적입니다. 또한, 현재 모든 PyTorch 연산에 도메인 병렬 디스패치 경로가 있는 것은 아니며, 정적 그래프 최적화를 위한 torch.compile과의 완전한 통합은 진행 중입니다.

결론적으로, ShardTensor 는 오늘날 데이터 병렬화가 일상적인 것처럼 과학적 머신러닝에서 고도 도메인 병렬화를 일상화하는 단계로, 과학적 파운데이션 모델 시대를 촉진합니다.