모델의 가중치 행렬을 여러 GPU에 분할해 각 GPU가 행렬 연산의 일부를 담당하도록 하는 분산 학습 기법. 단일 GPU에 올리기 너무 큰 모델을 처리할 수 있게 한다.

개념 설명

70B 모델은 FP16 기준 140GB로, 80GB A100 두 개에도 안 들어간다. 텐서 병렬화는 행렬 자체를 쪼개어 여러 GPU에 분산하는 방법이다. 예를 들어 큰 행렬 W를 열 방향으로 4등분해 GPU 4개에 하나씩 나눠주고, 각 GPU는 자신의 열에 해당하는 연산만 수행한 뒤 결과를 합친다. NVIDIA의 Megatron-LM이 이 방식을 체계화했다.

텐서 병렬화 외에 레이어를 GPU에 순서대로 배치하는 파이프라인 병렬화, 같은 모델 복사본에 다른 데이터를 넣는 데이터 병렬화를 조합해 대규모 학습을 구성한다.

사용 예시

vLLM 서빙 시 --tensor-parallel-size 4로 4개 GPU에 모델을 분산할 수 있다. 학습에서는 Megatron-LM, DeepSpeed의 ZeRO 전략이 텐서 병렬화를 지원한다. 8개 GPU 서버에서 70B 모델을 텐서 병렬화로 서빙하면 단일 GPU 대비 지연이 크게 줄어든다.