SP: TP와 함께 활성값 나누기

텐서 병렬화(TP)에서는 여러 GPU가 같은 입력의 연산을 나누어 계산했습니다. 어텐션은 헤드를 나누고, FFN은 중간 특징을 나눠 계산한 뒤 마지막 부분합을 합쳤습니다. 그러나 이렇게 큰 행렬 곱을 나누더라도, 레이어의 모든 데이터와 계산이 나뉘는 것은 아닙니다. 부분합을 All-Reduce하면 같은 완성 결과가 각 GPU에 생기고, 이어지는 정규화에서는 같은 입력을 중복해서 처리하게 됩니다.

모델 계산 중 만들어지는 이러한 중간값을 활성값(activation)이라고 합니다. 이번 글에서는 먼저 TP 사이에 어떤 활성값이 복제되어 남는지 살펴보겠습니다. 이어서 정규화와 잔차 연결을 토큰별로 나눠 처리할 수 있는 이유를 확인하고, TP 구간으로 들어가고 나올 때 데이터 배치를 어떻게 바꾸는지 따라갑니다. 마지막에는 이 경계들을 트랜스포머 레이어 전체에 연결하겠습니다.

여기서 시퀀스 병렬화(Sequence Parallelism, SP)는 Megatron에서 제안한 방식을 뜻합니다. 같은 입력의 토큰들을 나누며, DP처럼 서로 다른 요청이나 학습 샘플을 GPU마다 배정하는 방식과는 구별합니다. 그림은 GPU 두 개가 순전파를 수행하는 예시입니다.

TP 사이에 남는 활성값의 중복

먼저 TP만 적용한 레이어를 보겠습니다. 그림의 작은 행렬에서 행은 토큰, 열은 특징입니다. 어텐션 내부에서는 각 GPU가 담당 헤드의 값을, FFN 내부에서는 담당 중간 특징의 값을 갖습니다. 두 GPU가 서로 다른 부분을 계산하는 구간입니다.

하지만 어텐션의 Output Projection이나 FFN의 마지막 Linear가 만든 결과는 아직 부분합입니다. 앞서 본 기본 TP에서는 이를 All-Reduce로 더하고, 두 GPU가 같은 결과 전체를 갖게 했습니다.

그림에서 정규화와 잔차 연결 구간의 행렬은 두 GPU에 똑같이 있습니다. 토큰이 네 개이고 각 토큰의 특징이 네 개라면, 두 GPU 모두 같은 4×4 배열을 보관합니다. 정규화도 같은 네 토큰을 두 번 처리합니다.

SP는 이 구간에서 각 GPU가 일부 토큰의 활성값만 갖고 계산하도록 배치를 바꿉니다. TP가 이미 나누고 있는 어텐션과 FFN의 큰 행렬 곱은 유지하면서, 그 사이에 남은 복제와 중복 계산을 줄이는 것입니다.

토큰별로 정규화와 잔차 연결하기

토큰 t0·t1은 GPU 0에, t2·t3는 GPU 1에 둔다고 하겠습니다. 각 GPU가 갖는 활성값은 4×4 전체에서 2×4 조각으로 줄어듭니다. 여기서 중요한 것은 한 토큰의 특징 네 개는 함께 둔다는 점입니다. 특징을 반으로 자르는 것이 아니라, 온전한 토큰 벡터를 두 개씩 나눠 갖습니다.

이 배치라면 정규화에 필요한 값이 한 GPU 안에 있습니다. RMSNorm은 한 토큰의 특징값들을 이용해 그 토큰의 크기를 조정합니다. 다른 토큰의 특징값은 필요하지 않으므로, GPU 0은 t2·t3를 받지 않고도 t0·t1을 정규화할 수 있습니다.

잔차 연결도 같은 토큰의 대응하는 특징값끼리 더하는 연산입니다. GPU 0에 t0·t1의 연산 결과와 잔차가 함께 있으면 덧셈을 끝낼 수 있습니다. 따라서 잔차로 남겨 둔 입력 역시 결과와 같은 토큰 배치를 유지해야 합니다. 그림 아래의 덧셈은 이 조건을 보여줍니다.

여기서 “FFN은 토큰별로 독립적이지 않아서 TP를 쓴다”고 이해하면 안 됩니다. FFN도 각 토큰을 독립적으로 처리합니다. 다만 지금 구성에서는 큰 가중치 행렬과 그 계산을 TP로 나누고 있습니다. SP는 그 구성을 유지하면서 정규화·잔차 구간의 활성값을 나누는 방법입니다. 어떤 축으로 나눌지는 토큰 간 의존성뿐 아니라, 가중치가 어디에 있고 다음 연산이 어떤 입력을 필요로 하는지까지 함께 보고 정합니다.

SP와 TP 사이에서 데이터 배치 바꾸기

정규화를 마쳤다면 이제 TP의 첫 Linear를 계산해야 합니다. 각 GPU는 가중치의 서로 다른 열 조각을 가지고 있습니다. GPU 0은 모든 토큰의 앞쪽 출력 특징을, GPU 1은 모든 토큰의 뒤쪽 출력 특징을 계산해야 하므로, 두 GPU 모두 모든 토큰의 입력이 필요합니다.

그래서 All-Gather로 토큰 조각을 모읍니다. SP 구간의 t0·t1과 t2·t3가 합쳐져 두 GPU에 같은 입력 X가 생깁니다. 이후에는 앞서 본 TP의 열 분할과 행 분할을 그대로 연결합니다. 다음 그림의 ‘다음’을 눌러 작은 FFN의 계산을 따라가 보겠습니다.

첫 Linear에서는 X에 가중치 U의 열 조각을 곱합니다. GPU 0은 중간 특징 f0·f1을, GPU 1은 f2·f3를 계산합니다. 모든 토큰을 처리하되, 서로 다른 특징을 맡는 것입니다. 이 값들은 각각 완성된 출력 특징이므로 원소별 활성화를 바로 적용할 수 있습니다. 그림은 양수를 그대로 유지하는 ReLU의 작은 예시이며, 실제 모델에서는 다른 활성화를 사용할 수 있습니다.

둘째 Linear는 이 특징 조각과 가중치 V의 대응하는 행 조각을 곱합니다. 중간값을 모으는 통신은 필요하지 않지만, 그 결과 P0와 P1은 같은 출력 위치에 대한 부분합입니다. 예를 들어 t0의 첫 출력 특징은 GPU 0의 4와 GPU 1의 5를 더해야 완성된 값 9가 됩니다.

여기서 다음 구간도 SP로 처리한다면, 모든 GPU가 완성된 출력 전체를 가질 필요는 없습니다. Reduce-Scatter로 부분합을 더하면서 결과를 토큰별로 나눠 갖습니다. 그림의 마지막 단계에서 GPU 0에는 t0·t1의 완성된 출력이, GPU 1에는 t2·t3의 완성된 출력이 남습니다. 각 GPU의 부분합을 단순히 반으로 잘라 쓰는 것이 아니라, 다른 GPU의 기여까지 합쳐 자기 토큰의 값을 완성한 것입니다.

이 과정에서 GPU 그룹이 바뀌지는 않습니다. 같은 두 GPU가 SP 구간에서는 토큰을 나누고, TP 내부에서는 중간 특징을 나눕니다. All-Gather와 Reduce-Scatter가 두 배치를 이어 주는 경계입니다.

레이어 전체에서 SP와 TP 연결하기

지금까지 FFN을 따라갔지만 어텐션 앞뒤에도 같은 구조가 있습니다. 아래 그림은 정규화를 먼저 수행하는 Pre-Norm 레이어입니다. 점선은 같은 GPU에 남겨 둔 잔차가 나중에 더해지는 경로이고, 보라색은 두 GPU가 참여하는 집합 통신입니다.

먼저 각 GPU가 자기 토큰을 정규화한 뒤, All-Gather로 입력을 모읍니다. TP의 QKV Projection은 그 입력으로 담당 헤드의 Q·K·V를 만듭니다. 이 시점에는 담당 헤드의 모든 토큰이 한 GPU에 있으므로, 그 헤드의 어텐션을 계산하기 위해 다른 GPU의 토큰을 추가로 받아 올 필요가 없습니다.

헤드별 계산을 Output Projection까지 이어가면 각 GPU에 부분합이 생깁니다. Reduce-Scatter로 이를 합치고 토큰별로 나눈 뒤, SP 구간에서 잔차를 더하고 다시 정규화합니다. 이어지는 FFN도 All-Gather → TP 계산 → Reduce-Scatter 순서로 진행하며, 마지막 잔차 연결 후에는 담당 토큰의 활성값을 다음 레이어로 넘깁니다.

이 구조는 기존 TP의 All-Reduce에 통신을 단순히 덧붙인 것이 아닙니다. 집합 통신에서 살펴본 것처럼, All-Reduce의 결과는 Reduce-Scatter 뒤에 All-Gather를 수행해 얻을 수 있습니다. SP에서는 합산된 조각만 필요한 구간을 두 동작 사이에 놓고, 다음 TP 계산이 전체 입력을 요구할 때 모읍니다. 그 사이에 정규화와 잔차 연결을 조각 상태로 처리할 수 있어야 이 배치가 성립합니다.

활성값 절약과 전체 실행 비용

SP의 직접적인 이점은 해당 구간에서 각 GPU가 보관하고 처리하는 활성값을 줄이는 것입니다. 학습에서는 역전파를 위해 활성값을 남겨 두는 경우가 많아, 시퀀스가 길어질수록 이 절약이 중요해질 수 있습니다.

그렇다고 GPU의 전체 메모리 사용량이 모두 절반이 되는 것은 아닙니다. TP 구간에서는 All-Gather한 입력이 필요하고, 통신 버퍼나 다른 중간값도 존재합니다. 가중치 배치 역시 SP 자체가 바꾸는 대상은 아닙니다. 따라서 그림의 특정 배열 크기와 실제 실행 중의 최고 메모리 사용량을 구별해야 합니다.

통신도 없어지지 않습니다. 같은 데이터 크기의 All-Reduce와 Reduce-Scatter·All-Gather 조합이 전달량 측면에서 대응하더라도, 호출 시점과 버퍼, 계산과 겹치는 정도에 따라 실행 시간은 달라질 수 있습니다. 추론은 학습과 달리 한 번에 처리하는 토큰 수와 보관할 상태가 다르므로, 학습에서의 효과를 그대로 기대할 수는 없습니다.

이번 구성에서는 TP 어텐션으로 들어가기 전에 전체 토큰의 입력을 모았습니다. 그렇다면 어텐션에서도 각 GPU가 일부 토큰만 맡도록 하려면 어떻게 해야 할까요? 내 토큰의 Q가 다른 GPU의 K·V를 참조하는 문제를 해결해야 합니다. 다음 문맥 병렬화(CP)에서는 이 연결을 유지하며 긴 문맥을 나누는 방법을 살펴보겠습니다.