텐서 병렬화

Tensor parallelism, TP; 텐서 병렬화
신경망 한 층 안의 가중치 행렬을 여러 GPU에 쪼개 나눠 계산하는 모델 병렬화 방식. 층 내부(intra-layer) 모델 병렬화라고도 한다.

텐서 병렬화(tensor parallelism, TP)는 행렬 곱 하나를 여러 장치가 나눠 계산하도록 가중치 행렬을 행이나 열 방향으로 쪼개는 분산 학습 방식이다. 층 단위로 모델을 나누는 파이프라인 병렬화(GPipe, 파이프드림)가 층과 층 사이를 자른다면, 텐서 병렬화는 층 안을 자른다. 장치마다 가중치의 일부만 들고 있으므로 GPU 한 장에 들어가지 않는 층도 학습·추론할 수 있고, 활성값 일부도 나눠 저장된다.

트랜스포머용 텐서 병렬화의 기본형은 2019년 엔비디아의 Megatron-LM 논문이 정리한 방식이다. 트랜스포머의 MLP와 셀프 어텐션 블록을 각각 "열 분할 다음 행 분할" 구조로 나누면 블록마다 순전파에 올 리듀스 한 번, 역전파에 한 번만 통신하면 된다는 것이 핵심이다. 이 방법은 새 컴파일러나 라이브러리 없이 파이토치에 통신 연산 몇 개만 넣어 구현할 수 있었다.^[1] 그 이전에 구글의 Mesh-TensorFlow(2018)가 텐서 차원을 장치 격자에 나누는 일반적인 방식을 제안해 T5 학습에 쓰였다.^[2] 이 문서의 "메가트론 방식"은 Megatron 문서의 프레임워크에서 온 이름이다.

행렬 곱 나누기

입력 $X \in ℝ^{b \times h}$ , 가중치 $A$ 로 $Y = X A$ 를 계산할 때 $t$ 개 장치에 나누는 방법은 두 가지다.

방식	가중치 분할	각 장치의 계산	결과를 얻는 통신
열 병렬(column parallel)	$A = [A_{1}, A_{2}, \dots, A_{t}]$	$Y_{i} = X A_{i}$ (입력 $X$ 는 전체 필요)	출력이 열 방향으로 나뉨. 전체가 필요하면 올-개더
행 병렬(row parallel)	$A = [\begin{matrix} A_{1} \\ ⋮ \\ A_{t} \end{matrix}]$	$Y_{i} = X_{i} A_{i}$ (입력도 열 방향으로 나뉨)	$Y = \sum_{i} Y_{i}$ , 올 리듀스

핵심은 두 방식을 이어 붙이는 것이다. 첫 행렬을 열 병렬로 나누면 출력이 열 방향으로 나뉜 채 나오고, 이것이 바로 둘째 행렬의 행 병렬 입력 모양이다. 그래서 중간에 통신 없이 두 번째 곱까지 하고, 마지막에 한 번만 합치면 된다.

트랜스포머 블록 나누기

MLP

트랜스포머의 MLP는 $Z = GeLU (X A) B$ 다. $A$ 를 행 방향으로 나누면 $GeLU (X_{1} A_{1} + X_{2} A_{2}) \neq GeLU (X_{1} A_{1}) + GeLU (X_{2} A_{2})$ 라서 GeLU 전에 동기화가 필요하다. 반면 $A$ 를 열 방향으로 나누면 비선형 함수를 장치마다 따로 적용할 수 있다. $[Y_{1}, \dots, Y_{t}] = [GeLU (X A_{1}), \dots, GeLU (X A_{t})], Z = \sum_{i = 1}^{t} Y_{i} B_{i}$ 따라서 $A$ 는 열 병렬, $B$ 는 행 병렬로 나누고 마지막 합만 올 리듀스한다.^[1] 아래는 이 등식을 한 프로세스 안에서 확인하는 최소 예시다.

import torch
import torch.nn.functional as F
torch.manual_seed(0)

b, h, t = 4, 16, 4                    # 배치, 은닉 차원, 텐서 병렬 크기
X = torch.randn(b, h)
A = torch.randn(h, 4 * h)             # 첫 선형층 (h -> 4h)
B = torch.randn(4 * h, h)             # 둘째 선형층 (4h -> h)
ref = F.gelu(X @ A) @ B               # 병렬화하지 않은 결과

A_shards = A.chunk(t, dim=1)          # 열 분할: 장치 i는 A_i만 가진다
B_shards = B.chunk(t, dim=0)          # 행 분할: 장치 i는 B_i만 가진다
partials = [F.gelu(X @ A_i) @ B_i     # 장치마다 통신 없이 계산
            for A_i, B_i in zip(A_shards, B_shards)]
out = sum(partials)                   # 실제로는 all-reduce 한 번
print(torch.allclose(out, ref, atol=1e-4))   # True

SwiGLU를 쓰는 라마 계열 MLP $W_{2} (SiLU (X W_{1}) ⊙ X W_{3})$ 도 $W_{1}, W_{3}$ 를 열 병렬, $W_{2}$ 를 행 병렬로 나누면 똑같이 올 리듀스 한 번으로 끝난다.

셀프 어텐션

멀티 헤드 어텐션은 헤드끼리 독립이므로 헤드를 장치에 나누면 된다. $Q, K, V$ 투영 가중치를 열 병렬로 나눠 장치마다 $a / t$ 개 헤드를 맡고, 그 안에서 어텐션을 통신 없이 계산한 뒤, 출력 투영을 행 병렬로 해 올 리듀스한다. 헤드 수가 $t$ 로 나누어떨어져야 한다. 그룹 쿼리 어텐션이면 KV 헤드 수도 함께 고려해야 하며, KV 헤드가 $t$ 보다 적으면 KV 헤드를 여러 장치에 복제하는 식으로 처리한다.

통신 연산 f와 g

Megatron-LM은 블록 앞뒤에 서로 켤레(conjugate)인 두 연산을 둔다.

$f$ : 순전파에서는 항등, 역전파에서는 올 리듀스. 블록 입력 쪽에 둔다(모든 장치가 같은 $X$ 를 쓰므로 $X$ 의 기울기는 장치별 부분 기울기의 합).
$g$ : 순전파에서는 올 리듀스, 역전파에서는 항등. 블록 출력 쪽에 둔다.

import torch.distributed as dist

class CopyToTP(torch.autograd.Function):      # f
    @staticmethod
    def forward(ctx, x):
        return x
    @staticmethod
    def backward(ctx, grad):
        dist.all_reduce(grad, group=tp_group)
        return grad

class ReduceFromTP(torch.autograd.Function):  # g
    @staticmethod
    def forward(ctx, x):
        dist.all_reduce(x, group=tp_group)
        return x
    @staticmethod
    def backward(ctx, grad):
        return grad

결국 트랜스포머 층 하나(어텐션 + MLP)에서 순전파 두 번, 역전파 두 번 올 리듀스가 일어난다.^[1] 레이어 정규화, 드롭아웃, 잔차 연결은 장치마다 똑같이 복제해 계산한다.

임베딩과 출력층

어휘 수 $v$ 가 수만~수십만이라 임베딩 행렬 $E \in ℝ^{h \times v}$ 도 크다. 어휘 방향으로 나누고(어휘 병렬), 출력층 로짓을 모두 모으는 대신( $b \times s \times v$ 통신) 장치마다 부분 로짓으로 교차 엔트로피를 계산하는 데 필요한 최댓값과 합만 주고받으면 통신이 $b \times s$ 로 준다.^[1] PyTorch에서는 이를 loss_parallel이라 부른다.

통신 비용과 배치

텐서 병렬화의 올 리듀스는 층마다, 계산 경로 위에서 일어난다. 크기는 $b \times s \times h$ 로 활성값 전체다. 계산과 겹치기 어렵기 때문에 GPU 사이 대역폭이 매우 높아야 한다. 엔비디아 연구진은 $g$ 개 GPU가 든 서버에서는 텐서 병렬을 $g$ 까지만 쓰고, 더 큰 모델은 서버 사이에 파이프라인 병렬을 쓰라고 정리했다.^[3] 즉 텐서 병렬 그룹은 보통 NVLink로 묶인 한 서버 안에 둔다.

Megatron-LM 원 논문은 83억 매개변수 모델을 GPU 512장으로 학습해 15.1PFLOPS를 냈고, 이는 GPU 한 장(39TFLOPS, 최대 성능의 30%) 기준 76% 확장 효율이었다.^[1] 2021년 후속 연구는 텐서·파이프라인·데이터 병렬을 결합(PTD-P)해 1조 매개변수 모델을 GPU 3072장에서 502PFLOPS(GPU당 최대 성능의 52%)로 학습했다.^[3]

시퀀스 병렬화

기본 텐서 병렬에서 레이어 정규화와 드롭아웃 구간의 활성값은 모든 장치에 복제된다. 2022년 엔비디아 연구진은 이 구간을 시퀀스 방향으로 나누는 시퀀스 병렬화(sequence parallelism)를 제안했다.^[4]

레이어 정규화·드롭아웃은 토큰마다 독립이므로 시퀀스를 $t$ 조각으로 나눠 각 장치가 자기 조각만 처리한다.
텐서 병렬 구간에 들어갈 때 올-개더로 시퀀스를 모으고, 나올 때 올 리듀스 대신 리듀스-스캐터로 합치면서 다시 시퀀스 방향으로 나눈다.
올 리듀스 = 리듀스-스캐터 + 올-개더이므로 통신량은 기본 텐서 병렬과 같다.

층당 활성값 메모리(바이트, 16비트 기준)는 $s b h (10 + 24 / t + 5 a s / (h t))$ 에서 $s b h (34 / t + 5 a s / (h t))$ 로 줄어, 모든 활성값이 $t$ 로 나뉜다. 선택적 재계산과 함께 써서 활성값 메모리를 5배 줄였다.^[4] 자세한 메모리 식은 그래디언트 체크포인팅 문서 참고.

같은 "시퀀스 병렬"이라는 이름으로, 어텐션까지 시퀀스 방향으로 나누는 긴 문맥용 기법들도 있다. 링 어텐션(Ring Attention)은 KV 블록을 장치 사이에 고리 모양으로 돌리며 블록 단위 어텐션을 계산하고,^[5] DeepSpeed-Ulysses는 올-투-올 통신으로 시퀀스 분할을 헤드 분할로 바꿔 어텐션을 계산한다.^[6] Megatron-LM은 이런 방식을 컨텍스트 병렬(context parallelism)이라 따로 부르며, 텐서 병렬과 함께 쓰는 위의 시퀀스 병렬과 구분한다.

다른 병렬화와 비교

방식	나누는 축	GPU당 가중치	GPU당 활성값	통신 내용과 시점	약점
데이터 병렬(DDP)	배치	전체	배치 몫	역전파 끝에 기울기 올 리듀스	모델이 한 장에 들어가야 함
ZeRO·FSDP	배치 + 모델 상태	$1 / N$ (3단계)	배치 몫	층마다 가중치 올-개더, 기울기 리듀스-스캐터	배치가 작으면 통신 비율 증가
파이프라인 병렬	층 묶음(깊이)	단계 몫	단계 몫 × 동시 마이크로배치	단계 경계의 활성값·기울기 점대점 전송	파이프라인 버블 $(p - 1) / m$ , 부하 균형
텐서 병렬	층 안 행렬(너비)	$1 / t$	일부만 $1 / t$	층마다 활성값 올 리듀스, 계산 경로 위	고대역폭 필요, 노드 안으로 제한
시퀀스 병렬(TP와 함께)	정규화·드롭아웃 구간의 시퀀스	TP와 같음	전부 $1 / t$	TP와 같은 양(리듀스-스캐터 + 올-개더)	TP 없이는 의미 없음
컨텍스트 병렬	전체 시퀀스(어텐션 포함)	전체 또는 다른 방식과 결합	$1 / c$	어텐션 중 KV 교환	긴 문맥에서만 이득
전문가 병렬	MoE 전문가	전문가 몫	라우팅된 토큰	올-투-올	부하 불균형

파이프라인 버블 식의 $p$ 는 파이프라인 단계 수, $m$ 은 배치당 마이크로배치 수다.^[3] 실제 대규모 학습은 노드 안에서 텐서(+시퀀스) 병렬, 노드 사이에서 파이프라인 병렬, 그 바깥에서 데이터 병렬(또는 FSDP)을 쓰는 3차원 병렬화를 기본으로 한다. 이때 GPU 수는 $N = t \times p \times d$ 다.

PyTorch로 쓰기

PyTorch는 DTensor 기반 torch.distributed.tensor.parallel 모듈로 텐서 병렬화를 제공한다. 모듈 이름별로 ColwiseParallel, RowwiseParallel, SequenceParallel을 지정하면 통신은 자동으로 들어간다.^[7] 아래는 PyTorch 튜토리얼의 라마 2 블록 분할 계획이다.^[8]

from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import Shard, Replicate
from torch.distributed.tensor.parallel import (
    parallelize_module, ColwiseParallel, RowwiseParallel,
    SequenceParallel, PrepareModuleInput)

tp_mesh = init_device_mesh("cuda", (8,))        # 한 서버의 GPU 8장

layer_tp_plan = {
    "attention_norm": SequenceParallel(),       # 시퀀스 방향 분할
    "attention": PrepareModuleInput(
        input_layouts=(Shard(1), Replicate()),
        desired_input_layouts=(Replicate(), Replicate())),
    "attention.wq": ColwiseParallel(use_local_output=False),
    "attention.wk": ColwiseParallel(use_local_output=False),
    "attention.wv": ColwiseParallel(use_local_output=False),
    "attention.wo": RowwiseParallel(output_layouts=Shard(1)),
    "ffn_norm": SequenceParallel(),
    "feed_forward": PrepareModuleInput(
        input_layouts=(Shard(1),), desired_input_layouts=(Replicate(),)),
    "feed_forward.w1": ColwiseParallel(),
    "feed_forward.w2": RowwiseParallel(output_layouts=Shard(1)),
    "feed_forward.w3": ColwiseParallel(),
}
for block in model.layers:
    parallelize_module(block, tp_mesh, layer_tp_plan)

Shard(1)은 (배치, 시퀀스, 은닉) 텐서를 1번 차원, 곧 시퀀스 방향으로 나눈 배치를 뜻한다. 어텐션의 Q·K·V를 use_local_output=False로 두는 것은 헤드 수가 장치마다 $a / t$ 로 바뀌는 것을 DTensor가 처리하게 하려는 것이다. 튜토리얼은 서버 안에서는 텐서 병렬, 서버 사이에서는 FSDP를 쓰는 2차원 메시(init_device_mesh("cuda", (8, 8)))도 보여 준다.^[8]

Megatron-LM에서는 --tensor-model-parallel-size, --sequence-parallel, --pipeline-model-parallel-size, --context-parallel-size 같은 인자로 병렬 차원을 정한다. 추론 엔진(vLLM 등)도 같은 열·행 분할로 큰 모델을 여러 GPU에 나눠 서빙한다.

언제 쓰나

층 하나의 가중치나 활성값이 GPU 한 장에 너무 클 때, 또는 ZeRO-3·FSDP만으로는 GPU당 배치가 너무 작아져 효율이 떨어질 때.
추론에서 지연 시간을 줄일 때. 파이프라인 병렬은 처리량은 늘려도 토큰 하나의 지연은 줄이지 못하지만, 텐서 병렬은 한 층의 계산을 여러 GPU가 동시에 하므로 지연을 줄인다.
서버 사이 연결(이더넷·인피니밴드)을 넘어 텐서 병렬을 늘리면 통신 대기로 효율이 급격히 떨어지므로 피한다.
작은 모델에 높은 $t$ 를 쓰면 행렬이 잘게 쪼개져 GPU 활용률이 떨어진다.

같이 보기

각주

↑ ^1.0 ^1.1 ^1.2 ^1.3 ^1.4 Mohammad Shoeybi 외, 〈Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism〉, 2019. arXiv:1909.08053
↑ Noam Shazeer 외, 〈Mesh-TensorFlow: Deep Learning for Supercomputers〉, NeurIPS 2018. arXiv:1811.02084
↑ ^3.0 ^3.1 ^3.2 Deepak Narayanan 외, 〈Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM〉, SC21. arXiv:2104.04473
↑ ^4.0 ^4.1 Vijay Korthikanti 외, 〈Reducing Activation Recomputation in Large Transformer Models〉, 2022. arXiv:2205.05198
↑ Hao Liu, Matei Zaharia, Pieter Abbeel, 〈Ring Attention with Blockwise Transformers for Near-Infinite Context〉, 2023. arXiv:2310.01889
↑ Sam Ade Jacobs 외, 〈DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models〉, 2023. arXiv:2309.14509
↑ PyTorch, 〈Tensor Parallelism - torch.distributed.tensor.parallel〉, 2.14 문서. docs.pytorch.org
↑ ^8.0 ^8.1 PyTorch, 〈Large Scale Transformer model training with Tensor Parallel (TP)〉 튜토리얼. docs.pytorch.org

참고 문헌

[meg19-1] 1.0 ^1.1 ^1.2 ^1.3 ^1.4 Mohammad Shoeybi 외, 〈Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism〉, 2019. arXiv:1909.08053

[mtf-2] Noam Shazeer 외, 〈Mesh-TensorFlow: Deep Learning for Supercomputers〉, NeurIPS 2018. arXiv:1811.02084

[ptdp-3] 3.0 ^3.1 ^3.2 Deepak Narayanan 외, 〈Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM〉, SC21. arXiv:2104.04473

[kor22-4] 4.0 ^4.1 Vijay Korthikanti 외, 〈Reducing Activation Recomputation in Large Transformer Models〉, 2022. arXiv:2205.05198

[ring-5] Hao Liu, Matei Zaharia, Pieter Abbeel, 〈Ring Attention with Blockwise Transformers for Near-Infinite Context〉, 2023. arXiv:2310.01889

[ulysses-6] Sam Ade Jacobs 외, 〈DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models〉, 2023. arXiv:2309.14509

[pttp-7] PyTorch, 〈Tensor Parallelism - torch.distributed.tensor.parallel〉, 2.14 문서. docs.pytorch.org

[tptut-8] 8.0 ^8.1 PyTorch, 〈Large Scale Transformer model training with Tensor Parallel (TP)〉 튜토리얼. docs.pytorch.org

[1]

[2]

[3]

[4]

[5]

[6]

[7]

[8]