PP: 모델의 층을 나누어 실행하기

지금까지는 같은 레이어 안에서 특징이나 토큰을 나누었습니다. 모델에는 앞 레이어의 출력을 다음 레이어가 사용하는 순서도 있습니다. 파이프라인 병렬화(Pipeline Parallelism, PP)는 이 레이어들을 여러 GPU에 나누어 배치하는 방법입니다. 각 GPU가 모델의 일부를 맡으므로, 한 GPU에 담기 어려운 모델을 나누어 담을 수 있습니다.

하지만 레이어를 나누었다고 모든 GPU가 곧바로 함께 계산하는 것은 아닙니다. 뒤쪽 GPU는 앞쪽 GPU의 결과를 기다려야 합니다. 이번 글에서는 먼저 레이어 배치와 순전파·역전파의 전달 방향을 살펴보겠습니다. 이어서 배치를 통째로 넘길 때 GPU들이 얼마나 기다리는지 확인하겠습니다. 그 대기를 줄이기 위해 배치를 작은 묶음으로 나누고, 분할 전후를 비교해 여러 GPU가 함께 일하는 과정을 따라갑니다. 마지막에는 남는 대기와 분할 크기의 영향을 살펴보겠습니다.

층을 나누고 순전파·역전파 연결하기

Transformer 레이어가 여덟 개 있는 모델을 생각해 보겠습니다. GPU 0은 레이어 0·1, GPU 1은 레이어 2·3, GPU 2는 레이어 4·5, GPU 3은 레이어 6·7을 맡습니다. 각 GPU에 배치한 레이어 묶음을 파이프라인 단계(stage)라고 부릅니다. 여기서는 단계 하나에 GPU 하나를 사용하며, 입력과 출력에 필요한 나머지 구성은 생략하고 여덟 레이어의 연결에 집중합니다.

입력에서 출력을 만드는 방향의 계산이 순전파(forward pass)입니다. GPU 0이 첫 두 레이어를 계산하면, 그 출력인 활성값을 GPU 1로 보냅니다. GPU 1은 받은 값을 입력으로 다음 두 레이어를 계산합니다. 이런 식으로 GPU 3까지 이어집니다. 가중치를 매번 다음 GPU로 옮기는 것이 아니라, 각 단계에 가중치를 두고 중간값을 전달합니다.

학습에서는 모델의 출력과 정답으로 오차를 나타내는 손실을 계산합니다. 이어서 각 값이 손실에 미치는 영향을 기울기로 구하며 앞쪽으로 되짚는 계산이 역전파(backward pass)입니다. 다음 그림의 ‘다음’을 눌러 순전파와 역전파의 방향을 비교해 보겠습니다.

역전파에서는 GPU 3에서 GPU 0 쪽으로 진행합니다. 예를 들어 GPU 1은 뒤쪽에서 받은 기울기와 순전파에서 얻은 값을 이용해 자기 레이어를 역전파합니다. 그 과정에서 자기 가중치의 기울기는 GPU 1에서 계산하고, 입력 활성값의 기울기는 앞 단계인 GPU 0으로 보냅니다. 그림의 역방향 화살표는 이 활성값 기울기의 전달을 나타냅니다.

어느 방향이든 앞선 계산의 결과가 필요합니다. 입력 하나의 순전파에서 GPU 0이 아직 활성값을 만들지 못했다면 GPU 1은 시작할 수 없습니다. 여러 GPU가 같은 입력의 모든 레이어를 동시에 계산하는 구조는 아닙니다. 이 의존성이 여러 입력을 묶어 실행할 때 어떤 대기를 만드는지 살펴보겠습니다.

배치를 통째로 처리하면 생기는 대기

함께 처리할 입력 묶음을 배치(batch)라고 합니다. 입력 시퀀스 A·B·C·D 네 개를 한 배치로 처리한다면 배치 크기는 4입니다. 각 입력은 여러 토큰으로 이루어져 있습니다.

먼저 네 입력을 한 묶음으로 계산하고, 결과도 통째로 다음 GPU에 넘긴다고 가정하겠습니다. GPU 0이 A·B·C·D를 계산하는 동안 GPU 1은 입력을 받지 못해 기다립니다. GPU 2와 GPU 3도 마찬가지입니다. GPU 0이 배치 전체의 계산을 끝내야 GPU 1이 시작할 수 있습니다.

그림에서는 GPU 1이 계산을 시작한 뒤에도 GPU 2와 GPU 3은 기다리고, GPU 0은 이미 자기 계산을 끝낸 상태입니다. 이 배치만 실행하면 한 GPU가 계산하는 동안 다른 세 GPU는 계산할 일이 없습니다. 레이어를 나누어 담았지만, 계산은 여전히 차례로 진행되는 것입니다.

다음 GPU를 더 일찍 시작시킬 수는 없을까요? 네 입력의 계산이 모두 끝날 때까지 기다리는 대신, A의 계산이 끝나자마자 그 결과를 다음 GPU로 넘기면 됩니다. 그러려면 배치를 더 작은 묶음으로 실행할 수 있어야 합니다.

한 배치를 나누어 다음 GPU를 먼저 시작시키기

이 배치를 더 작은 묶음으로 나눈 것이 마이크로배치(microbatch)입니다. 여기서는 A 하나, B 하나, C 하나, D 하나씩 크기 1인 마이크로배치 네 개로 나누겠습니다. 배치 크기 4를 크기 1인 네 묶음으로 실행하는 것입니다. 처리할 입력은 그대로 네 개이고, 각 입력의 토큰 시퀀스도 그대로 유지합니다.

A를 GPU 0에, B를 GPU 1에 맡기고 끝내는 것도 아닙니다. GPU마다 모델의 일부 레이어만 있으므로 A·B·C·D 모두 GPU 0부터 GPU 3까지 통과해야 합니다. 바뀌는 것은 각 단계를 통과시키는 묶음의 크기입니다. 크기 1은 설명을 위한 예시이며, 배치 크기 8을 크기 2인 마이크로배치 네 개로 나눌 수도 있습니다.

학습에서 실행 묶음을 나누는 것과 가중치를 갱신하는 횟수도 구별해야 합니다. 이 글에서는 네 마이크로배치의 순전파·역전파 동안 같은 가중치를 사용하고, 전체 배치에 맞게 기울기를 누적한 뒤 가중치를 한 번 갱신하는 구성을 가정합니다. GPipe 논문이 설명하는 방식입니다. 입력 하나를 처리할 때마다 가중치를 바꾸는 네 번의 학습 단계로 바꾼다는 뜻은 아닙니다.

서로 다른 입력의 계산을 겹치기

배치를 통째로 전달하면 GPU 0이 네 입력을 모두 계산할 때까지 GPU 1은 기다립니다. 마이크로배치로 나누면 GPU 0이 A를 계산한 직후 그 결과를 GPU 1로 보낼 수 있습니다. GPU 1이 A를 계산하는 동안 GPU 0은 B를 계산합니다. 한 입력의 계산 순서는 유지하면서, 서로 다른 입력의 계산을 겹치는 것입니다.

같은 GPU 네 개와 같은 입력 네 개로 두 실행을 비교하겠습니다. 차이를 보기 쉽도록 각 단계에서 입력 하나를 계산하는 시간은 1칸, 네 입력을 통째로 계산하는 시간은 4칸이라고 가정합니다. 모든 단계의 계산 시간은 같고, 통신 시간은 생략합니다. 다음 그림은 이 가정 아래의 순전파만 비교하며, 역전파까지 포함한 전체 학습 시간은 아닙니다.

위쪽에서는 GPU 0이 배치 전체를 4칸 동안 계산합니다. 그 결과를 GPU 1이 받아 다시 4칸 동안 계산하고, GPU 2와 GPU 3도 차례로 계산합니다. 네 입력 모두 마지막 GPU를 통과하는 시점은 16칸 뒤입니다. 각 GPU의 한 블록은 네 입력을 함께 처리하는 작업입니다.

아래쪽에서는 1칸 뒤 GPU 0의 A가 GPU 1로 넘어갑니다. 그다음 1칸 동안 GPU 0은 B를, GPU 1은 A를 계산합니다. 뒤쪽 단계까지 차례로 채워지면, 시각 3부터 4까지는 GPU 0이 D, GPU 1이 C, GPU 2가 B, GPU 3이 A를 동시에 계산합니다. 4칸 뒤 A가 먼저 끝나고, 7칸 뒤에는 D까지 모두 끝납니다.

이 예시에서는 같은 네 입력을 16칸 대신 7칸에 끝내므로 처리량이 증가합니다. 처리량은 단위 시간에 완료한 작업량입니다. 입력 하나가 필요한 레이어를 건너뛰거나, 계산 자체를 없앤 것은 아닙니다. 네 GPU가 각각 수행하는 계산량은 두 경우 모두 4칸이고, 달라진 것은 서로 기다리는 시간입니다.

실제 배치 크기 4의 계산 시간이 배치 크기 1의 정확히 네 배인 것은 아닙니다. 큰 배치의 행렬 연산이 더 효율적일 수도 있고, 마이크로배치를 전달할 때 통신 비용도 듭니다. 따라서 16칸과 7칸은 실제 속도 향상의 보장이 아니라, 작은 묶음을 먼저 넘겨 다음 GPU를 일찍 시작시키는 원리를 보여주는 수치입니다.

남는 대기와 분할 크기의 균형

마이크로배치로 나눈 그림에도 빈칸이 남습니다. 처음에는 A가 뒷단에 도착하기 전까지 뒤쪽 GPU가 기다립니다. 마지막에는 D가 앞단을 떠난 뒤 뒤쪽 단계를 통과하는 동안 앞쪽 GPU가 빕니다. 이렇게 파이프라인을 채우고 비우는 동안 생기는 빈 구간을 버블(bubble)이라고 부릅니다. 그림에서 7칸 뒤의 아래쪽은 이 배치의 순전파가 이미 끝난 구간입니다.

더 많은 마이크로배치를 이어 넣으면 처음과 끝의 대기가 전체에서 차지하는 비중을 줄일 수 있습니다. 하지만 정해진 배치를 너무 잘게 나누면 한 번의 행렬 연산이 작아져 GPU를 효율적으로 쓰기 어려워지고, 전달 횟수와 실행을 관리하는 비용도 늘 수 있습니다. 마이크로배치 수를 늘리는 것만으로 항상 빨라지지는 않습니다.

단계의 계산 시간이 다르면 느린 단계 때문에 추가로 기다립니다. 레이어 두 개씩 배치했더라도 모든 단계의 실행 시간이 같다고 보장되지는 않습니다. 실제로는 각 단계의 계산량과 메모리 사용량, 경계에서 전달할 데이터 크기를 함께 확인해야 합니다.

추론에서도 서로 다른 입력을 여러 단계에 겹쳐 실행할 수 있습니다. 다만 한 요청이 토큰을 하나씩 생성할 때는 이전 토큰의 결과가 있어야 다음 입력을 정할 수 있으므로, 그 의존성은 남습니다. 함께 진행할 입력이 충분한지도 파이프라인을 채우는 데 영향을 줍니다.

PP는 모델의 레이어를 나누어 담고, 마이크로배치를 통해 서로 다른 입력의 계산을 겹치는 방법입니다.