텐서 병렬과 파이프라인 병렬

모델이 GPU 한 장에 들어가지 않으면 나눠야 한다. 들어가더라도 디코드 지연을 줄이려고 나누기도 한다. 나누는 방법은 크게 둘이다. 층 안의 행렬을 쪼개 여러 GPU가 같은 층을 함께 계산하는 텐서 병렬(tensor parallelism)과, 층들을 순서대로 나눠 GPU마다 일부 층을 맡기는 파이프라인 병렬(pipeline parallelism)이다. 통신 패턴에 따라 맞는 하드웨어와 목적이 나뉜다.

한 장에 들어가지 않는 모델

Llama 3.1 70B의 설정은 층 80개, 은닉 차원 8,192, 어텐션 헤드 64개, KV 헤드 8개, MLP 중간 차원 28,672다. 행렬 크기를 더하면 파라미터는 약 705억 개이고 BF16 가중치는 약 141GB다. H100 80GB 한 장에는 가중치조차 올라가지 않는다.

KV 캐시도 커진다. 층 80개, KV 헤드 8개, 헤드 차원 128이므로 토큰당 크기는 이렇다.

2 × 80 × 8 × 128 × 2바이트 = 327,680바이트 = 320KiB

8B 모델의 2.5배다. 가중치를 겨우 올려도 KV 캐시 자리가 없으면 동시에 처리할 수 있는 요청이 몇 개 안 된다. 그래서 70B급 모델은 보통 GPU 4장이나 8장에 나눠 올린다.

텐서 병렬

Megatron-LM이 정리한 텐서 병렬은 트랜스포머 층 안의 행렬 곱을 GPU들이 나눠 계산한다.

MLP는 행렬 곱 두 개로 이뤄진다. 첫 행렬을 열 방향으로 나누면 GPU마다 출력의 일부 열을 독립적으로 계산하고, 활성화 함수도 각자 적용할 수 있다. 둘째 행렬을 행 방향으로 나누면 GPU마다 부분 합을 내고 이 부분 합을 모든 GPU가 더하는 올리듀스(all-reduce) 한 번으로 출력이 완성된다.

어텐션은 헤드 단위로 나눈다. 쿼리·키·값 투영을 헤드 단위로 쪼개 GPU마다 일부 헤드를 맡고, 출력 투영을 행 방향으로 나눠 다시 올리듀스 한 번으로 합친다.

그래서 순전파 기준으로 층마다 올리듀스가 두 번 일어난다. 70B 모델이면 디코드 스텝 하나에 160번이다.

텐서 병렬의 이득은 GPU마다 가중치와 KV 캐시의 일부만 갖는다는 데 있다. TP=4로 나누면 GPU마다 가중치는 약 35GB, KV 헤드는 2개씩만 맡는다. 디코드 스텝에서 GPU 네 장이 각자 자기 몫을 동시에 읽으므로, 쓸 수 있는 대역폭이 네 배가 된다.

가중치 읽기 하한 ≈ 141GB ÷ (4 × 3.35TB/s) ≈ 10.5ms

텐서 병렬은 모델을 올리는 수단이자 디코드 지연을 줄이는 수단이다. 8B처럼 한 장에 들어가는 모델도 TP=2로 나누면 스텝 하한이 절반 가까이 준다.

통신 비용

텐서 병렬의 대가는 올리듀스다. 디코드에서 올리듀스 한 번이 주고받는 데이터는 배치 크기 × 은닉 차원 × 2바이트로, 배치 32면 약 0.5MB다. 크지 않지만 스텝마다 160번 일어나고 각 올리듀스가 끝나야 다음 계산을 시작할 수 있다. 메시지가 작으면 전송 시간보다 통신을 시작하고 동기화하는 고정 지연이 시간을 정한다.

그래서 텐서 병렬은 GPU 사이 연결이 빠를 때만 이득이다. H100 SXM은 NVLink로 초당 900GB를 주고받는다. 같은 서버 안에서 NVLink로 연결된 GPU끼리는 텐서 병렬을 쓰고, 서버를 넘는 네트워크(InfiniBand, 이더넷)에 텐서 병렬을 걸지 않는 이유가 이것이다. 서버 하나에 GPU가 8장이면 텐서 병렬의 크기도 대개 8이 상한이다.

TP를 키울수록 GPU마다 맡는 계산은 줄지만 올리듀스 횟수는 그대로이고 참여하는 GPU가 늘어 통신 비중이 커진다. 그래서 TP=2에서 4로, 4에서 8로 갈 때 지연이 줄어드는 폭은 점점 작아진다.

파이프라인 병렬

파이프라인 병렬은 층을 순서대로 나눠 GPU마다 연속된 층 묶음을 맡긴다. 80층을 4장에 나누면 GPU마다 20층씩이다. 앞 GPU가 자기 층을 다 계산하면 그 활성화를 다음 GPU에 넘긴다. 단계 사이에 오가는 것은 활성화 하나뿐이라 통신이 적고 느린 서버 간 네트워크로도 나눌 수 있다.

대신 토큰 하나의 계산은 GPU들을 차례로 거쳐야 하므로 요청 하나만 보면 한 번에 GPU 한 장만 일한다. 디코드 지연은 줄지 않는다. 나머지 GPU를 놀리지 않으려면 여러 배치를 시차를 두고 흘려보내는 마이크로배치가 필요하다. GPipe가 학습에서 보인 이 방식은 서빙에서도 같다. 배치 1이 2번 GPU로 넘어가면 1번 GPU는 배치 2를 시작한다. 파이프라인이 다 차면 모든 GPU가 일하므로 처리량은 오르지만 단계마다 넘겨주는 시간과 단계별 부하 불균형만큼 손실이 생긴다.

	텐서 병렬	파이프라인 병렬
나누는 단위	층 안의 행렬	층 묶음
통신	층마다 올리듀스 2회 (작고 잦다)	단계 경계마다 활성화 전달 (드물다)
필요한 연결	NVLink 같은 고속 연결	서버 간 네트워크로도 가능
디코드 지연	줄어든다	줄지 않는다
주로 쓰는 곳	한 서버 안	한 서버에 담기지 않는 모델의 서버 간 분할

한 서버에 담기지 않는 모델은 서버 안에서 텐서 병렬, 서버 사이에서 파이프라인 병렬을 함께 쓴다.

복제본을 늘리는 방식과의 비교

모델이 한 장에 들어간다면 GPU를 여러 장 쓰는 또 다른 방법은 같은 모델을 장마다 따로 올리는 것이다. 복제본 사이에는 통신이 없고 로드 밸런서가 요청을 나눠 준다.

8B 모델과 H100 4장이 있을 때를 비교하면 이렇다.

구성	디코드 스텝 하한	KV 캐시 공간	특징
복제본 4개 (TP=1 × 4)	약 4.8ms	장마다 약 55GB, 따로 쓴다	통신 없음. 처리량이 가장 높다. 한 장이 죽어도 나머지가 돈다
TP=4 하나	약 1.2ms + 통신	약 265GB를 한 묶음이 쓴다	지연이 가장 짧다. 긴 문맥 요청을 더 많이 담는다
TP=2 × 2	약 2.4ms + 통신	묶음마다 약 125GB	중간

처리량이 목표라면 복제본을 늘리는 편이 대개 낫다. 통신이 없고 서버마다 배치가 따로 차기 때문이다. TPOT 목표가 빡빡해 한 장의 스텝 시간으로는 지킬 수 없거나, 문맥이 아주 긴 요청을 받아야 해서 KV 캐시를 한곳에 크게 모아야 할 때 텐서 병렬을 쓴다. 이 경우에도 고르는 기준은 같은 부하로 잰 장비 수당 굿풋이다.