CP: 긴 문맥을 여러 GPU로 나누기

문맥이 길어지면 처리할 토큰과 활성값도 늘어납니다. 토큰을 여러 GPU에 나누면 각 GPU의 부담을 줄일 수 있지만, 어텐션에서는 한 가지 문제가 생깁니다. 내가 맡은 토큰을 계산하는 데 다른 GPU가 맡은 토큰의 정보도 필요하다는 점입니다. 토큰을 나눈 뒤에도 원래 모델의 참조 관계를 유지해야 합니다.

이번 글에서는 먼저 시퀀스 병렬화(SP)와 문맥 병렬화의 분할 범위를 비교하고, 어텐션에서 왜 GPU 사이의 정보 교환이 필요한지 확인하겠습니다. 이어서 Q를 두고 K·V를 순환시키는 Ring과, 토큰별 데이터를 헤드별로 재배치하는 Ulysses를 살펴봅니다. 마지막에는 인과적 어텐션에서 같은 수의 토큰을 나눠도 작업량이 달라지는 이유를 보겠습니다.

그림은 여러 토큰을 함께 처리하는 순전파를 기준으로 합니다. CP만 적용해 모델 가중치는 각 GPU가 동일하게 가지고 있다고 가정합니다. TP와 CP를 함께 사용하는 구성은 이후에 연결할 수 있지만, 여기서는 토큰의 배치와 어텐션에 필요한 정보의 이동에 집중하겠습니다.

SP에서 CP로 분할 범위 넓히기

앞서 다룬 Megatron 방식의 SP는 TP 사이의 정규화와 잔차 연결 구간에서 활성값을 토큰별로 나눴습니다. TP로 들어가기 전에는 All-Gather로 입력을 모았으므로, TP 어텐션 안에서는 각 GPU가 담당 헤드의 모든 토큰을 가지고 계산했습니다.

문맥 병렬화(Context Parallelism, CP)는 어텐션을 포함한 모델 실행에서 긴 문맥의 토큰들을 나눠 맡도록 하는 접근입니다. 같은 가중치를 가진 GPU들이 자기 토큰의 Q·K·V Projection이나 FFN을 계산하고, 어텐션에서는 다른 토큰과의 참조를 통신으로 연결합니다. NVIDIA의 CP 설명도 기존 SP와 CP를 적용하는 연산 범위로 구별합니다.

두 방식 모두 그림 위쪽의 정규화에서는 일부 토큰의 전체 특징을 갖습니다. 차이가 드러나는 곳은 어텐션과 FFN입니다. 왼쪽의 SP+TP는 모든 토큰에 대해 특징이나 헤드를 나눠 계산합니다. 오른쪽의 CP 예시는 같은 모델 가중치로 담당 토큰을 계산합니다. FFN은 다른 토큰의 값을 사용하지 않으므로 자기 토큰만으로 수행할 수 있지만, 어텐션에는 다른 GPU의 정보가 필요합니다.

여기서 “CP는 모든 순간 같은 토큰 배치를 유지한다”는 뜻은 아닙니다. 담당 토큰의 출력을 만들기 위해 K·V를 가져올 수도 있고, 어텐션 내부에서 분할 축을 잠시 바꿀 수도 있습니다. 그림 오른쪽은 각 GPU가 어떤 토큰의 Q와 최종 출력을 맡는지를 보여줍니다.

또한 문헌마다 이름의 범위가 다릅니다. 뒤에서 다룰 Ulysses도 원문에서는 Sequence Parallelism이라는 이름을 사용합니다. 이 글에서는 앞선 글의 SP를 Megatron 방식으로 한정하고, Ulysses는 긴 문맥의 어텐션을 여러 GPU에 나누는 방법으로 함께 살펴보겠습니다.

어텐션에서 다른 GPU의 정보가 필요한 이유

토큰 t0·t1을 GPU 0에, t2·t3를 GPU 1에 둔다고 하겠습니다. 각 GPU는 자기 토큰에 가중치를 곱해 Q·K·V를 만듭니다. 이 Projection 단계는 다른 토큰의 값을 필요로 하지 않습니다.

다음 어텐션 계산은 다릅니다. Q는 K와의 점수로 어떤 위치를 얼마나 참조할지 정하고, 그 가중치로 V를 조합합니다. 예를 들어 인과적 어텐션에서 t3는 자기 자신과 앞선 t0·t1·t2를 모두 참조할 수 있습니다. GPU 1에 t3의 Q가 있어도, t0·t1의 K·V가 없으면 원래 어텐션을 완성할 수 없습니다.

그림의 행은 Q의 토큰, 열은 K의 토큰입니다. ●는 계산할 수 있는 참조이고, 미래 토큰은 마스킹되어 있습니다. Q(t3)의 행을 따라가면 GPU 0에 있는 K(t0)·K(t1)까지 필요하다는 점이 드러납니다. 그 점수에 대응하는 V(t0)·V(t1)도 출력 계산에 참여해야 합니다.

이 행렬은 토큰 사이의 관계를 표시한 것입니다. 실제 구현에서 전체 점수 행렬을 한 번에 메모리에 만들어야 한다는 뜻은 아닙니다. 필요한 참조를 유지하면서 데이터를 블록 단위로 가져오거나, 계산에 맞게 재배치할 수 있습니다.

Ring으로 K·V 블록 순환시키기

첫 번째 방법은 Q를 담당 GPU에 두고, K·V 블록을 차례로 가져오는 것입니다. GPU들을 원형으로 연결하고, 각 GPU가 가진 K·V 블록을 다음 GPU에 전달합니다. 도착한 블록으로 자기 Q의 어텐션 계산을 진행한 뒤, 다음 블록을 이어서 처리합니다. Ring Attention이 사용하는 핵심 구조입니다.

다음 그림은 한 헤드의 토큰 여덟 개를 GPU 네 개가 두 개씩 맡는 예시입니다. 전달 순서를 명확히 보기 위해 모든 토큰을 참조하는 경우를 먼저 보여줍니다. 인과적 마스크가 있는 경우의 작업량은 뒤에서 따로 살펴보겠습니다.

GPU 0을 따라가 보겠습니다. Q0는 t0·t1의 질의이며 계산 내내 GPU 0에 남습니다. 먼저 로컬 K0·V0로 t0·t1에 대한 기여를 계산합니다. 한 번 전달하면 GPU 3에서 K3·V3가 도착해 t6·t7을 처리하고, 이어 K2·V2, K1·V1로 나머지 토큰을 처리합니다. 세 번의 전달을 마치면 로컬 블록을 포함한 네 블록의 기여가 모두 반영됩니다. 다른 GPU도 자기 Q에 대해 같은 과정을 진행합니다.

그림 아래의 채워지는 칸은 GPU 0이 어떤 Q–K 토큰 쌍을 처리했는지 보여줍니다. K·V는 이동하지만, 완성하려는 출력은 계속 자기 Q의 토큰에 대한 출력입니다. 전체 K·V를 한꺼번에 모아 두는 대신, 도착한 블록을 계산하고 누적할 수 있습니다.

다만 블록별 어텐션 출력을 단순히 더하거나 평균내면 안 됩니다. Softmax의 가중치는 해당 Q가 참조하는 전체 점수를 기준으로 정해지기 때문입니다. Online Softmax에서 살펴본 것처럼, 가중합과 함께 최댓값·분모를 관리하고 새 블록의 기준에 맞춰 누적값을 보정해야 합니다. FlashAttention의 블록별 계산이 GPU 사이의 데이터 전달로 이어지는 지점입니다.

앞서 살펴본 Ring All-Reduce와는 전달하는 데이터와 계산 목적이 다릅니다. Ring All-Reduce는 GPU들의 값을 합산해 공유하는 집합 연산입니다. 여기서는 K·V 블록을 전달받아 각 GPU가 자기 Q의 어텐션 출력을 계산합니다. K·V 자체를 서로 더하는 과정은 아닙니다.

현재 블록을 계산하는 동안 다음 블록을 받을 수 있다면 통신 대기를 줄일 여지가 있습니다. 그러나 실제로 얼마나 겹칠 수 있는지는 블록의 계산 시간, 연결의 속도, 버퍼와 실행 방식에 달려 있습니다. 그림의 ‘다음’은 데이터 의존성을 따라가는 단계이며, 측정한 실행 시간표를 뜻하지는 않습니다.

Ulysses로 토큰 분할을 헤드 분할로 바꾸기

다른 방법은 한 헤드의 전체 문맥을 한 GPU에 모아 어텐션을 계산하는 것입니다. DeepSpeed Ulysses는 이를 위해 Q·K·V의 배치를 바꿉니다.

토큰 네 개와 헤드 두 개를 예로 들겠습니다. 처음에는 GPU 0이 t0·t1의 모든 헤드를, GPU 1이 t2·t3의 모든 헤드를 갖습니다. 각 GPU가 자기 토큰에서 Q·K·V를 만든 뒤, 헤드 h0의 데이터는 GPU 0으로 보내고 h1의 데이터는 GPU 1로 보냅니다.

이때 사용하는 All-to-All은 각 GPU가 상대별로 다른 조각을 보내고 받는 집합 통신입니다. All-Gather처럼 모든 데이터를 모든 GPU에 복제하는 것이 아니라, 각 GPU에 필요한 헤드 조각을 배정합니다.

첫 All-to-All이 끝나면 GPU 0에는 t0부터 t3까지의 h0 데이터가, GPU 1에는 같은 네 토큰의 h1 데이터가 있습니다. 토큰별로 나눴던 배열이 헤드별로 나뉜 배열로 바뀐 것입니다. 각 GPU는 담당 헤드의 Q·K·V를 전체 문맥에 걸쳐 갖고 있으므로, 그 헤드의 어텐션을 로컬에서 완성할 수 있습니다.

계산 후에는 배치를 다시 돌려놓습니다. GPU 0에서 계산한 h0 출력 중 t2·t3 부분은 GPU 1로, GPU 1에서 계산한 h1 출력 중 t0·t1 부분은 GPU 0으로 보냅니다. 이 역방향 All-to-All이 옮기는 것은 Q·K·V가 아니라 계산된 어텐션 출력 O입니다.

그 결과 GPU 0에는 t0·t1의 모든 헤드 출력이, GPU 1에는 t2·t3의 모든 헤드 출력이 모입니다. 헤드 출력을 이어 붙여 Output Projection을 계산하고, 잔차 연결과 정규화를 거쳐 FFN으로 진행할 수 있습니다. 다시 자기 토큰을 담당하는 배치로 돌아온 것입니다.

TP에서도 헤드를 나눴기 때문에 비슷하게 느껴질 수 있습니다. 하지만 이 그림에서는 가중치를 TP로 나눈 것이 아니라, 복제된 가중치로 만든 활성값의 배치를 바꿨습니다. Ring이 고정된 Q 쪽으로 K·V를 가져온다면, Ulysses는 Q·K·V를 함께 옮겨 담당 헤드의 전체 문맥을 구성합니다.

인과적 어텐션의 작업량 맞추기

이번에는 미래 토큰을 참조하지 않는 인과적 어텐션을 보겠습니다. 앞쪽 토큰은 참조할 위치가 적고 뒤쪽 토큰은 많습니다. 따라서 같은 수의 토큰을 맡는다고 어텐션 계산량까지 같아지는 것은 아닙니다.

토큰 여덟 개를 GPU 두 개가 네 개씩 맡는다면, 앞의 t0부터 t3까지를 맡은 GPU가 계산할 토큰 쌍은 1+2+3+4 = 10개입니다. 뒤의 t4부터 t7까지를 맡은 GPU는 5+6+7+8 = 26개를 계산합니다.

‘다음’을 누르면 GPU 0이 t0·t1·t6·t7을, GPU 1이 t2·t3·t4·t5를 맡습니다. 앞쪽과 뒤쪽 토큰을 묶어 배치하면 이 예시의 계산할 쌍 수는 각각 18개로 같아집니다. 그림에서 바뀌는 것은 행의 담당 GPU를 나타내는 색이며, 참조를 허용하는 삼각형은 그대로입니다.

즉, 토큰의 원래 위치와 인과적 마스크는 바꾸지 않고 담당 GPU만 바꿉니다. 이 구분을 유지해야 원래 문맥의 어텐션을 계산할 수 있습니다. 실제 구현에서는 이런 배치를 통신 순서와 계산 블록 크기에 맞춰 구성해야 합니다. 토큰 쌍 수를 맞췄다고 실행 시간까지 정확히 같아지는 것은 아니지만, 불균형이 생기는 이유와 조정 방향을 보여줍니다.

메모리 절약과 통신 비용 함께 보기

Ring과 Ulysses는 모두 문맥을 나눈 상태에서 원래 어텐션을 완성하려는 방법이지만, 필요한 정보를 배치하는 방식이 다릅니다.

방식	어텐션 계산을 위한 배치	주요 통신
Ring	자기 Q를 유지하고 K·V 블록을 차례로 처리	GPU 사이의 K·V 블록 전달
Ulysses	담당 헤드의 전체 문맥을 모아 처리	Q·K·V All-to-All과 출력의 역 All-to-All

이 글의 CP 단독 예시는 토큰별 활성값과 계산을 나누지만, 모델 가중치는 각 GPU가 동일하게 갖습니다. 전체 GPU 메모리가 분할 수에 반비례하는 것은 아닙니다. 통신 중 동시에 보관하는 블록과 버퍼, 역전파를 위해 남길 활성값까지 포함해 최고 메모리 사용량을 봐야 합니다.

Ulysses의 기본 구성은 헤드를 GPU들에 나눠 맡길 수 있어야 합니다. 그림처럼 헤드 두 개를 GPU 두 개에 하나씩 배정하는 경우는 단순하지만, 헤드 수와 병렬화 크기, Q와 K·V의 헤드 구조가 달라지면 적용 조건도 달라집니다. Ring 역시 GPU를 늘릴수록 로컬 계산은 작아지고 전달 단계는 늘어날 수 있으므로, 계산과 통신의 균형을 확인해야 합니다.

학습에서는 역전파의 기울기를 전달하는 과정도 필요합니다. 또한 한 번에 여러 입력 토큰을 처리하는 경우와, 토큰을 하나씩 생성하며 기존 K·V를 읽는 추론 단계는 작업 구성이 다릅니다. 여기서 본 순전파 그림을 그대로 모든 상황의 실행 구조로 옮길 수는 없습니다.

지금까지 TP는 같은 입력의 특징별 계산을, SP는 TP 사이의 활성값을, CP는 긴 문맥과 그 참조 관계를 나누는 관점으로 살펴봤습니다. 다음 파이프라인 병렬화(PP)에서는 같은 레이어 안의 분할에서 벗어나, 모델의 서로 다른 레이어를 다른 GPU에 배치하는 방법을 보겠습니다.