어떠한 표준 병렬화 (Parallelism) 방식보다 낮은 GPU당 피크 메모리 (Peak Memory)를 달성하며, 대규모 모델을 GPU에 분할하는 새로운 방식인 folded Tensor and Sequence Parallelism (TSP)를 소개합니다.
@AMD MI300x에서 확장되었습니다.
더 큰 모델, 더 긴 컨텍스트 (Context), 그리고 더 높은 처리량 (Throughput)
현대의 LLM (Large Language Models)은 단일 GPU에 담기에는 너무 크기 때문에 분할됩니다. 두 가지 대중적인 접근 방식은 다음과 같습니다:
- Tensor Parallelism (TP): 모델 가중치 (Weights)를 분할
- Sequence Parallelism (SP): 시퀀스 (Sequence)를 따라 활성화 값 (Activations)을 분할
모델과 컨텍스트가 커짐에 따라 두 방식 모두 필수적일 수 있습니다.
TP와 SP를 결합하는 전통적인 방식은 독립적인 축 (Axes)을 사용하는 것입니다. 만약 가중치를 8개로 분할하고 시퀀스를 2개로 분할하고 싶다면, 단일 모델 복제본 (Replica)을 위해 8 × 2 = 16개의 GPU가 필요합니다.
이는 데이터 병렬화 (Data Parallelism)와 같은 다른 병렬화 축으로 확장할 수 있는 여유 GPU를 줄어들게 만듭니다.
일반적인 노드 (Node)는 고대역폭 노드 내 링크 (Intra-node links, 예: NVLink, AMD Infinity Fabric)로 연결된 8개의 GPU를 보유합니다.
복제본이 하나의 노드를 초과하면, 통신은 더 느린 노드 간 상호 연결 (Inter-node interconnects, 예: Ethernet/InfiniBand)을 통해 이루어지며, 이는 처리량/지연 시간 (Throughput/Latency)을 저하시키고 오버랩 (Overlap)을 어렵게 만듭니다.
TSP의 핵심 통찰은 TP와 SP를 동일한 장치 축 (Device axis)으로 접는 (Fold) 것입니다.
각 GPU가 가중치의 슬라이스 (Slice)와 시퀀스의 슬라이스를 동시에 보유함으로써, TSP는 16개가 아닌 단 8개의 GPU만으로 두 방식의 메모리 이점을 모두 제공합니다. 전체 모델 복제본이 하나의 노드에 들어갑니다.
TSP는 두 가지 설계로 이루어져 있으며, 각 레이어의 어텐션 블록 (Attention block)과 MLP 블록 (MLP block)을 위한 설계가 각각 존재합니다:
- Attention: GPU들이 자신의 가중치 슬라이스를 번갈아 가며 브로드캐스팅 (Broadcasting)한 다음, 어텐션이 전체 컨텍스트를 볼 수 있도록 시퀀스 샤드 (Sequence shards)를 공유합니다.
- MLP: 가중치 샤드 (Weight shards)가 링 (Ring) 구조로 순환하며 누적됩니다.
MI300X에서 TSP는 테스트된 어떤 방식보다 가장 낮은 메모리 사용량을 기록했습니다.
8K 컨텍스트에서는 가중치 중심 (Weight-dominated)이므로 TSP ≈ TP입니다.
64K 컨텍스트에서는 가중치와 활성화 값 모두가 중심 (Weight- and activation-dominated)이 되므로, TSP가 두 방식의 장점을 모두 제공합니다.
TSP는 TP의 가중치 샤딩 (Weight-sharding)과 SP의 활성화 값 샤딩 (Activation-sharding)을 모두 상속받습니다.
TSP는 더 높은 처리량 (Throughput)을 제공합니다. 1024개의 MI300X GPU 환경에서 - 128K 토큰의 컨텍스트 - 모델 복사본당 8개의 GPU를 사용할 때, TSP는 173M tok/sec를 달성하며, 이는 동일한 조건의 TP+SP 방식인 86M와 대비됩니다.
TSP는 EP 및 PP와 같은 기존 방식들과 병행하여 사용할 수 있습니다.
논문:
https://arxiv.org/abs/2604.26294
블로그:
https://zyphra.com/post/tsp