AI VIDEO BRIEFING

대형 모델 분산 학습 정리: 데이터·파이프라인·텐서 병렬과 FSDP, ZeRO가 필요한 이유

대규모 언어 모델 학습 인프라를 만드는 엔지니어가 버클리 강연에서 풀어놓은 분산 학습의 기본기를 정리했다. 데이터·파이프라인·텐서 병렬의 원리와 한계, 그리고 FSDP와 제로가 메모리 벽을 어떻게 우회하는지 살펴본다.

출처: Sourish Kundu2024년 12월 27일AI 보조 요약

GPU 한 장으로는 못 만든다, 대형 모델 분산 학습의 세 가지 전략 영상 대표 이미지

핵심 메시지

AI 성능을 끌어올리는 축은 데이터, 알고리즘, 연산 세 가지인데 알고리즘의 도약은 5~6년에 한 번꼴이라 그 사이를 메우는 것은 연산 효율화다.
데이터 병렬은 GPU마다 모델을 복제해 서로 다른 배치를 학습시킨 뒤 기울기를 합치는 방식으로, 모델이 GPU 한 장에 들어가야만 쓸 수 있다는 한계가 있다.
파이프라인 병렬은 층을 나눠 담지만 GPU 상당수가 노는 버블이 생기고, 텐서 병렬은 행렬 자체를 쪼개 층 하나가 안 들어가는 문제를 푼다.
FSDP와 딥스피드의 제로는 옵티마이저 상태, 기울기, 파라미터를 차례로 분산 저장해 데이터 병렬의 메모리 한계를 우회한다.
통신 비용은 모든 병렬화의 숙적이지만, FSDP는 기존 통신 연산을 분해해 재조합하는 방식이라 추가 통신 부담이 거의 없다.

쉽게 이해하기

발표자는 대규모 언어 모델의 학습 인프라를 만드는 일을 하고 있다며, 연구자들이 연구에만 집중할 수 있도록 학습 기반을 다지는 것이 자기 일이라고 소개했다. 그가 보는 AI 성능 향상의 축은 세 가지다. 데이터의 양과 품질, 알고리즘의 혁신, 그리고 연산 자원이다. 2012년 알렉스넷과 2017년 트랜스포머처럼 판을 뒤집는 알고리즘은 대략 5~6년에 한 번 나오기 때문에, 그 사이 기간에 확실하게 개선할 수 있는 영역이 바로 연산을 얼마나 효율적으로 키우느냐라는 것이다.

GPU가 학습의 중심이 된 배경도 짚는다. 원래 그래픽 카드의 코어는 화면의 픽셀 색을 초당 수십 번씩 계산하기 위한 물건이었고, 연구자들은 처음에 이 그래픽 전용 파이프라인에 기계학습을 억지로 끼워 넣어야 했다. 2006년 쿠다가 나오면서 GPU에서 직접 도는 코드를 짜기 쉬워졌고, 딥러닝의 핵심 연산인 행렬 곱셈이 서로 독립적으로 계산 가능한 이른바 민망할 정도로 병렬적인 문제라는 점이 GPU와 맞아떨어졌다. 세 겹 반복문으로 짜던 코드가 수천 개의 인스턴스를 동시에 띄우는 한 겹으로 줄어든다.

문제는 모델이 커지면서 시작된다. 4천억 파라미터급 모델은 양자화한 상태로 추론만 돌려도 800GB가 넘는 메모리를 요구한다. 소비자용 GPU가 16~24GB, 데이터센터급이 80GB 수준이니 추론에만 최소 여덟 장이 필요하고, 옵티마이저 상태와 기울기, 데이터 적재까지 감당해야 하는 학습은 차원이 다른 문제가 된다. 여기서 첫 번째 전략인 데이터 병렬이 등장한다. GPU마다 모델을 복제하고 서로 다른 배치를 먹인 뒤, 역전파로 구한 기울기를 전체 평균으로 맞추는 동기화 단계가 핵심이다.

데이터 병렬의 치명적 한계는 모델 자체가 GPU 한 장에 들어가야 한다는 점이다. 그래서 모델을 자르는 방법이 나온다. 파이프라인 병렬은 층을 세로로 나눠 GPU마다 몇 개씩 맡기는데, 앞쪽 GPU가 계산하는 동안 뒤쪽은 놀게 되어 활용률이 크게 떨어진다. 이 유휴 구간을 버블이라 부르고, 배치를 더 잘게 쪼개는 마이크로 배칭으로 줄인다. 층 하나조차 메모리에 안 들어가는 경우에는 행렬을 열이나 행 단위로 쪼개 여러 GPU에 나눠 계산한 뒤 결과를 이어 붙이는 텐서 병렬이 답이 된다.

마지막으로 데이터 병렬을 살리면서 메모리 문제를 푸는 접근이 소개된다. 메타 계열의 FSDP와 마이크로소프트의 딥스피드가 대표적이며, 옵티마이저 상태만 나누는 단계, 기울기까지 나누는 단계, 파라미터까지 모두 나누는 단계로 최적화 수준이 나뉜다. 아담 옵티마이저가 파라미터마다 모멘텀과 분산 두 값을 더 들고 있어 메모리 사용이 세 배로 불어난다는 점을 떠올리면 왜 옵티마이저 상태 분산이 첫 단추인지 분명해진다. FSDP에서는 각 GPU가 자기 몫의 파라미터를 필요한 순간에 나머지에 뿌리고 쓰고 나면 지우는 식으로 순전파를 함께 진행한다.

주요 인사이트

세 가지 병렬화 전략은 서로의 단점을 메워 주는 관계라, 실제 대규모 학습에서는 하나를 고르는 것이 아니라 여러 축으로 동시에 쪼개는 조합이 쓰인다.
GPU 여러 장이 없어도 배울 방법은 있다. 발표자는 도커 컨테이너 하나하나를 별도 GPU처럼 취급해 데이터 병렬을 직접 구현해 보거나, 게이밍 PC 여러 대를 묶어 실험해 보라고 권한다.
만 장이 넘는 GPU로 학습을 돌리면 같은 모델의 카드끼리도 성능 편차가 눈에 띄게 나타난다. 대규모 학습에서는 알고리즘만큼이나 장애 감지와 이중화가 중요해진다.
요즘 GPU는 연산보다 메모리에서 먼저 막힌다. 플래시 어텐션이 계산량을 오히려 늘리면서도 실제 시간은 단축하는 이유가 여기에 있다.
이 분야의 성패는 결국 사용성이다. 연구자는 분산 학습 내부 사정에 관심이 없기 때문에, 몇 줄만 추가하면 여러 GPU로 확장되는 도구가 실제 채택을 좌우한다.

자주 묻는 질문

데이터 병렬만으로는 왜 대형 언어 모델을 학습시킬 수 없나?

데이터 병렬은 GPU마다 모델을 통째로 복제하는 방식이라 모델과 옵티마이저 상태, 기울기가 GPU 한 장의 메모리에 모두 들어가야 한다. 수천억 파라미터 모델은 추론만으로도 수백 GB가 필요해 이 조건을 만족할 수 없다.

파이프라인 병렬의 버블은 무엇이고 어떻게 줄이나?

층을 나눠 맡긴 GPU들이 자기 차례를 기다리며 노는 구간을 버블이라고 한다. 배치를 더 작은 마이크로 배치로 쪼개 앞뒤 GPU가 겹쳐 일하게 만들면 유휴 시간이 줄어들지만 완전히 없어지지는 않는다.

FSDP는 추가 통신 비용이 얼마나 드나?

발표자는 사실상 늘지 않는다고 설명한다. 데이터 병렬에서 쓰던 전체 합산 연산이 내부적으로 축소 분배와 전체 수집의 결합이기 때문에, 파라미터를 나눠 갖는 방식으로 바꿔도 통신량이 같은 수준에 머문다는 것이다.

GPU가 한 장뿐인데 분산 학습을 연습할 방법이 있나?

발표자는 도커 컨테이너 각각을 독립된 GPU처럼 다뤄 데이터 병렬을 직접 구현해 보는 방법, 그리고 친구들의 게이밍 PC를 묶어 실험하는 방법을 현실적인 대안으로 제시했다.

원문과 출처

이 글은 원본 영상의 자막을 바탕으로 한국어 독자를 위해 요약했습니다. 전체 맥락과 최신 정보는 원문에서 확인하세요.

YouTube 원본 영상 보기 ↗