TP: 하나의 연산을 여러 GPU로 나누기

데이터 병렬화(DP)에서는 모델을 복제하고 GPU마다 서로 다른 입력을 처리했습니다. 텐서 병렬화(Tensor Parallelism, TP)는 여러 GPU가 같은 입력의 연산을 나누어 처리하는 방법입니다. 하나의 요청을 처리하는 중에도 여러 GPU가 협력할 수 있고, 하나의 가중치 행렬도 나누어 보관할 수 있습니다. 여기서 텐서는 벡터나 행렬처럼 여러 값을 배열한 데이터를 뜻합니다.

연산을 나누면 각 GPU의 계산량은 줄지만, 나뉜 결과를 어떻게 연결할지가 중요해집니다. 매번 전체 결과를 모은다면 통신에 많은 시간을 쓸 수 있기 때문입니다. 이 글에서는 먼저 가중치의 열을 나눌 때와 행을 나눌 때 결과가 어떻게 다른지 살펴보겠습니다. 이어서 두 방식을 FFN과 어텐션에 연결해 중간 통신을 생략하고, 트랜스포머 레이어 전체에서 어디에 통신이 남는지 확인하겠습니다.

그림은 GPU 두 개가 같은 입력을 처리하는 순전파를 보여줍니다. 파란색은 GPU 0, 주황색은 GPU 1이 맡은 부분입니다. ‘다음’을 누르면 계산이 진행되며, 숫자가 없는 행렬 격자는 값 대신 크기와 분할 위치를 나타냅니다. 세부 내용을 크게 보고 싶다면 그림을 눌러 원본으로 열 수 있습니다.

열을 나누어 출력 조각 완성하기

먼저 하나의 행렬 곱 Y = XW를 두 GPU로 나누어 보겠습니다. X는 입력, W는 가중치, Y는 출력입니다. 그림의 X는 두 토큰이 각각 네 개의 특징값을 가진 2×4 행렬입니다. 행 하나가 토큰 하나이고, 열 하나가 특징 하나입니다. 두 GPU가 서로 다른 토큰을 맡는 것이 아니라, 같은 두 토큰에 대한 계산을 함께 나눈다는 점을 기억하면 됩니다.

W는 4×4 행렬입니다. W의 각 열은 입력의 네 특징을 사용해 출력 특징 하나를 계산합니다. 따라서 W를 열 방향으로 나누면 각 GPU가 서로 다른 출력 특징을 맡을 수 있습니다. GPU 0은 앞의 두 열을, GPU 1은 뒤의 두 열을 보관합니다. 두 GPU 모두 입력 X 전체를 갖습니다.

그림에서 첫 토큰의 입력은 [1, 2, 0, 1]이고, 완성된 출력은 [3, 2, 5, 2]입니다. GPU 0은 앞의 [3, 2]를, GPU 1은 뒤의 [5, 2]를 계산합니다. 각 값에는 해당 출력 특징을 계산하는 데 필요한 곱셈과 덧셈이 모두 반영되어 있습니다.

즉, 열 분할의 결과는 서로 다른 위치의 완성된 출력 조각입니다. 두 조각을 더하는 것이 아니라 열 방향으로 나란히 놓아야 전체 Y가 됩니다. 두 번째 토큰도 같은 방식으로 나누어 계산합니다.

그렇다고 실제 실행에서 바로 조각을 모아야 하는 것은 아닙니다. 다음 연산이 각 GPU에서 Y 전체를 필요로 한다면 All-Gather로 모을 수 있습니다. 반대로 다음 연산을 자기 조각만으로 시작할 수 있다면, 나뉜 상태를 유지하면 됩니다. 그림의 마지막 단계는 두 조각과 전체 출력의 관계를 보여주며, 반드시 그 시점에 통신한다는 뜻은 아닙니다.

행을 나누어 출력의 부분합 계산하기

이번에는 같은 X와 W를 다른 방향으로 나눠 보겠습니다. W의 앞 두 행을 GPU 0에, 뒤 두 행을 GPU 1에 둡니다. W의 행은 입력 특징에 대응하므로, X도 그에 맞춰 특징 열을 나눕니다. GPU 0은 각 토큰의 앞 두 특징을, GPU 1은 뒤 두 특징을 사용합니다. 여기서 ‘행 분할’은 가중치 W의 행을 나눈다는 뜻이며, 입력 토큰을 나누는 것이 아닙니다.

행렬 곱의 출력값 하나를 얻으려면 입력 특징마다 곱한 값을 모두 더해야 합니다. 이번에는 각 GPU가 그중 일부 입력 특징만 담당하므로, 계산 결과도 아직 덜 더해진 값입니다. 이런 결과를 부분합이라고 부르겠습니다.

첫 토큰에서 GPU 0의 부분합은 [1, 2, 4, 1], GPU 1의 부분합은 [2, 0, 1, 1]입니다. 같은 위치의 값끼리 더하면 원래 출력인 [3, 2, 5, 2]가 됩니다. 예를 들어 첫 출력값 3은 GPU 0이 계산한 1과 GPU 1이 계산한 2를 더한 값입니다.

열 분할과 달리 각 GPU의 결과는 전체 Y와 크기가 같습니다. 하지만 같은 크기라고 해서 완성된 결과는 아닙니다. 각 결과에는 입력 특징의 일부만 반영되어 있으므로, 두 결과를 이어 붙여서는 원래 행렬 곱을 얻을 수 없습니다. 반드시 같은 위치끼리 더해야 합니다.

이 합산 결과를 두 GPU 모두에서 사용하려면 All-Reduce가 필요합니다. All-Reduce는 각 GPU의 값을 같은 위치끼리 합산하고, 완성된 결과를 모두에게 돌려주는 집합 통신입니다. 뒤의 예시에서는 이 방식으로 전체 출력을 각 GPU에 준비하겠습니다.

이제 두 분할의 차이를 정리할 수 있습니다. 열 분할은 출력의 서로 다른 조각을 완성하고, 행 분할은 출력의 같은 위치에 기여하는 부분합을 만듭니다. 다음으로 이 차이를 이용해 두 연산을 연결해 보겠습니다.

FFN에서 열 분할과 행 분할 연결하기

트랜스포머의 FFN(Feed-Forward Network)은 각 토큰의 특징을 변환하는 부분이며, MLP라고도 부릅니다. 여기서는 두 선형 변환, 즉 두 Linear 사이에 활성화 함수가 있는 기본 형태를 보겠습니다. 첫 Linear는 특징 수를 늘리고, 두 번째 Linear는 다시 줄입니다. 활성화 함수는 그 사이에서 값을 비선형으로 바꿉니다.

이를 Y = f(XU)V로 쓸 수 있습니다. X에 첫 가중치 U를 곱하고, 활성화 함수 f를 적용한 뒤, 두 번째 가중치 V를 곱한다는 뜻입니다. 그림에서는 입력 특징 4개를 중간 특징 8개로 늘렸다가 다시 4개로 줄입니다. 설명을 간단히 하기 위해 편향 항은 생략합니다.

첫 가중치 U를 열로 나누면 중간 결과도 특징 방향으로 나뉩니다. GPU 0은 앞 네 특징을, GPU 1은 뒤 네 특징을 계산합니다. 이때 각 GPU가 맡은 특징값은 다른 GPU의 결과를 더할 필요 없이 이미 완성되어 있습니다.

두 번째 Linear로 넘어가기 전에는 이 중간값에 활성화 함수를 적용해야 합니다. 여기서는 GeLU처럼 각 값을 독립적으로 바꾸는 원소별 활성화 함수를 사용합니다. 특징값 하나를 바꾸는 데 다른 특징값은 필요하지 않으므로, GPU 0은 앞 네 특징에, GPU 1은 뒤 네 특징에 각각 활성화를 적용하면 됩니다. 여덟 특징을 한곳에 모아 활성화를 적용한 뒤 다시 나눈 것과 같은 결과입니다. 따라서 활성화를 거친 뒤에도 중간값을 각 GPU에 나누어 둔 채 다음 Linear로 넘어갈 수 있습니다.

그다음이 핵심입니다. 두 번째 가중치 V를, 각 GPU가 이미 가진 중간 특징에 맞춰 행으로 나눕니다. GPU 0에는 앞 네 특징에 대응하는 행을, GPU 1에는 뒤 네 특징에 대응하는 행을 둡니다. 그러면 각 GPU는 자기 중간값과 자기 V 조각만으로 두 번째 행렬 곱을 계산할 수 있습니다.

이 결과는 행 분할의 출력이므로 부분합입니다. 마지막에 All-Reduce로 두 부분합을 더하면 두 GPU 모두 완성된 Y를 갖게 됩니다.

첫 Linear를 열로 나누고 두 번째 Linear를 행으로 나누면, 두 Linear 사이에서 중간값 전체를 모으는 통신을 생략할 수 있습니다. 첫 연산이 만든 조각을 다음 연산이 그대로 사용할 수 있도록 가중치의 분할 방향을 맞춘 것입니다. 중간 결과를 전부 모았다가 다시 나누는 과정이 필요하지 않습니다.

어텐션에서 헤드별 계산 이어가기

어텐션에도 입력을 바꾸는 선형 변환과 출력 결과를 바꾸는 선형 변환이 있습니다. 입력 X에서 Q, K, V를 만드는 과정을 QKV Projection, 각 헤드의 어텐션 결과를 최종 출력 특징으로 바꾸는 과정을 Output Projection이라고 부릅니다. 여기서 Projection은 가중치 행렬을 곱해 특징을 변환한다는 뜻입니다.

일반적인 멀티헤드 어텐션에서는 각 헤드가 자기 Q, K, V로 어텐션을 계산합니다. 따라서 헤드 하나의 계산에 필요한 값을 한 GPU에 함께 배치하면, 그 헤드의 어텐션은 다른 GPU의 헤드 결과 없이 계산할 수 있습니다. 그림에서는 두 헤드 중 하나씩을 두 GPU가 맡습니다.

이를 위해 Q, K, V의 가중치 WQ, WK, WV를 각각 헤드에 대응하는 열로 나눕니다. 두 GPU는 같은 X를 사용하지만, 서로 다른 헤드의 Q, K, V를 만듭니다. 각 GPU의 Q, K, V에는 자기 헤드가 사용할 두 토큰의 정보가 모두 들어 있습니다. 토큰을 반으로 나누는 방식과는 다릅니다.

각 GPU는 자기 Q와 K로 점수를 구하고, Softmax를 적용한 뒤 V를 가중합해 헤드 출력을 얻습니다. GPU 0의 출력을 O₀, GPU 1의 출력을 O₁이라고 하겠습니다. 여기까지 각 헤드의 계산은 자기 GPU 안에서 진행됩니다.

이제 모든 헤드 출력을 모아서 Output Projection을 수행할 수도 있지만, FFN에서처럼 모으는 단계를 생략할 수 있습니다. 출력 가중치 WO의 행을 각 GPU가 가진 헤드 출력에 맞춰 나누면 됩니다. GPU 0은 O₀와 그에 대응하는 WO의 행을 곱하고, GPU 1은 O₁과 나머지 행을 곱합니다. 두 결과는 같은 최종 출력에 대한 부분합이므로, 마지막에 All-Reduce로 더합니다.

전체 헤드 출력을 이어 붙인 뒤 WO를 곱한 결과와, 각 헤드 출력에 대응하는 WO 조각을 곱하고 합산한 결과는 같습니다. 그래서 헤드 출력을 중간에 모으지 않고 Output Projection까지 진행할 수 있습니다.

FFN과 어텐션은 가운데 계산이 다르지만 연결 방식은 비슷합니다. 첫 투영은 열로 나누고, GPU별로 독립적인 계산을 이어간 다음, 출력 투영은 행으로 나누어 최종 부분합을 합산합니다. 다만 어텐션 전체가 행렬 곱 두 번으로 끝난다는 뜻은 아닙니다. Q·K·V를 만드는 투영과 출력 투영 사이에는 점수 계산과 V의 가중합도 있습니다.

이 그림은 헤드마다 Q, K, V가 있는 일반적인 멀티헤드 어텐션을 전제로 합니다. K·V를 여러 헤드가 공유하는 구조에서는 어떤 헤드와 데이터를 함께 배치하거나 복제할지 추가로 살펴봐야 합니다.

레이어 전체에서 통신 위치 확인하기

이제 두 GPU가 트랜스포머 레이어 하나를 함께 처리하는 흐름으로 연결해 보겠습니다. 그림은 어텐션과 FFN 앞에 정규화(Norm)를 두는 Pre-Norm 구조입니다. 정규화는 입력 특징의 크기를 조정하는 연산이며, 잔차 연결은 블록에 들어가기 전 입력을 블록의 출력에 더하는 경로입니다.

레이어에 들어올 때 두 GPU는 같은 입력 X 전체를 갖습니다. 정규화 후 어텐션에서는 QKV 투영과 헤드별 계산, 출력 투영까지 각자 진행합니다. 그다음 All-Reduce로 부분합을 합산해 두 GPU에 같은 완성 결과를 준비합니다. 각 GPU는 자신이 보관한 입력을 여기에 더해 잔차 연결을 계산합니다.

FFN에서도 같은 흐름이 반복됩니다. 정규화 후 첫 Linear와 활성화, 두 번째 Linear까지 각자 계산하고, All-Reduce로 합산합니다. 잔차 연결까지 마치면 두 GPU에 같은 레이어 출력이 남아 다음 레이어로 넘어갑니다.

따라서 그림에 나타낸 기본 TP의 레이어 하나에서는 순전파 All-Reduce가 두 번 필요합니다. 어텐션의 출력 투영 뒤에 한 번, FFN의 두 번째 Linear 뒤에 한 번입니다. 어텐션의 헤드 출력을 모으는 통신과 FFN의 중간 특징을 모으는 통신은 열·행 분할을 연결해 생략했습니다. 이 배치는 Megatron-LM 논문에서 설명하는 기본 TP 구조입니다.

트랜스포머는 이런 레이어를 여러 번 쌓습니다. 같은 구조의 레이어가 L개라면, 그림의 두 합산만으로도 순전파에 All-Reduce가 2L회 반복됩니다. 레이어 하나에서 중간 통신을 없애는 효과도 레이어 수만큼 누적됩니다. 한 번의 행렬 곱을 잘 나누는 것에 더해, 다음 연산까지 조각을 그대로 이어 쓰는 것이 중요한 이유입니다. 이 횟수는 그림에 표시한 두 블록의 순전파만 센 것으로, 역전파나 임베딩 등 모델의 다른 부분에서 필요한 통신까지 포함한 값은 아닙니다.

TP는 한 입력을 처리하는 연산과 가중치를 나눠 GPU당 부담을 줄일 수 있습니다. 하지만 계산을 나눈 만큼 실행 시간이 줄어든다고 보장할 수는 없습니다. GPU 사이의 통신과 대기에 시간이 들고, 너무 작은 행렬로 나누면 각 GPU의 계산 효율도 떨어질 수 있습니다. 계산에서 절약한 시간이 통신 비용보다 충분히 큰지, 같은 입력의 전체 완료 시간이 실제로 줄었는지를 함께 확인해야 합니다.

또한 그림의 정규화와 잔차 연결 구간을 보면, 두 GPU가 같은 전체 입력과 중간값을 보관하며 같은 계산을 합니다. 가중치를 나눴다고 모든 데이터가 나뉘는 것은 아닙니다. 이렇게 복제되어 남은 일부 중간값을 토큰 방향으로 나누는 방법은 SP를 다룰 때 살펴보겠습니다.