분산 학습 3가지 전략 정리: 데이터·파이프라인·텐서 병렬이 대형 AI 모델을 학습시키는 법

핵심 메시지

• 영상은 GPT-4를 엔비디아 A100 한 장으로 학습시키면 약 6000년이 걸리지만 GPU 2만 5000장에 분산해 90일 만에 끝냈다는 예시로 시작한다
• 데이터 병렬은 모델을 각 GPU에 그대로 복제하고 데이터만 나눠 처리한 뒤 기울기를 평균 내 동기화한다
• 모델이 GPU 한 장의 메모리에 들어가지 않으면 층 단위로 세로로 자르는 파이프라인 병렬을, 행렬 연산 단위로 가로로 자르는 텐서 병렬을 쓴다
• 병렬화의 최대 적은 GPU 간 통신 시간으로, 동기화 횟수를 줄이는 것이 성능의 핵심이다
• 실제 대규모 학습은 세 전략을 함께 쓰며, 최적 조합은 모델 구조와 클러스터 네트워크 구성에 따라 크게 달라진다

쉽게 이해하기

발표자는 규모의 감각을 먼저 잡아준다. GPT-4를 엔비디아 A100 한 장으로 학습시키면 약 6000년이 걸리지만, 실제로는 2만 5000장에 나눠 90일 만에 끝났다는 것이다. 이 격차를 만드는 것이 데이터 병렬, 파이프라인 병렬, 텐서 병렬이라는 세 가지 분산 학습 전략이다.

가장 기본은 데이터 병렬이다. 모델 전체를 각 GPU에 똑같이 복제해 두고 데이터셋만 나눠 각자 다른 배치로 순전파와 역전파를 수행한 뒤, 옵티마이저가 값을 갱신하기 전에 각 GPU의 기울기를 모아 평균 내고 다시 돌려주는 동기화를 거친다. 이 통신 방식을 올리듀스(all-reduce)라 하고, 파이토치는 엔비디아의 통신 라이브러리 NCCL로 처리한다. 구현은 스레드 기반이라 기기 한 대에 묶이는 DP와 다중 프로세스로 여러 대까지 확장되는 DDP로 갈리며, 발표자는 허깅페이스 accelerate로 코드 네 줄만 고쳐 RTX 3090 두 장에서 약 40% 속도 향상을 얻었다. 다만 데이터 병렬에는 모델이 GPU 한 장의 메모리에 통째로 들어가야 한다는 결정적 제약이 있다.

그래서 모델 자체를 자른다. 파이프라인 병렬은 2018년 GPipe 논문에서 나온 방식으로, 층을 기준으로 모델을 세로로 잘라 앞부분은 첫 GPU, 뒷부분은 다음 GPU에 올리고 중간 활성값만 넘긴다. 문제는 첫 GPU가 계산하는 동안 나머지가 놀고 있다는 점으로, GPU가 8장이면 매 순간 7장이 대기한다. 해결책이 마이크로 배치로, 배치를 잘게 쪼개 각 GPU가 서로 다른 조각을 동시에 처리하도록 겹쳐 유휴 구간(버블)을 줄이며, 순전파와 역전파를 교차시키는 1F1B 기법은 이 버블을 더 줄인다.

텐서 병렬은 같은 모델을 가로로 자른다. 메가트론 논문에서 나온 방식으로, 행렬 곱을 이루는 내적 연산들이 서로 독립이라는 점을 이용한다. 오른쪽 행렬을 열 기준으로 나누면 결과를 이어 붙이는 올개더(all-gather)로, 행 기준으로 나누면 원소별로 더하는 올리듀스로 합친다. 핵심은 여러 층에서 열 방향 분할과 행 방향 분할을 번갈아 이으면 입력을 처음에 한 번 나누고 마지막에 한 번만 합치면 되므로, 통신 횟수가 줄어드는 만큼 학습 속도에 큰 이득이 된다는 점이다.

실전에서는 셋을 함께 쓴다. GPU 8장이라면 먼저 파이프라인 병렬로 세로로 자르고, 텐서 병렬로 가로로 자른 뒤, 그렇게 만든 덩어리를 데이터 병렬로 복제하는 식이다. 최적 조합은 모델 구조와 클러스터의 네트워크 구성에 따라 달라지는데, 갤바트론 같은 연구는 동적 계획법으로 탐색 공간을 줄여 좋은 조합을 자동으로 찾는다. 서로 멀리 떨어진 노드 사이에는 통신량이 가장 적은 파이프라인 병렬을 쓰고, 모델 파라미터가 작으면 데이터 병렬, 크면 텐서 병렬이 유리하다는 것이 그 연구가 정리한 경험칙이다. FSDP, 딥스피드, 메가트론-LM 같은 라이브러리가 이 전략들을 감싸 제공한다.

주요 인사이트

• 분산 학습의 진짜 병목은 계산이 아니라 통신이다. GPU가 서로 대화하는 시간은 연산하지 않는 시간이므로, 동기화 횟수를 줄이는 설계가 곧 속도다.
• GPU를 두 장으로 늘려도 속도가 두 배가 되지는 않는다. 영상의 실측에서는 약 40% 개선에 그쳤고, 차이는 동기화와 데이터 분할 비용에서 온다.
• 데이터 병렬은 구현이 쉽고 확장도 매끄럽지만 '모델이 GPU 한 장에 들어가야 한다'는 제약을 넘지 못한다. 이 지점이 모델 자체를 쪼개는 전략으로 넘어가는 분기점이다.
• 텐서 병렬에서 열 분할과 행 분할을 번갈아 배치하면 층마다 합치지 않고 처음과 끝에서만 통신하면 된다. 나누는 방향을 설계하는 것만으로 통신 비용이 달라진다.
• 최적 병렬 구성은 모델과 하드웨어 구성마다 다르다. 그래서 사람이 손으로 고르는 대신 탐색 공간을 줄여 자동으로 찾는 연구가 별도 분야로 자리 잡았다.

자주 묻는 질문

관련 AI 소식