Folded Tensor and Sequence Parallelism (TSP) 소개: 대규모 모델을 GPU에 분할하는 새로운 방식

어떠한 표준 병렬화 (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