[vLLM] Tensor Parallelism: 거대 모델을 여러 GPU에 나누는 텐서 병렬화

2026년 4월 7일수정: 2026년 4월 7일

들어가며

70B, 405B 급 모델은 단일 GPU 메모리에 올릴 수 없다. **Tensor Parallelism(TP)**은 모델의 각 레이어 내부에서 가중치 행렬을 여러 GPU에 분할하여 이 문제를 해결한다. Megatron-LM에서 제안된 이 기법은 통신 오버헤드를 최소화하면서 Transformer의 Linear 레이어를 효율적으로 분산한다.

논문: Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (arxiv 1909.08053)
공식 문서: https://docs.vllm.ai

핵심 구조/코드 분석

두 가지 분할 전략

vLLM의 텐서 병렬화는 vllm/model_executor/layers/linear.py에 구현되어 있다. 핵심은 두 가지 Linear 레이어이다:

ColumnParallelLinear: 가중치의 열(output dimension)을 분할
RowParallelLinear: 가중치의 행(input dimension)을 분할

이 두 가지를 교대로 사용하면 MLP와 Attention 레이어를 단 2번의 all-reduce(또는 all-gather)로 병렬화할 수 있다.

ColumnParallelLinear: Y = X[A₁, A₂, ..., Aₚ]

class ColumnParallelLinear(LinearBase):
    """Linear layer with column parallelism.
    The linear layer is defined as Y = XA + b.
    A is parallelized along its second dimension as A = [A_1, ..., A_p].
    """

    def __init__(self, input_size, output_size, bias=True,
                 gather_output=False, ...):
        self.tp_rank = get_tensor_model_parallel_rank()
        self.tp_size = get_tensor_model_parallel_world_size()
        # 출력 차원을 TP 수로 나눔
        self.output_size_per_partition = divide(output_size, self.tp_size)
        self.output_partition_sizes = [self.output_size_per_partition]
        self.gather_output = gather_output

각 GPU는 전체 가중치의 1/tp_size만큼만 보유한다. 입력 X는 모든 GPU에 복제되고, 각 GPU는 자신의 파티션 A_i로 Y_i = XA_i를 계산한다.

RowParallelLinear: Y = [X₁, X₂, ..., Xₚ]A

class RowParallelLinear(LinearBase):
    """Linear layer with row parallelism.
    The linear layer is defined as Y = XA + b.
    A is parallelized along its first dimension and X along
    its second dimension.
    """

    def __init__(self, input_size, output_size, bias=True,
                 input_is_parallel=True, reduce_results=True, ...):
        self.tp_rank = get_tensor_model_parallel_rank()
        self.tp_size = get_tensor_model_parallel_world_size()
        # 입력 차원을 TP 수로 나눔
        self.input_size_per_partition = divide(input_size, self.tp_size)
        self.reduce_results = reduce_results

각 GPU는 입력의 일부(X_i)와 가중치의 일부(A_i)를 가지고, Y_i = X_iA_i를 계산한다. 최종 결과는 all-reduce로 합산한다.

RowParallelLinear의 forward: all-reduce

# RowParallelLinear.forward
output_parallel = self.quant_method.apply(self, input_parallel, bias_)
if self.reduce_results and self.tp_size > 1:
    output = tensor_model_parallel_all_reduce(output_parallel)
else:
    output = output_parallel

bias는 rank 0에서만 더한다:

bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias

TP > 1일 때 모든 rank에서 bias를 더하면 결과가 tp_size배가 되므로, rank 0에서만 더하고 all-reduce 후 모든 rank가 동일한 결과를 갖게 된다.

QKVParallelLinear: Attention의 Q, K, V를 한 번에

class QKVParallelLinear(ColumnParallelLinear):
    ...

Attention 레이어의 Q, K, V 프로젝션은 QKVParallelLinear로 하나의 행렬곱으로 처리된다. ColumnParallelLinear를 상속하므로 output dimension(Q+K+V)이 GPU별로 분할된다. 각 GPU가 전체 Q, K, V의 일부 헤드를 담당하게 된다.

MergedColumnParallelLinear: MLP의 gate/up 통합

class MergedColumnParallelLinear(ColumnParallelLinear):
    ...

Transformer MLP의 gate_proj와 up_proj를 하나의 행렬곱으로 합쳐서 처리한다. 커널 호출 횟수를 줄여 GPU 활용률을 높인다.

Transformer 블록의 통신 패턴

Megatron-LM 스타일 TP의 통신 패턴은 다음과 같다:

Attention:
  QKV (ColumnParallel) → Attention 연산 → O_proj (RowParallel) → all-reduce

MLP:
  gate+up (ColumnParallel) → 활성화 → down (RowParallel) → all-reduce

한 Transformer 블록당 all-reduce가 정확히 2번 발생한다. all-reduce는 GPU 간 통신이므로 최소화하는 것이 핵심이며, ColumnParallel → RowParallel 교대 배치가 이를 달성한다.

양자화와의 통합

class LinearBase(PluggableLayer):
    def __init__(self, input_size, output_size, ...):
        if quant_config is None:
            self.quant_method = UnquantizedLinearMethod()
        else:
            self.quant_method = quant_config.get_quant_method(
                self, prefix=prefix
            )

quant_method를 플러그인으로 주입하여, FP8, GPTQ, AWQ 등 다양한 양자화 방법이 TP와 자연스럽게 결합된다. 양자화된 가중치도 동일한 TP 분할 로직을 따른다.

왜 이 설계인가

최소 통신: ColumnParallel과 RowParallel을 교대로 배치하여 Transformer 블록당 all-reduce를 2번으로 제한한다. Pipeline Parallelism과 달리 레이턴시 증가가 최소이다.
가중치 메모리 분산: 각 GPU가 1/tp_size의 가중치만 보유하므로, TP=8이면 8배 큰 모델을 올릴 수 있다. KV 캐시도 헤드 수가 분할되어 함께 줄어든다.
양자화 호환: quant_method 플러그인 패턴으로 TP와 양자화가 독립적으로 조합된다. FP8 + TP=4 같은 조합이 코드 변경 없이 동작한다.
NCCL 최적화: vLLM은 all-reduce에 NCCL을 사용하여 NVLink 같은 고속 인터커넥트를 최대한 활용한다. NVLink가 있으면 all-reduce 오버헤드가 매우 작아 TP 확장 효율이 높다.

Tensor Parallelism은 LLM 서빙에서 가장 기본적이면서도 필수적인 분산 기법이다. vLLM의 구현은 Megatron-LM의 검증된 패턴을 충실히 따르면서, 양자화와 다양한 어텐션 메커니즘과의 조합을 유연하게 지원한다.

논문 핵심 내용

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (1909.08053) 논문은 레이어 내부(intra-layer) 모델 병렬화를 통해 수십억 파라미터 트랜스포머를 효율적으로 학습하는 방법을 제안했다.

핵심 아이디어: PyTorch에 최소한의 코드 변경만으로 구현 가능한 텐서 병렬화 전략이다. 파이프라인 병렬과 독립적으로 동작하며, Transformer의 MLP와 Self-Attention을 ColumnParallel/RowParallel로 교대 분할하여 블록당 all-reduce를 정확히 2회로 제한한다.

스케일링 효율성

메트릭	수치
단일 GPU 성능	39 TFLOPS (피크의 30%)
전체 애플리케이션 성능	15.1 PetaFLOPS
스케일링 효율	76% (512 GPU 기준)
최대 모델 크기	83억(8.3B) 파라미터
사용 GPU 수	512

벤치마크 결과

모델	벤치마크	결과	기존 SOTA
GPT-2 (8.3B)	WikiText103 Perplexity	10.8	15.8
GPT-2 (8.3B)	LAMBADA Accuracy	66.5%	63.2%
BERT (3.9B)	RACE Accuracy	90.9%	89.4%

512개 GPU에서 76% 스케일링 효율을 달성한 것이 핵심이다. 이론적 최대 대비 76%라는 숫자는 all-reduce 통신 오버헤드를 고려하면 상당히 높은 수치다. GPT-2 8.3B 모델은 WikiText103에서 perplexity 10.8로 기존 15.8을 크게 앞섰고, BERT 3.9B는 RACE에서 90.9%를 기록했다. 이 결과는 모델 크기 확대가 성능 향상으로 직결된다는 것을 보여줬고, 이후 GPT-3(175B), LLaMA(65B/70B) 등 대규모 모델 학습의 기반이 됐다.

vLLM 의 다른글

이전글 [vLLM] Sampler: logits에서 토큰까지, 샘플링 파이프라인 전체 분석
현재글 : [vLLM] Tensor Parallelism: 거대 모델을 여러 GPU에 나누는 텐서 병렬화
다음글 [vLLM] GPTQ: 2차 정보 기반 후훈련 양자화