[워크삽] AWS 트레니엄 워크삽(3)-텐서병렬화

이번 블로그는 vLLM과 AWS Trainium을 활용한 LLM 추론 확장하는 데 있어서 텐서 병렬화(Tensor Parallelism)에 대한 개념을 알아보도록 하자.

1. 메모리 문제의 해결책: 텐서 병렬화

Trainium이 상당한 메모리 용량을 제공하더라도 Llama 3.1 8B와 같은 모델은 단일 가속기의 메모리 용량을 초과할 수 있다. 이 문제는 텐서 병렬화(Tensor Parallelism)를 사용해 해결할 수 있다. 텐서 병렬화는 모델의 가중치를 여러 NeuronCore에 나누어 배치하는 기술이다.

텐서 병렬화는 다음과 같은 방식으로 동작한다.

• 가중치 텐서 분할(Splitting Weight Tensors): 모델의 가중치 텐서를 여러 가속기에 나누어 저장한다.
• 병렬 계산(Parallel Computation): 각 가속기는 자신에게 할당된 가중치와 데이터 부분을 동시에 처리한다.
• 출력 결합(Output Combination): 각 가속기에서 계산한 결과를 하나로 결합해 최종 출력을 생성한다.

이 방식을 사용하면 단일 가속기에는 올릴 수 없는 대규모 모델도 여러 가속기의 메모리와 연산 자원을 함께 활용해 효율적으로 추론할 수 있다.

2. 핵심 정리

텐서 병렬화의 목적은 하나의 모델 연산을 여러 가속기에 분산하는 것이다. 모델 복제본을 가속기마다 하나씩 배치하는 데이터 병렬화(Data Parallelism)와 달리, 텐서 병렬화는 하나의 모델을 구성하는 가중치 텐서 자체를 분할한다.

각 NeuronCore는 전체 모델 중 자신에게 배정된 부분을 계산하며, 필요한 지점에서 코어 간 통신을 수행해 중간 결과나 최종 출력을 결합한다. 따라서 단일 가속기의 메모리 한계를 넘어설 수 있지만, 코어 사이의 통신 비용과 동기화 부담도 함께 고려해야 한다.

3. 참고 자료

• 원문: AWS Workshop Studio – Tensor Parallelism