분산 학습(Distributed training)은 머신러닝과 딥러닝에서 모델을 GPU나 서버 전체와 같은 여러 컴퓨팅 장치에 걸쳐 협력적으로 학습시키는 방법이다. 이 접근 방식은 수십억 또는 수조 개의 매개변수를 가지며 방대한 양의 데이터를 필요로 하는 대규모 언어 모델과 같은 현대 인공지능 시스템에 필수적이다. 계산 작업량을 분산함으로써 학습 시간을 수개월에서 수일 또는 수시간으로 단축할 수 있으며, 단일 장치에 맞출 수 없는 크기로 모델을 확장할 수 있다. 이 관행은 수십 년간 고성능 컴퓨팅의 초석이 되어 온 병렬 컴퓨팅의 원리를 활용하며, 프로세서 속도의 물리적 한계와 모델 크기의 증가로 인해 AI 분야에서 주류적인 필수 요소가 되었다.
분산 학습의 필요성은 메모리 용량과 계산 속도라는 두 가지 주요 제약에서 발생한다. 단일 GPU 또는 TPU는 유한한 메모리를 가지므로 저장하고 학습할 수 있는 모델의 크기가 제한된다. 마찬가지로 단일 장치에서 수백만 개의 학습 예제를 처리하는 데 필요한 시간은 지나치게 길 수 있다. 분산 학습은 모델과 데이터를 여러 장치에 분할하여 병렬 처리를 가능하게 함으로써 이러한 문제를 해결한다. 그러나 이는 통신, 동기화, 내결함성과 관련된 복잡성을 도입하며, 이는 이 분야의 핵심 과제이다. 이론적 속도 향상은 병렬화할 수 없는 작업 부하의 비율(예: 통신 오버헤드 및 순차적 종속성)에 의해 최대 개선이 제한된다는 Amdahl의 법칙에 의해 제한된다.
데이터 병렬 처리
데이터 병렬 처리는 가장 널리 사용되는 분산 학습 형태이다. 이 접근 방식에서는 모델이 각 장치에 복제되고 학습 데이터 세트가 더 작은 배치로 분할되며, 각 장치는 서로 다른 데이터 하위 집합을 동시에 처리한다. 각 장치가 로컬 그래디언트를 계산한 후, 이러한 그래디언트는 모든 장치에 걸쳐 집계되어 모델 매개변수를 업데이트한다. 이는 일반적으로 모든 장치의 그래디언트를 합산하고 결과를 브로드캐스트하는 all-reduce와 같은 집합 통신 연산을 사용하는 동기화 단계를 필요로 한다.
데이터 병렬 처리의 주요 장점 중 하나는 단순성과 확장성이다. 트랜스포머 기반 모델을 포함한 모든 모델 아키텍처에 상당한 수정 없이 적용할 수 있다. 그러나 장치 수가 증가함에 따라 통신 오버헤드가 병목 현상이 될 수 있으며, 특히 매개변수 수가 많은 모델의 경우 더욱 그렇다. 그래디언트 압축, 비동기 업데이트, 로컬 그래디언트 누적과 같은 기법이 이 문제를 완화하기 위해 개발되었다. PyTorch 및 TensorFlow와 같은 프레임워크는 데이터 병렬 처리를 위한 내장 지원을 제공하여 실무자가 쉽게 사용할 수 있게 한다.
모델 병렬 처리
모델 병렬 처리는 모델이 단일 장치의 메모리에 맞추기에는 너무 클 때 사용된다. 이 접근 방식에서는 모델의 서로 다른 부분이 서로 다른 장치에 배치되고, 데이터는 모델을 통해 순차적으로 흐르며 각 장치는 순방향 및 역방향 패스의 일부를 계산한다. 이는 NVIDIA나 AMD의 가장 진보된 GPU의 메모리 용량을 초과하는 수천억 개의 매개변수를 가질 수 있는 대규모 언어 모델과 특히 관련이 있다.
모델 병렬 처리는 각 장치가 레이어 하위 집합을 처리하는 레이어별 분할과 단일 레이어의 계산을 여러 장치에 분할하는 레이어 내부 분할을 포함한 여러 방식으로 구현될 수 있다. 후자는 어텐션 메커니즘과 피드포워드 네트워크를 병렬화할 수 있는 트랜스포머 모델에서 일반적이다. 그러나 모델 병렬 처리는 이전 레이어가 있는 장치가 이후 레이어가 계산되는 동안 유휴 상태일 수 있으므로 불균형한 활용을 초래하는 경우가 많다. 이는 장치 간에 계산을 중첩하는 파이프라인 병렬 처리를 통해 완화할 수 있다.
파이프라인 병렬 처리
파이프라인 병렬 처리는 데이터 및 모델 병렬 처리의 요소를 결합한 하이브리드 접근 방식이다. 이 방법에서는 모델이 단계로 나뉘고 각 단계가 서로 다른 장치에 할당된다. 학습 데이터는 마이크로 배치로 처리되며, 이는 파이프라인을 통해 지그재그 방식으로 흘러 여러 장치가 서로 다른 마이크로 배치를 동시에 작업할 수 있게 한다. 이는 순수 모델 병렬 처리와 관련된 유휴 시간을 줄이고 하드웨어 활용도를 향상시킨다.
파이프라인 병렬 처리의 주목할 만한 예로는 2019년 Google 연구자들이 도입한 GPipe 프레임워크가 있으며, 이는 이 기법을 사용하여 대규모 모델을 효율적으로 학습할 수 있음을 입증했다. 또 다른 변형으로는 Microsoft가 개발한 PipeDream이 있으며, 비동기 업데이트를 사용하여 처리량을 더욱 향상시킨다. 파이프라인 병렬 처리는 매우 깊은 모델에 특히 효과적이지만, 파이프라인 일정 관리와 단계 간 통신 처리에 어려움을 도입한다. 단계 수와 마이크로 배치 크기의 선택은 성능에 큰 영향을 미칠 수 있다.
집합 통신
집합 통신은 분산 학습의 중추로, 장치 간 데이터 교환과 계산 동기화를 가능하게 한다. 가장 일반적인 연산에는 all-reduce, all-gather, broadcast, reduce-scatter가 포함된다. 이러한 연산은 InfiniBand 및 Ethernet과 같은 고속 상호 연결에 최적화된 NVIDIA의 NCCL(NVIDIA Collective Communications Library) 및 MPI(Message Passing Interface)와 같은 라이브러리에서 구현된다.
데이터 병렬 처리에서 all-reduce 연산은 그래디언트를 집계하는 데 사용된다. 예를 들어, N개의 장치가 있는 경우 각 장치는 그래디언트 벡터를 계산하고, all-reduce는 모든 장치에 걸쳐 요소별 합을 계산한 다음 결과를 각 장치에 다시 분배한다. 이 연산은 통신 시간을 줄이는 트리 기반 또는 링 기반 알고리즘을 사용하여 최적화할 수 있다. 통신 토폴로지의 선택과 네트워크 대역폭은 분산 학습의 확장성에 중요한 요소이다. 2024년 기준으로 CoreWeave 또는 Amazon Web Services가 제공하는 수천 개의 GPU를 갖춘 클러스터는 통신 오버헤드를 최소화하기 위해 고대역폭 상호 연결에 의존한다.
동기 및 비동기 학습
분산 학습은 동기 및 비동기 접근 방식으로 분류할 수 있다. 동기 학습에서는 모든 장치가 로컬 데이터에 대한 그래디언트를 계산한 다음 모델을 업데이트하기 전에 다른 모든 장치가 완료될 때까지 기다린다. 이는 모델이 모든 장치에서 일관되게 유지되도록 보장하지만, 하드웨어 변동이나 네트워크 혼잡으로 인해 느린 장치인 지연 장치로 인해 속도가 저하될 수 있다. 동기 학습은 수렴 속성을 보장하므로 대부분의 딥러닝 프레임워크에서 표준이다.
반면 비동기 학습은 장치가 다른 장치를 기다리지 않고 독립적으로 모델을 업데이트할 수 있게 한다. 이는 처리량을 향상시킬 수 있지만, 장치가 오래된 모델 매개변수를 사용하여 수렴 속도가 느려지거나 불안정해질 수 있는 오래된 그래디언트의 위험을 도입한다. 그래디언트 오래됨 제한 및 탄성 평균화와 같은 기법이 이러한 문제를 해결하기 위해 제안되었다. 실제로 대부분의 대규모 학습 작업에서는 all-reduce를 사용한 동기 학습이 선호되며, 이는 단순성과 성능 사이의 좋은 균형을 제공한다.
하드웨어 및 인프라
분산 학습은 높은 성능을 달성하기 위해 특수 하드웨어와 인프라를 필요로 한다. 가장 일반적인 설정은 각각 여러 GPU를 갖춘 서버 클러스터로, 고속 네트워크로 연결된다. NVIDIA와 같은 회사는 GPU 시장을 지배하며 A100 및 H100 시리즈가 AI 학습에 널리 사용된다. AMD는 MI250X와 같은 경쟁력 있는 GPU를 제공하며, Intel은 Gaudi 가속기로 이 분야에 진입했다. Amazon Web Services, Azure, Google Cloud와 같은 클라우드 제공업체는 사용자가 주문형으로 분산 학습 클러스터를 임대할 수 있는 관리형 서비스를 제공하여 사내 인프라의 필요성을 줄인다.
GPU 외에도 Cerebras의 웨이퍼 스케일 엔진과 Graphcore의 IPU(Intelligence Processing Unit)와 같은 특수 하드웨어가 분산 학습을 가속화하기 위해 개발되었다. 이러한 시스템은 종종 외부 통신의 필요성을 줄이는 온칩 통신 네트워크를 특징으로 한다. 또한 상호 연결 기술의 선택이 중요하다. InfiniBand는 낮은 지연 시간과 높은 대역폭을 제공하는 반면, Ethernet은 비용 효율적이지만 더 느리다. 2023년 기준으로 GPT-4와 같은 일부 대규모 AI 학습 실행은 수만 개의 GPU를 사용한 것으로 추정되며, 이는 현대 분산 학습의 규모를 강조한다.
소프트웨어 프레임워크 및 기법
분산 학습 구현을 단순화하기 위해 여러 소프트웨어 프레임워크가 개발되었다. PyTorch의 DistributedDataParallel(DDP)은 데이터 병렬 처리를 위한 간단한 API를 제공하는 인기 있는 선택이다. TensorFlow는 MirroredStrategy 및 MultiWorkerMirroredStrategy를 포함한 다양한 전략을 지원하는 tf.distribute 모듈을 제공한다. Uber가 개발한 Horovod와 Microsoft가 개발한 DeepSpeed와 같은 더 고급 라이브러리는 그래디언트 압축, 혼합 정밀도 학습, 그리고 옵티마이저 상태, 그래디언트 및 매개변수를 장치 간에 분할하여 메모리 사용을 줄이는 ZeRO(Zero Redundancy Optimizer)와 같은 추가 최적화를 제공한다.
이러한 프레임워크는 분산 통신의 많은 복잡성을 추상화하여 연구자가 모델 개발에 집중할 수 있게 한다. 그러나 디버깅 및 성능 최적화를 위해 기본 원리를 이해하는 것은 여전히 중요하다. 예를 들어, 올바른 배치 크기, 학습률 일정 및 통신 백엔드를 선택하는 것은 학습 효율성에 큰 영향을 미칠 수 있다. 모델이 계속 성장함에 따라 분산 학습 알고리즘과 하드웨어의 혁신은 인공지능 분야의 발전에 중요한 역할을 할 것이다.
과제 및 향후 방향
성공에도 불구하고 분산 학습은 여러 과제에 직면해 있다. 통신 오버헤드는 특히 수십억 개의 매개변수를 가진 모델에서 주요 병목 현상으로 남아 있다. 그래디언트 압축 및 저정밀 통신과 같은 기법이 전송되는 데이터 양을 줄이기 위해 탐구되고 있다. 내결함성도 또 다른 문제로, 장치나 네트워크 링크의 장애가 학습을 중단시킬 수 있다. 체크포인팅과 장치 수를 동적으로 조정하는 탄력적 학습은 활발한 연구 분야이다.
미래를 바라보면 OpenAI와 Google DeepMind가 개발 중인 것과 같은 더 큰 모델로의 추세는 더 효율적인 분산 학습 방법의 필요성을 계속 주도할 것이다. 데이터, 모델 및 파이프라인 병렬 처리를 결합한 3D 병렬 처리의 등장은 이미 대규모 학습 실행에서 사용되고 있다. 또한 AWS Trainium 및 Google Cloud TPU와 같은 특수 하드웨어의 개발은 분산 학습을 더 비용 효율적이고 접근 가능하게 만드는 것을 목표로 한다. 분야가 발전함에 따라 분산 학습의 원리는 인공지능의 진보에 기초적인 역할을 계속할 것이다.