모델 병렬화는 머신 러닝 및 딥 러닝에서 단일 장치(예: GPU 또는 TPU)의 메모리에 맞지 않는 대규모 신경망을 훈련하거나 추론하는 데 사용되는 분산 컴퓨팅 기법이다. 각 장치가 모델의 전체 사본을 보유하고 서로 다른 데이터 배치를 처리하는 데이터 병렬화와 달리, 모델 병렬화는 모델 자체를 분할하며, 일반적으로 서로 다른 레이어 또는 레이어 하위 집합을 서로 다른 장치에 할당한다. 이 접근 방식은 수십억 또는 수조 개의 매개변수를 가질 수 있는 Transformer (architecture) 기반 대규모 언어 모델과 같은 모델을 확장하는 데 필수적이다.
개요
모델 병렬화에서는 신경망이 깊이(레이어 단위)를 따라 분할되거나, 더 고급 형태에서는 너비 또는 개별 텐서 연산과 같은 다른 차원을 따라 분할된다. 가장 간단한 형태는 레이어별 분할로, 연속된 레이어가 서로 다른 장치에 배치된다. 순방향 전파 중에는 활성화가 한 장치에서 다음 장치로 흐르고, 역방향 전파 중에는 기울기가 반대 방향으로 흐른다. 이는 파이프라인과 유사한 실행을 만들며, 신중하게 스케줄링하지 않으면 유휴 시간(버블)이 발생할 수 있다. 모델 병렬화는 종종 데이터 병렬화와 대조되지만, 대규모 모델 훈련에 사용되는 3D 병렬화와 같은 하이브리드 접근 방식에서 두 가지를 결합할 수 있다.
동기
현대 신경망, 특히 대규모 언어 모델은 기하급수적으로 크기가 증가했다. 예를 들어, GPT-3(2020)는 1750억 개의 매개변수를 가지며, 32비트 정밀도에서 가중치만으로도 350GB 이상의 메모리가 필요하다. 이는 고성능 GPU인 A100의 80GB 용량을 훨씬 초과한다. 혼합 정밀도 및 기울기 체크포인팅과 같은 메모리 최적화를 사용하더라도 단일 장치 훈련은 불가능하다. 모델 병렬화는 고속 상호 연결로 연결된 장치 클러스터에 모델을 분산하여 단일 장치가 보유할 수 있는 것보다 훨씬 큰 모델을 연구자들이 훈련할 수 있게 한다.
모델 병렬화의 유형
레이어별(파이프라인) 병렬화
이것은 가장 간단한 형태로, 모델을 순차적 단계로 나누고 각 단계를 서로 다른 장치에 할당한다. 예를 들어, 100개 레이어 네트워크를 각각 10개 레이어로 구성된 10개 단계로 나누고 각 단계를 별도의 GPU에 배치할 수 있다. 순방향 패스 중에는 데이터가 1단계, 2단계, 그리고 그 다음으로 흐른다. 주요 과제는 부하 균형과 파이프라인 버블 감소이다. GPipe 및 PipeDream과 같은 기술은 마이크로 배치 및 스케줄링을 도입하여 활용도를 개선한다.
텐서 병렬화
텐서 병렬화는 개별 연산(예: 행렬 곱셈)을 여러 장치에 분할한다. 예를 들어, 트랜스포머의 어텐션 메커니즘에서 쿼리, 키, 값 투영을 GPU에 분할하고 결과를 all-reduce 연산으로 결합할 수 있다. 이 접근 방식은 Megatron-LM에서 사용되며 장치당 메모리를 줄이면서 높은 계산 효율성을 유지하는 데 효과적이지만, 고대역폭 통신이 필요하다.
전문가 병렬화
혼합 전문가(MoE) 모델에서 사용되는 전문가 병렬화는 서로 다른 전문가 네트워크를 서로 다른 장치에 배치하고, 라우터(게이팅 네트워크)가 토큰을 적절한 전문가에게 분배한다. 이는 토큰당 전문가의 하위 집합만 활성화되므로 계산 증가 없이 방대한 매개변수 수를 가능하게 하는 모델 병렬화의 한 형태이다.
구현 과제
모델 병렬화는 여러 가지 과제를 도입한다:
- 통신 오버헤드: 장치는 활성화 및 기울기를 교환해야 하며, 특히 느린 상호 연결에서는 병목 현상이 될 수 있다.
- 부하 불균형: 레이어 크기나 계산 요구 사항이 고르지 않으면 일부 장치는 유휴 상태가 되고 다른 장치는 작업을 수행한다.
- 파이프라인 버블: 파이프라인 병렬화에서 장치는 이전 단계의 데이터를 기다려야 할 수 있어 활용도가 감소한다.
- 메모리 단편화: 모델 분할은 메모리 사용량이 고르지 않게 만들 수 있어 신중한 배치와 스케줄링이 필요하다.
- 내결함성: 장치가 많을수록 실패 확률이 증가하므로 체크포인팅 및 복구 메커니즘이 필요하다.
데이터 병렬화와의 비교
데이터 병렬화에서는 각 장치가 모델의 전체 사본을 보유하고 서로 다른 미니 배치의 데이터를 처리한다. 각 단계 후에 기울기가 장치 간에 평균화된다. 이는 단일 장치 메모리에 맞는 모델에 대해 간단하고 확장성이 좋다. 반면 모델 병렬화는 모델 자체가 너무 클 때 필요하다. 그러나 모델 병렬화는 직렬 종속성으로 인해 통신 비용이 더 높고 효율성이 낮은 경우가 많다. 실제로 대규모 훈련에서는 노드 간 데이터 병렬화와 노드 내 모델 병렬화를 모두 사용한다.
응용 분야
모델 병렬화는 GPT-4, PaLM, LLaMA와 같은 대규모 언어 모델을 훈련하고 프로덕션에서 서비스하는 데 중요하다. OpenAI, Google(Google Cloud를 통해), Anthropic과 같은 기업은 모델 병렬화에 의존하여 모델을 훈련하고 배포한다. 또한 단일 가속기에 너무 큰 모델이 있는 컴퓨터 비전(예: 3D CNN) 및 과학 컴퓨팅과 같은 다른 영역에서도 사용된다.
소프트웨어 지원
여러 프레임워크가 모델 병렬화에 대한 기본 지원을 제공한다:
- PyTorch: 파이프라인 병렬화(예:
torch.distributed.pipeline.sync.Pipe) 및tensor_parallel라이브러리를 통한 텐서 병렬화를 포함한torch.distributed를 제공한다. - TensorFlow: 모델 병렬화를 통합할 수 있는
TPUStrategy및MultiWorkerMirroredStrategy와 같은 전략을 포함한tf.distribute를 제공한다. - Megatron-LM: 텐서 및 파이프라인 병렬화를 위한 NVIDIA의 특수 라이브러리이다.
- DeepSpeed: Microsoft에서 제공하며, 옵티마이저 상태, 기울기 및 매개변수를 분할하는 모델 병렬화 형태인 ZeRO(Zero Redundancy Optimizer)를 제공한다.
- JAX:
jax.sharding을 통해 명시적 장치 배치 및 병렬 실행을 허용한다.
하드웨어 고려 사항
모델 병렬화는 고속 장치 간 통신에 크게 의존한다. 단일 노드 내에서는 NVLink와 PCIe가 빠른 연결을 제공하고, 노드 간에는 InfiniBand 또는 고속 이더넷이 사용된다. NVIDIA GPU 및 Google TPU와 같은 특수 하드웨어는 모델 병렬화를 용이하게 하기 위해 전용 상호 연결(예: NVSwitch, TPU 링크)을 포함하는 경우가 많다. Cerebras 및 SambaNova와 같은 기업은 모델 병렬화의 필요성을 줄이기 위해 대형 온칩 메모리를 갖춘 시스템을 설계했지만, 극단적인 규모에서는 분산 접근 방식이 여전히 필요하다.
최근 발전
최근 연구는 모델 병렬화의 효율성을 개선하는 데 초점을 맞추고 있다. 시퀀스 병렬화(시퀀스 차원 분할) 및 컨텍스트 병렬화(긴 시퀀스용)와 같은 기술이 개발되었다. Alpa 및 FlexFlow의 자동 모델 분할 알고리즘은 장치 간 연산 배치를 최적화한다. 또한 데이터, 텐서 및 파이프라인 병렬화를 결합한 하이브리드 병렬화는 NVIDIA와 Microsoft가 훈련한 Megatron-Turing NLG 모델(5300억 매개변수)에서 볼 수 있듯이 대규모 훈련에서 표준이 되었다.
한계 및 향후 방향
강력함에도 불구하고 모델 병렬화에는 한계가 있다. 통신 오버헤드는 작은 모델이나 느린 네트워크에서 지배적일 수 있다. 분할 및 스케줄링의 복잡성은 모델 크기와 클러스터 이질성에 따라 증가한다. 향후 방향에는 더 지능적인 자동 병렬화, 더 나은 통신 압축, 본질적으로 더 병렬화 가능한 새로운 아키텍처가 포함된다. 생성형 AI의 부상과 점점 더 커지는 모델과 함께 모델 병렬화는 확장 가능한 딥 러닝의 초석으로 남을 것이다.
같이 보기
참고 문헌
- Narayanan, D., et al. (2019). PipeDream: Generalized Pipeline Parallelism for DNN Training.
- Shoeybi, M., et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.
- Rajbhandari, S., et al. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.
- Huang, Y., et al. (2019). GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.