Tensor Parallelism

Tensor Parallelism

하나의 모델을 여러 GPU에 나누어 추론하는 방식은 어떻게 작동하는가.

Scope

핵심 질문

Tensor Parallelism은 하나의 Transformer 레이어 연산을 여러 GPU가 어떻게 나눠 계산하며, 어떤 통신비용을 발생시키는가?

포함 범위

• matrix partition
• column/row parallel
• collective communication
• multi-GPU inference
• interconnect 병목

제외 범위

• 특정 프레임워크 설정법
• 특정 GPU 조합의 최신 벤치마크

Research

핵심 근거

• Tensor Parallelism은 하나의 큰 tensor operation을 여러 장치가 동시에 나눠 수행한다.
• Megatron-LM은 Transformer의 MLP와 attention projection을 column/row partition하는 대표 방식을 제시했다.
• 각 레이어의 부분결과를 결합하기 위해 all-reduce/all-gather 같은 collective communication이 필요하다.
• GPU 수를 늘려도 통신비용 때문에 속도가 선형으로 증가하지 않는다.

참고 자료

• Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism: https://arxiv.org/abs/1909.08053
• NVIDIA NCCL documentation: https://docs.nvidia.com/deeplearning/nccl/

Draft

Tensor Parallelism은 하나의 모델 레이어 안에서 수행되는 큰 tensor 연산을 여러 GPU가 동시에 나눠 계산하는 모델 병렬화 방식이다.

모델을 레이어 단위로 서로 다른 GPU에 배치하는 pipeline형 분할과 달리, 같은 레이어를 여러 GPU가 함께 처리한다.

행렬을 나눈다

Transformer의 대부분의 연산은 큰 matrix multiplication이다.

예를 들어 입력 X와 weight matrix W를 곱할 때 W의 열을 GPU별로 나누면 각 GPU가 출력의 일부를 계산할 수 있다.

반대로 행을 나누어 각 GPU가 부분합을 계산한 뒤 결과를 합칠 수도 있다.

Megatron-LM은 이런 column parallel과 row parallel을 조합해 Transformer 레이어를 여러 GPU로 분산하는 대표적인 방식을 제시했다.

왜 빠를 수 있는가

한 GPU가 수행해야 할 matrix multiplication 양을 여러 GPU가 나누기 때문에 각 GPU의 계산량은 줄어든다.

또한 모델 가중치도 GPU에 분산되므로 한 장의 VRAM에 들어가지 않는 모델을 실행할 수 있다.

통신비용

문제는 각 GPU의 부분결과를 다시 합쳐야 한다는 것이다.

레이어마다 all-reduce, all-gather 같은 collective communication이 필요할 수 있다.

GPU가 계산을 빨리 끝내도 다른 GPU의 결과나 통신이 끝날 때까지 기다려야 하므로 interconnect가 느리면 성능이 크게 떨어질 수 있다.

왜 GPU를 두 배 늘려도 두 배 빨라지지 않는가

병렬화에는 항상 overhead가 있다.

GPU 수가 늘수록 각 GPU의 계산량은 줄지만 통신하는 상대와 synchronization이 증가한다.

작은 모델이나 batch에서는 계산시간보다 통신시간이 더 커질 수도 있다.

그래서 Tensor Parallelism의 효율은 모델 크기, batch, GPU 수, PCIe/NVLink 같은 interconnect에 따라 달라진다.

Pipeline Parallelism과 차이

Pipeline Parallelism은 레이어 묶음을 GPU별로 배치한다.

Tensor Parallelism은 같은 레이어를 동시에 계산하므로 한 토큰의 latency를 줄일 가능성이 더 크지만 통신빈도가 높다.

실제 대규모 시스템에서는 Tensor, Pipeline, Data Parallelism을 조합하기도 한다.

추론에서의 역할

대규모 LLM serving에서는 Tensor Parallelism을 이용해 모델을 여러 GPU에 올리고 하나의 요청을 병렬 처리할 수 있다.

특히 한 GPU의 VRAM에 모델이 들어가지 않거나 높은 처리량이 필요한 경우 유용하다.

다만 단일 사용자 소규모 환경에서는 GPU 간 통신비용 때문에 단순한 layer split이나 단일 GPU 실행이 더 효율적일 수 있다.

핵심

Tensor Parallelism은 하나의 연산을 여러 GPU가 함께 계산하는 방식이다. 계산과 메모리를 분산할 수 있지만 레이어마다 통신이 필요하기 때문에 성능은 GPU 개수보다 interconnect와 workload 특성에 크게 좌우된다.