모델 연산의 병렬성: 원소별 연산, Reduction, 행렬 곱

지난 글에서는 GPU의 연산 장치와 메모리, 그리고 데이터를 읽고 계산해 저장하는 흐름을 살펴봤습니다. 많은 연산 장치를 활용하려면 동시에 처리할 수 있는 계산이 있어야 합니다. 이번에는 모델에서 사용하는 연산을 다시 살펴보며 어떤 결과를 독립적으로 계산할 수 있고, 결과 하나를 만들려면 어떤 계산이 필요한지 알아보겠습니다.

벡터나 행렬을 이루는 각각의 값을 요소 또는 원소라고 부릅니다. 벡터에서는 각 위치의 값, 행렬에서는 각 행과 열이 만나는 칸의 값입니다. 이번 글에서는 출력의 각 요소를 기준으로 원소별 연산, reduction, 행렬 곱을 비교하겠습니다.

원소별 연산: 각 위치에서 독립적으로 계산하기

모델 편에서 살펴본 Residual 연결은 두 벡터의 같은 위치에 있는 값을 더합니다. 첫 번째 위치의 두 값을 더해 첫 번째 출력을 만들고, 두 번째 위치의 두 값을 더해 두 번째 출력을 만드는 식입니다. 이처럼 각 위치의 입력에 연산을 적용해 같은 위치의 출력을 만드는 연산을 원소별 연산(element-wise operation)이라고 합니다.

그림 1에서 첫 번째 결과인 11을 구하려면 입력 1과 10이 필요합니다. 두 번째 결과인 22를 구하려면 2와 20이 필요합니다. 22를 계산하는 데 11은 필요하지 않습니다. 따라서 각 위치의 입력이 준비되어 있다면 네 결과를 서로 독립적으로 계산할 수 있습니다.

MLP에서 사용하는 활성화 함수도 각 요소에 독립적으로 적용할 수 있습니다. Gating에서 두 벡터의 같은 위치끼리 곱하는 연산 역시 원소별 연산입니다. 입력 하나에 함수를 적용하는지, 입력 두 개를 더하거나 곱하는지는 다르지만, 한 위치의 출력이 다른 위치의 출력에 의존하지 않는다는 점은 같습니다.

Reduction: 여러 입력을 모아 결과 하나 만들기

이번에는 벡터의 모든 값을 더해 합계 하나를 만든다고 생각해 보겠습니다. 결과를 구하려면 특정 위치의 값 하나만으로는 부족하고, 합산 대상인 모든 값을 반영해야 합니다. 이처럼 여러 입력 값을 합이나 최댓값 같은 하나의 결과로 모으는 연산을 reduction이라고 합니다. 여기서는 합계를 구하는 경우를 살펴보겠습니다.

행렬의 각 행을 따로 합산하면 행마다 결과가 하나씩 나옵니다. 두 행이 있다면 첫 번째 행의 합과 두 번째 행의 합을 독립적으로 계산할 수 있습니다. 그림 2는 두 행의 합을 각각 구하는 과정입니다.

첫 번째 행에서는 1부터 8까지를 더해 36을 만듭니다. 두 번째 행에서는 2와 1이 네 번 반복된 값을 더해 12를 만듭니다. 첫 번째 행의 합을 구하는 데 두 번째 행의 합은 필요하지 않습니다. 출력 요소를 기준으로 보면 두 결과를 나누어 계산할 수 있습니다.

하지만 출력 하나를 만드는 과정은 원소별 덧셈과 다릅니다. 첫 번째 행의 결과인 36에는 여덟 입력 값이 모두 반영되어야 합니다. 그림에서는 먼저 이웃한 값끼리 더해 3, 7, 11, 15라는 부분합을 만듭니다. 그다음 부분합을 더해 10과 26을 만들고, 마지막으로 두 값을 더합니다.

합계를 만드는 과정 안에도 동시에 계산할 수 있는 부분이 있습니다. 첫 단계의 1 + 2와 3 + 4는 서로 독립적입니다. 다만 다음 단계에서 3 + 7을 계산하려면 두 부분합이 먼저 준비되어야 합니다. 같은 단계의 독립적인 계산과, 단계 사이에서 앞선 결과를 기다리는 관계를 구분할 수 있습니다. 부동소수점 값을 더할 때는 합산 순서에 따라 반올림 차이가 생길 수 있지만, 여기서는 작은 정수로 계산의 관계에 집중하겠습니다.

RMSNorm에서는 한 토큰 벡터의 요소들을 제곱한 뒤 합하는 과정에 reduction이 들어갑니다. Softmax에서도 지수 함수를 적용한 값들의 합을 구할 때 reduction이 필요합니다. 두 연산 모두 이 합계뿐 아니라 원소별 계산도 포함합니다. 모델에서 이름 붙인 연산 하나가 여러 종류의 계산으로 이루어질 수 있습니다.

행렬 곱: 출력 요소마다 곱하고 누적하기

MLP와 Attention의 projection에서는 입력 행렬과 가중치 행렬을 곱합니다. X의 크기가 M × K이고 W의 크기가 K × N이면, 출력 Y의 크기는 M × N입니다. M과 N은 출력의 행과 열 개수이고, K는 출력 하나를 구할 때 곱해서 더하는 입력 쌍의 개수입니다.

그림 3에서 X는 2 × 3, W는 3 × 4이므로 출력은 2 × 4이고 결과 요소는 여덟 개입니다. 첫 번째 출력인 7은 X의 첫 행 [1, 2, 3]과 W의 첫 열 [1, 0, 2]를 사용합니다. 같은 위치의 값을 곱하면 1, 0, 6이 되고, 이를 더하면 7입니다.

바로 옆의 출력인 4는 X의 같은 행과 W의 두 번째 열 [2, 1, 0]으로 구합니다. 계산은 1 × 2 + 2 × 1 + 3 × 0 = 4입니다. 이 계산에는 앞에서 구한 출력 7이 필요하지 않습니다. 나머지 출력도 각각 해당하는 X의 행과 W의 열만 있으면 구할 수 있으므로, 행렬 곱은 출력 요소 단위로 병렬화할 수 있습니다.

한편 출력 하나를 만드는 과정에는 여러 곱을 합치는 reduction이 포함됩니다. 그림에서는 K가 3이므로 곱 세 개를 더합니다. K가 커지면 출력의 행과 열 개수가 같더라도 출력 하나에 필요한 곱과 누적이 늘어납니다.

또한 출력이 독립적이라는 말이 사용하는 입력까지 서로 다르다는 뜻은 아닙니다. 출력 Y의 같은 행에 있는 요소들은 X의 같은 행을 사용하고, Y의 같은 열에 있는 요소들은 W의 같은 열을 사용합니다. 서로의 출력은 기다릴 필요가 없으면서도 입력의 일부는 공통으로 사용하는 것입니다. 이 성질은 뒤에서 데이터를 재사용하며 계산하는 방법을 이해할 때 중요해집니다.

독립적인 결과의 개수와 결과 하나의 계산량

앞의 세 연산을 출력 요소를 기준으로 나란히 비교해 보겠습니다. 원소별 덧셈은 같은 위치의 입력 두 개를, 행별 합계는 한 행의 여러 입력을, 행렬 곱은 한 행과 한 열의 입력을 사용합니다. 세 경우 모두 서로 다른 출력 요소를 독립적으로 계산할 수 있지만, 출력 하나에 필요한 입력의 범위와 계산량은 다릅니다.

그림 4의 왼쪽에서는 출력 11을 만들기 위해 1과 10을 한 번 더합니다. 가운데에서는 출력 36을 만들기 위해 한 행의 여덟 값을 합합니다. 오른쪽에서는 출력 7을 만들기 위해 세 쌍의 값을 곱하고 그 결과를 더합니다. 같은 출력 요소 하나라도 그 안에 들어 있는 계산은 같지 않습니다.

따라서 연산의 병렬성을 볼 때는 먼저 독립적으로 구할 결과 요소가 몇 개 있는지 살펴볼 수 있습니다. 그다음 각 결과를 만들려면 어떤 입력과 계산이 필요한지 확인합니다. 행렬 곱에서는 M·N이 출력 요소의 개수를 정하고, K가 각 요소에 필요한 곱과 누적의 양을 정합니다. 행별 합계에서도 행 수와 한 행의 길이는 서로 다른 역할을 합니다.

이 구분은 GPU에서 계산을 나누어 실행하는 출발점입니다. 독립적인 출력들을 따로 계산할 수도 있고, 앞서 합계의 트리에서 본 것처럼 결과 하나를 만드는 과정 안에서 계산을 나눌 수도 있습니다. 다만 병렬로 계산할 수 있다는 사실만으로 모든 계산이 실제로 같은 순간에 실행되거나 실행 시간이 정해지는 것은 아닙니다. 이를 수행할 연산 장치와 필요한 데이터가 함께 준비되어야 합니다.

다음 글에서는 이렇게 나눌 수 있는 계산을 GPU에 어떻게 맡기는지 살펴보겠습니다. 작업을 구성하는 실행 단위와, GPU가 그 작업을 SM에 배정해 실행하는 흐름을 연결해 보겠습니다.