GShard 논문 정리

한 줄 요약

GShard는 거대 모델을 여러 장치에 자동으로 나누어 학습시키는 시스템이다.

핵심은 두 가지다.

• 모델에는 가벼운 샤딩 주석만 붙인다.
• XLA 컴파일러가 실제 병렬 실행 계획을 만든다.

이 조합으로 저자들은 6,000억 파라미터 규모의 다국어 번역 모델을 2,048개 TPU v3 코어에서 4일 만에 학습했다.

결과도 강했다. 100개 언어를 영어로 번역하는 단일 모델에서 기존 이중언어 기준 모델보다 평균 BLEU가 크게 올랐다.

 

 

Crop 포인트: 모델 크기는 크게 늘어도 학습 비용은 훨씬 완만하게 증가하고, 품질은 꾸준히 상승하는 흐름에 주목하면 된다.

왜 중요한가

신경망은 커질수록 성능이 좋아지는 경우가 많다.

문제는 비용이다.

단순히 모델을 깊게 만들거나 넓게 만들면 다음 문제가 생긴다.

• 한 장치 메모리에 모델이 들어가지 않는다.
• 장치 간 통신이 병목이 된다.
• 병렬화 코드를 모델마다 다시 짜야 한다.
• 장치 수가 많아질수록 컴파일 그래프가 폭발한다.

GShard의 목표는 명확하다.

거대 모델을 만들되, 모델 개발자가 병렬화 세부 구현에 묶이지 않게 하는 것이다.

논문의 핵심 아이디어

GShard는 세 가지 설계를 결합한다.

1. 조건부 계산

모든 입력이 전체 모델을 통과하지 않는다.

각 토큰은 필요한 일부 전문가 네트워크만 사용한다.

즉, 모델 전체 크기는 매우 크지만 실제 계산량은 제한된다.

이 방식이 Mixture-of-Experts, 줄여서 MoE다.

2. 자동 샤딩

개발자는 텐서를 어떻게 나눌지 일부 힌트만 준다.

예를 들어 다음과 같은 의도를 표시한다.

• 이 텐서는 모든 장치에 복제한다.
• 이 텐서는 특정 차원으로 나눈다.
• 이 텐서는 장치 배치에 맞춰 더 일반적으로 나눈다.

그 뒤 컴파일러가 실제 통신과 연산 배치를 만든다.

3. 단일 프로그램 병렬화

전통적인 방식은 장치마다 다른 프로그램을 만들 수 있다.

하지만 장치가 수천 개가 되면 이 방식은 컴파일 단계부터 부담이 커진다.

GShard는 모든 장치가 같은 프로그램을 실행하게 만든다.

장치마다 맡는 데이터 조각만 다르다.

 

 

Crop 포인트: 왼쪽은 장치마다 연산 그래프가 늘어나는 구조이고, 오른쪽은 하나의 프로그램을 유지해 확장성을 확보하는 구조다.

모델 구조: Transformer에 MoE를 끼워 넣기

논문은 Transformer 번역 모델을 기반으로 한다.

일반 Transformer는 attention 층과 feed-forward 층을 반복한다.

GShard 모델은 여기서 일부 feed-forward 층을 MoE 층으로 바꾼다.

MoE 층 안에는 많은 전문가 네트워크가 있다.

각 토큰은 그중 일부만 선택해 통과한다.

이렇게 하면 파라미터 수는 크게 늘릴 수 있다.

반면 토큰 하나가 실제로 사용하는 계산 경로는 짧게 유지된다.

 

 

Crop 포인트: MoE 층만 여러 장치로 나뉘고, 나머지 Transformer 층은 주로 복제되는 구조가 핵심이다.

MoE 라우팅은 어떻게 동작하는가

MoE에서 가장 중요한 부분은 라우팅이다.

각 토큰을 어느 전문가에게 보낼지 정해야 한다.

논문은 토큰마다 가장 적합한 전문가 두 개를 후보로 둔다.

하지만 단순히 점수가 높은 전문가만 고르면 문제가 생긴다.

인기 전문가에게 토큰이 몰리고, 다른 전문가는 거의 학습되지 않는다.

그래서 GShard는 다음 장치를 쓴다.

• 전문가마다 처리 가능한 토큰 수를 제한한다.
• 배치를 여러 그룹으로 나누어 병렬로 라우팅한다.
• 특정 전문가 쏠림을 줄이는 보조 손실을 둔다.
• 두 번째 전문가 선택은 확률적으로 수행해 여유 용량을 아낀다.

이 설계의 의도는 단순하다.

전문가를 많이 두되, 모두가 어느 정도 학습되게 만드는 것이다.

샤딩 주석: 모델 코드와 병렬 실행의 분리

GShard의 프로그래밍 모델은 “거대한 단일 장치가 있는 것처럼 모델을 작성하라”에 가깝다.

개발자는 중요한 텐서에만 샤딩 정책을 붙인다.

그 뒤 컴파일러가 나머지를 추론한다.

이 분리는 실용적이다.

모델 구조를 바꿀 때마다 통신 코드를 전부 고칠 필요가 줄어든다.

MoE에서는 특히 중요하다.

입력 토큰은 데이터 병렬처럼 나눌 수 있다.

하지만 전문가 가중치는 너무 커서 모든 장치에 복제하기 어렵다.

따라서 일반 Transformer 층과 MoE 층은 서로 다른 병렬화 방식을 사용한다.

GShard는 이 전환을 컴파일러가 처리하게 한다.

 

 

Crop 포인트: 텐서 분할 축이 바뀌는 지점에서 AllToAll 같은 통신이 들어가며, 이 전환을 컴파일러가 자동으로 삽입한다.

컴파일러가 해결해야 하는 실제 문제

자동 샤딩은 단순한 텐서 자르기가 아니다.

연산마다 필요한 통신이 다르다.

예를 들어 다음 상황이 생긴다.

• 쪼갠 결과를 다시 합쳐야 한다.
• 텐서의 분할 축을 바꿔야 한다.
• 이웃 장치의 경계 데이터가 필요하다.
• 차원이 장치 수로 딱 나누어떨어지지 않는다.

GShard는 XLA 안에서 이런 경우를 다룬다.

핵심 통신 패턴은 다음과 같다.

• 일부 결과를 전체 장치에서 합산한다.
• 모든 장치의 조각을 모아 복제된 텐서를 만든다.
• 장치 간 조각을 서로 교환해 분할 축을 바꾼다.
• 이웃 장치와 경계 데이터만 주고받는다.

 

 

Crop 포인트: 전체 데이터를 복제하지 않고 경계에 필요한 일부 데이터만 교환해, 이웃 파티션 의존성을 처리하는 점을 보면 된다.

실험: 100개 언어를 영어로 번역

논문은 다국어 기계번역을 검증 과제로 선택했다.

이 과제는 거대 모델 실험에 적합하다.

이유는 두 가지다.

• 고자원 언어는 충분한 모델 용량을 요구한다.
• 저자원 언어는 다른 언어에서 오는 전이를 필요로 한다.

모델이 작으면 고자원 언어에서 병목이 생긴다.

반대로 공유가 부족하면 저자원 언어에서 전이가 약해진다.

MoE Transformer는 이 균형을 노린다.

많은 전문가로 용량을 키우고, 라우팅으로 언어와 토큰에 맞는 경로를 학습한다.

결과: 큰 MoE 모델이 가장 강했다

가장 큰 주요 모델은 다음 조건으로 학습됐다.

• 6,000억 파라미터
• 2,048개 TPU v3 코어
• 약 4일 학습
• 100개 언어에서 영어로 번역

평균 BLEU는 44.3이었다.

이중언어 기준 모델 평균은 30.8이었다.

즉, 단일 다국어 모델이 각 언어별 개별 모델보다 평균적으로 크게 앞섰다.

Crop 포인트: 더 큰 MoE 모델의 선이 전반적으로 위에 있으며, 고자원부터 저자원 언어까지 품질 향상이 이어지는 흐름이 중요하다.

깊이와 전문가 수의 효과

논문이 보여준 경향은 분명하다.

깊은 모델은 더 적은 토큰으로 수렴한다

같은 전문가 수라면 36층 모델이 12층 모델보다 학습 손실을 더 빨리 낮췄다.

이는 깊이가 샘플 효율에 기여한다는 해석을 가능하게 한다.

전문가 수는 고자원 언어에 특히 중요하다

전문가 수를 늘리면 모델 용량 병목이 완화된다.

고자원 언어는 데이터가 많다.

따라서 더 큰 용량을 활용할 여지가 크다.

저자원 언어에는 공유가 중요하다

저자원 언어는 독립 용량보다 전이가 중요할 수 있다.

논문은 깊고 조밀한 모델이 저자원 언어 전이에 강한 모습을 보인다고 해석한다.

다만 조밀한 거대 모델은 학습 비용이 훨씬 크다.

학습 효율: 단순히 큰 모델이 아니다

이 논문의 핵심은 “크다”가 아니다.

“큰데도 학습이 가능하다”가 핵심이다.

6,000억 파라미터 MoE 모델의 학습 비용은 약 22.4 TPU v3 core-year였다.

논문에서 비교한 96층 조밀 Transformer는 약 235.5 TPU v3 core-year가 필요했다.

품질도 MoE가 더 높았다.

즉, 조건부 계산은 단순한 모델 확장이 아니라 비용 구조를 바꾸는 장치다.

메모리 효율: 장치당 부담을 고정에 가깝게 유지

GShard에서 메모리는 세 덩어리로 나뉜다.

• 복제되는 일반 Transformer 가중치
• 분산되는 MoE 전문가 가중치
• 순전파와 역전파에 필요한 activation

전문가 수가 늘어도 각 장치가 맡는 전문가 조각만 관리하면 된다.

그래서 같은 깊이에서는 장치당 메모리가 크게 증가하지 않는다.

Crop 포인트: 전문가 수가 늘어나는 모델들 사이에서 장치당 메모리 막대가 크게 폭증하지 않는 점을 보면 된다.

실행 시간: 병목은 통신이지만 감당 가능한 수준

MoE 모델에서는 토큰을 전문가에게 보내고 다시 모으는 과정이 필요하다.

이 부분이 통신 비용을 만든다.

전문가 수가 늘면 이 통신 비용도 커진다.

하지만 전체 증가폭은 모델 크기 증가에 비해 작다.

논문은 128개 전문가에서 2,048개 전문가로 확장할 때 실행 시간이 약 1.7배 증가했다고 보고한다.

모델 규모 증가에 비하면 낮은 증가폭이다.

Crop 포인트: MoE dispatch와 combine 영역은 커지지만, 전체 막대의 증가가 모델 크기 증가만큼 폭발하지 않는 점이 핵심이다.

통신 프리미티브의 역할

GShard의 성능은 통신 프리미티브에 크게 좌우된다.

논문에서 특히 중요한 것은 두 가지다.

• 부분 결과를 전체 장치에서 합치는 통신
• 장치들이 서로 데이터 조각을 교환하는 통신

첫 번째는 장치 수가 늘어도 비용이 거의 일정하게 유지됐다.

두 번째는 비용이 증가하지만 비교적 완만했다.

이 때문에 MoE 라우팅과 재샤딩이 실용적인 수준에 머물렀다.

Crop 포인트: AllReduce 계열은 거의 낮게 유지되고, AllToAll 계열만 완만히 올라가는 대비를 보면 된다.

장치 배치도 성능 요소다

같은 텐서 분할이라도 장치 배치를 어떻게 하느냐에 따라 통신 비용이 달라진다.

하드웨어 토폴로지와 데이터 이동 방향이 맞아야 한다.

GShard의 일반 샤딩 API는 이런 배치를 표현할 수 있게 설계됐다.

Crop 포인트: 같은 논리적 텐서 분할도 하드웨어 연결 구조에 맞춰 장치 번호 배치를 다르게 가져갈 수 있다는 점이 중요하다.

MoE 밖에서도 쓸 수 있는가

논문은 GShard가 MoE 전용 도구가 아니라고 강조한다.

텐서 차원을 나누는 방식은 배치, 특징, 전문가, 공간 차원 모두에 적용될 수 있다.

예를 들어 큰 이미지나 3D 볼륨을 다루는 convolution 모델에도 사용할 수 있다.

여기서 핵심은 경계 처리다.

이미지를 공간 방향으로 나누면 인접 파티션의 일부 픽셀이 필요하다.

이때 전체 이미지를 복제하지 않고 필요한 경계만 교환한다.

Crop 포인트: 파티션 위치에 따라 필요한 경계 데이터 폭이 달라질 수 있다는 점을 보면 된다.

Crop 포인트: 최대 경계 데이터를 먼저 가져온 뒤, 각 파티션이 실제로 필요한 영역만 잘라 쓰는 흐름이 핵심이다.

더 복잡한 convolution 설정에서는 dilation과 stride 때문에 파티션별 정렬 문제가 생긴다.

GShard는 패딩, 슬라이싱, 마스킹을 조합해 이를 처리한다.

Crop 포인트: dilation 조건에 따라 같은 샤딩이라도 패딩과 윈도우 정렬 방식이 달라지는 세 구간을 보면 된다.

이 논문의 기여

GShard의 기여는 모델 하나에만 있지 않다.

더 중요한 기여는 추상화다.

개발자는 모델을 쓴다.

컴파일러는 병렬화를 만든다.

이 분리 덕분에 거대 모델 실험 주기가 짧아진다.

논문은 수천 개 장치에서 컴파일과 실행이 가능한 구조를 보여준다.

또한 조건부 계산이 단순한 아이디어가 아니라 실제 대규모 번역 품질을 끌어올리는 방법임을 보였다.

해석할 때 주의할 점

이 결과는 TPU와 XLA 기반 구현에서 나온 것이다.

다른 하드웨어나 컴파일러 환경에서는 통신 비용과 구현 난도가 달라질 수 있다.

또한 1조 파라미터 모델도 실험했지만, 논문은 재현성과 안정성 문제 때문에 주요 품질 결과에 포함하지 않았다.

따라서 이 논문의 안정적인 결론은 6,000억 파라미터 MoE Transformer 결과를 중심으로 보는 것이 적절하다.

결론

GShard는 거대 모델 학습에서 세 가지 병목을 동시에 줄인다.

• 모델 코드를 병렬화 세부 구현에서 분리한다.
• 조건부 계산으로 실제 계산량 증가를 억제한다.
• 단일 프로그램 방식으로 수천 장치 컴파일을 가능하게 한다.

그 결과 6,000억 파라미터 다국어 번역 모델을 4일 안에 학습했고, 100개 언어 번역에서 강한 품질 향상을 보였다.

이 논문의 메시지는 간단하다.

모델을 키우는 것만으로는 부족하다.

크게 만들면서도 효율적으로 학습할 수 있는 구조가 필요하다.

GShard는 그 방향을 보여준 대표적인 시스템 논문이다.

Source

• Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, Zhifeng Chen. “GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.” arXiv:2006.16668v1, 2020.
• arXiv: https://arxiv.org/abs/2006.16668
• PDF: https://arxiv.org/pdf/2006.16668