GPipe — 층을 나누고, 배치를 쪼개고, 갱신은 한 번에

1. 한 장치에 안 들어가는 모델

논문은 그림 1로 시작합니다. ImageNet 분류 모델은 몇 해 사이 용량이 36배 커졌고 정확도도 그만큼 따라 올랐습니다. 번역에서도 Transformer를 키울수록 BLEU가 올랐습니다. 문제는 모델이 커지는 속도를 가속기 메모리가 따라가지 못한다는 것입니다. 학습할 때 장치 메모리에는 가중치만 올라가는 게 아닙니다. 기울기와 옵티마이저 상태가 함께 올라가고, 역전파에 쓸 순전파 중간 결과(활성)도 남겨 둬야 합니다.

논문 표 1의 Transformer 설정으로 계산해 봅니다. 층마다 모델 차원 2048, 피드포워드 차원 8192, 어휘 32k, 시퀀스 길이 1024, 배치 32이고, RMSProp을 써서 파라미터 하나에 12바이트가 든다고 표 1 설명에 적혀 있습니다(가중치·기울기·RMSProp 누적값이 4바이트씩이라고 보면 맞는 숫자입니다).

파라미터만 놓고 보면 16GB 한 장에는 약 14억 개, 25층까지 들어갑니다. 그런데 표 1에서 GPipe 없이 한 장치로 학습할 수 있었던 건 3층뿐이었습니다. 나머지 자리를 활성이 차지했기 때문입니다¹. 128장으로 늘려도 같은 일이 일어납니다. 128 × 16GB = 2,048GB면 파라미터만으로는 1,833억 개가 들어가지만, 논문이 실제로 올린 것은 839억 개였고 그때 활성이 796.1GB를 차지했습니다. 모델을 장치에 나눠 싣는 것만큼 활성을 줄이는 것이 중요하다는 뜻이고, 이건 6절에서 다룹니다.

2. 층으로 자르기: 단순 모델 병렬과 유휴

GPipe가 모델에 요구하는 것은 하나입니다. 층을 순서대로 나열할 수 있을 것. L개 층의 순서열을 K개의 연속된 덩어리(논문은 셀이라고 부릅니다)로 자르고, k번째 셀을 k번째 가속기에 올립니다.

장치 사이에는 셀 경계의 활성만 오갑니다. 순전파는 장치 0에서 K−1로, 역전파는 거꾸로 흐릅니다. B_k는 위쪽 셀의 B_k+1과 자기 셀의 F_k 결과가 둘 다 있어야 계산할 수 있습니다. 이 의존 때문에 미니배치 하나를 통째로 흘리면 장치들이 차례로 일하고 차례로 놉니다(그림 2b).

K를 올려도 한 스텝에 걸리는 시간은 그대로이고, 각 장치가 일하는 비율만 1/K로 줄어듭니다. 이 방식이 쓸모없는 건 아닙니다. 한 장치에 안 들어가던 모델이 들어가게는 되니까요. 다만 장치를 K배 써도 빨라지지는 않습니다. 논문 표 2의 M = 1 행이 이것을 잽니다. Transformer의 처리량이 K = 2, 4, 8에서 1, 1.07, 1.3으로 거의 그대로였고, 저자들은 “어느 순간에든 한 장치만 계산하고 있다는 뜻”이라고 적었습니다.

3. 배치를 쪼개기: 마이크로배치 파이프라인 (그림 2c)

GPipe의 해법은 미니배치를 쪼개는 것입니다. 크기 N인 미니배치를 크기 N/M인 마이크로배치 M개로 나눠 차례로 흘려보냅니다. 장치 0이 첫 조각의 F를 끝내고 장치 1에 넘기면, 장치 1이 그 조각을 처리하는 동안 장치 0은 둘째 조각을 시작할 수 있습니다. 조립 라인과 같습니다.

역전파도 같은 식으로 흐릅니다. 각 마이크로배치의 기울기는 순전파 때와 같은 파라미터로 계산하고, 미니배치 끝에서 M개 조각의 기울기를 모두 더해 모든 장치가 동시에 한 번 갱신합니다(2.2절). 갱신이 한 번뿐이므로 그 사이에 가중치가 바뀌는 일이 없습니다.

M = 1에서는 데모 2와 똑같이 계단 하나가 올라갔다 내려옵니다. M을 키우면 계단의 경사는 그대로인데 가운데 모든 장치가 동시에 일하는 구간이 길어집니다. 비어 있는 칸은 두 삼각형, 순전파가 파이프라인을 채우는 동안의 왼쪽 위와 역전파가 비워 나가는 오른쪽 아래, 그리고 순전파와 역전파 사이에 끼인 가운데 부분에만 남습니다. 논문은 이 빈 공간을 버블이라고 부릅니다.

4. 버블은 얼마나 남는가

모든 셀의 계산 시간이 같다고 두고 세어 봅니다. 한 장치가 마이크로배치 하나의 F에 t_F, B에 t_B를 씁니다. 장치 k의 첫 F는 앞 장치 k개가 첫 조각을 넘겨줄 때까지 기다려야 하므로 k·t_F 늦게 시작하고, 마지막 장치는 (K−1)·t_F 늦게 시작합니다. 순전파가 모두 끝나는 데 (M + K − 1)·t_F, 역전파도 같은 이유로 (M + K − 1)·t_B입니다. 그동안 장치 하나가 실제로 일한 시간은 M·(t_F + t_B)입니다.

논문은 실험에서 M ≥ 4K이면 버블이 무시할 만했다고 보고합니다. 다만 위 식을 그대로 넣으면 M = 4K일 때 (K−1)/(5K−1)이라 K = 8에서 18%, K가 커지면 20%에 다가가서 “무시할 만하다”고 하기에는 큰 값입니다. 논문은 그 이유의 일부로, 재계산(6절)을 쓰면 역전파 앞부분의 순전파 재계산은 위쪽 셀의 기울기를 기다릴 필요가 없어 버블 안으로 미리 당겨 할 수 있다는 점을 듭니다. 실제 측정은 아래 표 2처럼 식보다 나은 경우가 있습니다.

M = 32 행은 모형과 8% 안에서 맞습니다. K를 2에서 8로 4배 늘렸을 때 1.8 → 6.3, 즉 3.5배라는 논문의 말이 여기서 나옵니다. M = 4 행은 측정이 모형보다 훨씬 높은데, 위에서 말한 재계산의 앞당김과 배치 크기 조정이 섞였을 것으로 보입니다. M = 1 행의 K = 8 값 1.3은 모형(1.0)보다 30% 높은데, 논문에는 이 차이에 대한 설명이 없습니다. 3절은 이 행을 두고 “M = 1이면 사실상 파이프라인 병렬이 없고, 장치 수와 상관없이 처리량이 비교적 일정해 한 번에 한 장치만 계산하고 있음을 보여 준다”고만 적었습니다. 표 2 주석의 “필요하면 메모리에 맞게 배치 크기를 조정했다”는 문장이 원인일 수는 있지만, 어느 칸의 배치를 얼마로 바꿨는지는 밝히지 않아 확인할 수 없습니다.

이후 논문들(Megatron-LM 등)은 버블을 전체 시간에 대한 비율이 아니라 이상적인 계산 시간에 대한 비율로 (K−1)/M이라고 적습니다. 분모가 다를 뿐 같은 양입니다. 위 readout에 둘 다 표시했습니다.

5. 고르게 자르기: 분할과 불균형

4절의 식은 모든 셀이 똑같이 오래 걸린다는 가정 위에 있습니다. 실제로는 그렇지 않습니다. 파이프라인은 가장 느린 셀의 속도로 흐르기 때문에, 다른 셀들은 그 셀을 기다리며 놉니다. 그래서 GPipe는 층별 비용 추정치 c_l로 셀 비용 C_k를 구하고, 셀 비용들의 분산이 최소가 되도록 자릅니다(2.2절). 셀 비용의 합은 자르는 방식과 상관없이 같으므로, 분산을 최소화하는 것은 제곱합 Σ C_k²를 최소화하는 것과 같고, 연속 구간 분할이라 동적 계획법으로 정확히 풀 수 있습니다.

“층 수로 똑같이”에 뒤로 갈수록 무거운 비용을 주면 마지막 셀이 다른 셀의 몇 배가 되고, 이상적인 가속이 K보다 한참 낮아집니다. 분산 최소화로 바꾸면 앞쪽 셀은 층을 많이, 뒤쪽 셀은 적게 가져가 다시 고르게 됩니다. 그래도 층 하나보다 잘게 자를 수는 없어서, 무거운 층이 하나 있으면 그 층이 상한이 됩니다. 논문은 AmoebaNet의 가속이 선형에 못 미친 이유(표 2에서 K = 8, M = 32일 때 Transformer 6.3, AmoebaNet 3.48)와 최대 모델 크기가 장치 수에 비례해 늘지 않은 이유(표 1)를 모두 층마다 파라미터와 계산이 고르지 않은 것에서 찾고, 더 나은 분할 알고리즘이 성능을 올릴 여지가 있다고 적었습니다.

6. 재계산으로 활성 메모리 갚기

역전파에서 층 i의 기울기 b_i를 구하려면 위층에서 내려온 기울기 b_i+1과 함께 순전파 때의 활성 f_i(x)가 필요합니다. 그래서 보통은 순전파 동안 모든 층의 활성을 남겨 두고, 그 메모리는 O(N × L)입니다. 파이프라인에서는 사정이 더 나쁩니다. 장치 0은 첫 마이크로배치의 역전파가 돌아올 때까지 M개 조각의 활성을 모두 들고 있어야 합니다.

GPipe는 Chen et al.(2016)의 재계산(re-materialization, gradient checkpointing이라고도 부릅니다)을 장치마다 적용합니다. 순전파 때는 셀 경계의 출력만 저장하고 셀 내부 활성은 버립니다. 역전파 차례가 오면 저장해 둔 경계 입력으로 F_k를 다시 계산해 내부 활성을 복원하고, 그 마이크로배치의 역전파가 끝나면 다시 버립니다.

재계산이 없으면 M을 키워도 정점은 N·L/K로 그대로입니다. 마이크로배치를 쪼개도 결국 M개 조각을 전부 들고 있어야 하니까요. 재계산을 켜면 정점이 N + (L/K)(N/M)으로 내려가고, 이제 M을 키울수록 둘째 항이 줄어듭니다. 버블을 줄이려고 키운 M이 메모리까지 줄여 주는 셈입니다. 이것이 6절 끝에서 논문이 “재계산과 결합하면 많은 수의 마이크로배치로 확장할 수 있다”고 말하는 이유입니다. 대가는 시간입니다. 역전파가 순전파의 두 배라면 재계산은 계산량을 3 → 4, 약 33% 늘립니다.

이 효과는 장치 하나에서도 나타납니다. 표 1의 Pipeline-1 열은 장치 1대에 재계산과 배치 쪼개기만 적용한 경우인데, AmoebaNet의 활성 메모리가 6.26GB에서 3.46GB로 줄어 82M 대신 318M 모델이 들어갔고, Transformer는 2.7배 큰 모델(3층 → 13층)이 들어갔습니다.

7. 쪼개도 같은 기울기: 동기식 누적

마이크로배치로 쪼갠 뒤 기울기를 더하는 것이 왜 원래 미니배치의 기울기와 같은지는 손실의 정의에서 바로 나옵니다. 미니배치 손실은 예제별 손실의 평균이고, 미분은 덧셈을 지나갑니다.

예외는 배치 전체를 보는 층입니다. BatchNorm은 학습 중 통계를 미니배치가 아니라 마이크로배치마다 계산하게 되므로 결과가 달라집니다. 논문은 이 경우 통계를 마이크로배치(필요하면 복제본들)에 대해 계산하고, 평가용 이동 평균은 미니배치 전체에 대해 따로 추적한다고 밝힙니다(2.2절).

②의 손실 곡선은 ①과 정확히 겹치고, 두 모델 가중치의 최대 차이는 10⁻¹⁵ 근처입니다. 덧셈 순서가 달라서 생기는 부동소수점 반올림 차이일 뿐입니다. M을 바꿔 다시 돌려도 마찬가지입니다. ③은 M에 따라 다른 경로로 학습됩니다. 더 빨리 내려갈 때도 있지만 그건 “다른 알고리즘”이라는 뜻이고, 장치 수를 바꿀 때마다 하이퍼파라미터를 다시 맞춰야 한다는 뜻이기도 합니다. ④는 9절의 PipeDream이 피하려고 애쓴 문제를 보여 줍니다. 지연 d를 키우면 손실이 출렁이거나 잘 내려가지 않습니다.

이 계산이 정확히 같으려면 층에 배치 통계가 없어야 합니다. 이 MLP에는 BatchNorm이 없습니다. 또 학습률을 크게 올리면 학습이 혼돈스러워져 10⁻¹⁵의 반올림 차이도 수백 스텝 뒤 눈에 보이게 커질 수 있습니다. 수학은 같아도 부동소수점 결과까지 비트 단위로 같다는 보장은 아닙니다.

8. 논문의 결과

얼마나 큰 모델이 들어갔나 (표 1)

AmoebaNet은 224 × 224 입력, 미니배치 128에서 장치당 8GB, Transformer는 앞의 설정으로 장치당 16GB에서 잰 값입니다. Naive-1은 GPipe 없는 한 장치, Pipeline-k는 GPipe로 k개 장치에 나눈 경우입니다.

모든 열이 장치 메모리를 80~95%쯤 채운 곳에서 멈춰 있습니다. 표가 “들어가는 한 최대로 키운 모델”을 잰 것이니 당연한 모양입니다. Transformer는 층마다 크기가 같아서 장치 수에 정확히 비례해 커졌고(8 → 32 → 128장에서 103 → 415 → 1663층), 128장에서 한 장치 대비 298배인 839억 파라미터에 도달했습니다. AmoebaNet은 8장에서 18억으로, 논문 표현으로 GPipe 없이 가능한 것의 25배에 그쳤습니다. 층마다 파라미터가 고르지 않아 분할이 완벽하지 않았기 때문입니다(5절).

고속 연결 없이도 (표 3)

NVLink 없이 PCI-E로만 연결된 P100 GPU 여러 장에서 M = 32로 잰 처리량입니다. 장치 사이에 오가는 것이 셀 경계의 활성뿐이라 통신 대역폭이 병목이 되지 않았고, 가속은 TPU에서와 비슷했습니다.

키운 모델로 얻은 것: 이미지 분류 (4절, 표 4)

AmoebaNet-B(18, 512)를 채널 수를 늘리고 입력을 480 × 480으로 키워 5억 5,700만 파라미터로 만들고, 4개 파티션으로 나눠 ImageNet을 학습했습니다. 단일 크롭 top-1 84.4%, top-5 97%입니다. 이 모델을 다른 데이터셋에 미세 조정한 결과(5회 평균)는 다음과 같습니다.

표 4. 별표(*)는 논문 표기를 그대로 옮긴 것입니다. ImageNet의 85.4%는 비공개 Instagram 데이터로 사전학습한 Mahajan et al.의 결과입니다.

키운 모델로 얻은 것: 다국어 번역 (5절, 표 5)

102개 언어 ↔ 영어, 언어당 10⁴~10⁹개, 총 250억 개의 문장 쌍으로 Transformer 하나를 학습했습니다. 4억 파라미터의 T(6, 8192, 16)에서 시작해 깊게(T(24, 8192, 16), 13억), 넓게(T(12, 16384, 32), 13억), 그리고 30억(T(32, 16384, 32))과 60억(T(64, 16384, 32))까지 키웠고, 큰 모델들은 2·4·8·16개 가속기에 나눴습니다. T(L, H, A)는 인코더·디코더 각 L층, 피드포워드 차원 H, 어텐션 헤드 A이고 모델 차원은 1024로 고정입니다.

• 용량 — 4억 → 13억에서 모든 언어가 크게 좋아졌고, 13억 → 60억은 고자원 언어에서 더 좋아졌지만 수확은 줄어들었습니다(그림 3). 60억 모델 하나가 언어쌍 100개에서 각각 따로 학습한 3억 5,000만 파라미터 이중 언어 모델보다 나았습니다.
• 깊이 대 너비 — 같은 13억에서 고자원 언어는 비슷했지만 저자원 언어는 깊은 모델이 큰 차이로 앞섰습니다. 논문은 깊이가 일반화와 전이에 더 나을 수 있다고 해석합니다.
• 깊은 모델의 학습 불안정 — 수천 스텝 뒤 예측이 지나치게 뾰족해지며 기울기가 폭발했습니다. 피드포워드 층 초기화를 층 수로 줄이고(Fixup을 따름), 소프트맥스 직전 로짓의 크기를 잘라 막았습니다.
• 큰 배치 — 독일어→영어에서 배치를 토큰 26만 → 100만 → 400만으로 키우자 BLEU 30.92 → 31.86 → 32.71, 검증 NLL 2.58 → 2.51 → 2.46으로 좋아졌습니다(표 5). 저자들이 아는 한 NMT에서 쓰인 가장 큰 배치였습니다.

9. 다른 병렬화와 비교: 데이터·텐서 병렬, PipeDream, 1F1B

논문 6절은 대규모 모델 병렬의 여러 접근을 “각자 다른 절충을 골랐다”는 관점에서 비교합니다. 여기에 논문 이후에 자리 잡은 1F1B 스케줄까지 더해 정리합니다.

GPipe가 셋째 줄에서 내세우는 것은 두 가지입니다. 통신이 셀 경계에서만 일어나 고속 연결이 없어도 되고(표 3), 갱신이 동기식이라 7절처럼 결과가 분할 방식과 무관합니다. 각주 3에서는 층 하나가 한 장치에 안 들어가는 한계를 피하는 방법으로 큰 행렬곱을 작은 것들로 쪼개 여러 층에 나눠 두는 것을 제안합니다. 실제로 이후의 초거대 모델 학습은 데이터·텐서·파이프라인 병렬을 함께 쓰는 쪽으로 갔습니다(Megatron-LM, 2021).

1F1B: 같은 버블, 더 적은 활성

GPipe 스케줄은 M개 조각의 순전파를 모두 끝낸 뒤에야 역전파를 시작하므로, 장치 0은 M개 조각의 경계 활성을 한꺼번에 들고 있습니다. PipeDream이 쓴 1F1B(one-forward-one-backward) 순서는 파이프라인이 찬 뒤로 순전파 하나, 역전파 하나를 번갈아 합니다. 역전파가 일찍 시작되니 끝난 조각의 활성을 일찍 버릴 수 있습니다. PipeDream은 이것을 비동기로 돌렸지만, 미니배치 끝에서 파이프라인을 비우고(flush) 한 번에 갱신하면 GPipe와 같은 동기식 학습이 됩니다(PipeDream-Flush, Narayanan et al. 2021).

두 타임라인의 길이는 같습니다. 버블은 둘 다 (K−1)/(M+K−1)입니다. 다른 것은 오른쪽 숫자입니다. GPipe는 모든 장치가 M개를, 1F1B는 장치 k가 K−k개 이하를 들고 있습니다. M을 24까지 올려도 1F1B의 장치 0은 K개에서 멈춥니다. 그래서 버블을 줄이려고 M을 크게 잡을 때 1F1B는 활성 메모리가 M과 무관해지고, 이후의 파이프라인 병렬 구현은 대부분 이 순서를 씁니다. 반대로 GPipe는 이 부담을 재계산(6절)으로 줄였습니다. 두 방법은 함께 쓸 수도 있습니다.

같은 시리즈의 How SGD Works는 미니배치 기울기를, How Adam Works는 옵티마이저 상태가 학습 메모리에서 차지하는 몫을, How Transformer Works는 이 페이지가 파라미터를 센 Transformer 블록을 다룹니다.

10. 확인 퀴즈