Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

"메모리 효율적인 트랜스포머 학습 및 추론을 위한 텐서 및 시퀀스 병렬성 접기 (Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference)" 논문에 대한 상세한 기술 요약은 다음과 같습니다.

1. 문제 제기

대규모 언어 모델 (LLM) 의 학습 및 추론은 두 가지 주요 메모리 병목 현상에 직면해 있습니다:

1. 파라미터 메모리: 모델 가중치, 기울기, 옵티마이저 상태를 저장하는 데 필요한 공간.
2. 활성화 메모리: 중간 활성화 값을 저장하는 데 필요한 공간으로, 어텐션 메커니즘에서 시퀀스 길이 ($S^2$) 에 따라 2 차적으로 증가합니다.

기존의 병렬화 전략들은 이러한 문제들을 개별적으로 해결하거나 트레이드오프를 수반합니다:

• 텐서 병렬성 (TP): 모델 가중치를 디바이스 간에 분할하여 파라미터 메모리를 줄이지만, 활성화 메모리는 완전히 복제된 상태로 남깁니다. 높은 활성화 메모리와 통신 볼륨으로 인해 긴 시퀀스 길이에서는 비효율적이게 됩니다.
• 시퀀스 병렬성 (SP): 입력 시퀀스를 디바이스 간에 분할하여 활성화 메모리를 줄이지만, 모델 가중치는 완전히 복제된 상태로 남깁니다. 파라미터 메모리가 병목이 되는 경우에는 효과가 없습니다.
• 전통적인 TP+SP: 장치를 격자 (예: $T \times \Sigma$ 그리드) 의 직교 축에 할당하여 두 가지를 결합합니다. 효과적이기는 하지만, 모델 병렬성을 위해 많은 수의 장치 ($T \times \Sigma$) 를 소모하여 데이터 병렬성을 위한 장치가 부족해집니다. 또한, 종종 병렬 그룹이 느린 노드 간 링크를 가로지르도록 강제하여 NVLink 나 AMD Infinity Fabric 과 같은 고속 노드 내 상호 연결을 충분히 활용하지 못하게 합니다.

핵심 과제: prohibitive 한 통신 비용을 치르지 않으면서 단일 장치 축에서 파라미터 및 활성화 메모리를 동시에 줄여 메모리 효율성을 극대화하고 데이터 병렬성을 위한 장치를 확보하는 방법은 무엇인가?

2. 방법론: 텐서 및 시퀀스 병렬성 (TSP)

저자들은 텐서 병렬성과 시퀀스 병렬성을 단일 장치 축에 "접는 (folds)" 전략인 TSP를 제안합니다. TP 와 SP 를 직교하는 차원에 할당하는 대신, TSP 는 모든 랭크 (GPU) 에 모델 가중치의 쉐어드와 시퀀스의 쉐어드를 모두 할당합니다.

주요 설계 구성 요소:

• 통합 쉐이딩 (Unified Sharding): 각 장치는 모델 파라미터의 $1/D$와 시퀀스 토큰의 $1/D$를 보유합니다. 이는 단일 축에서 파라미터 비례 메모리와 활성화 메모리를 모두 $D$배 줄입니다.
• 런타임 스케줄: 가중치와 활성화가 다르게 쉐이딩된다는 사실을 처리하기 위해 TSP 는 레이어별로 특정 통신 스케줄을 사용합니다:
  • 어텐션 레이어: 브로드캐스트 (broadcast) 스케줄을 사용합니다.
    • 랭크들은 가중치 쉐어드 (패킹된 $W_Q, W_K, W_V, W_O$) 를 반복적으로 처리합니다.
    • 각 반복에서 하나의 랭크가 자신의 가중치 쉐어드를 다른 모든 랭크에게 브로드캐스트합니다.
    • 모든 랭크는 자신의 시퀀스 쉐어드에서 로컬 프로젝션을 계산합니다.
    • All-Gather 연산을 통해 시퀀스 차원에 걸쳐 전체 Key/Value 컨텍스트를 재구성하여 FlashAttention 을 수행합니다.
    • 지그재그 파티셔닝 (Zigzag partitioning) 방식을 사용하여 인과적 어텐션 워크로드를 균형 있게 분배합니다 (나중 토큰들이 불균형적으로 더 많은 작업을 수행하지 않도록 보장).
  • 게이트드 MLP 레이어: 링 (ring) 통신 스케줄을 사용합니다.
    • 가중치 쉐어드 ($W_1, W_2, W_3$) 가 링을 따라 포인트 투 포인트로 순환합니다.
    • 각 랭크는 로컬에서 부분 출력을 누적합니다.
    • 시퀀스는 로컬에 머무르고 가중치만 이동하므로, TP 에서 표준인 All-Reduce 연산이 필요하지 않습니다.
• 하드웨어 인식: AMD MI300X(Infinity Fabric) 와 같은 토폴로지에 최적화되어 있으며, 여기서 집단 대역폭은 참여 GPU 수에 따라 증가합니다. TSP 는 브로드캐스트 및 링과 같은 통신 원시 연산이 전체 그룹을 활용하거나 계산과 오버랩되도록 하여 효율성을 극대화합니다.

3. 주요 기여

1. 새로운 병렬성 전략: 두 개의 직교 병렬 축을 하나로 축소하여 파라미터 및 활성화 메모리를 동시에 줄일 수 있게 하는 TSP 도입.
2. 이론적 분석: DP, TP, SP, TP+SP 와 TSP 를 비교하는 엄격한 수학적 모델 제시:
   • 메모리: $D^2$개의 장치를 요구하지 않으면서 가중치 비례 메모리와 활성화 메모리를 모두 $1/D$만큼 줄이는 유일한 방식이 TSP 입니다.
   • 통신: TSP 는 TP+SP 가 피하는 가중치 이동 통신을 도입하지만, 저자들은 긴 시퀀스나 중간 크기의 배치 크기에서는 활성화 비례 통신이 지배적이므로 TSP 가 경쟁력이 있음을 증명합니다.
   • FLOPs: TSP 는 TP 및 SP 와 동일한 장치당 계산 효율성을 유지합니다.
3. 구현 및 오버랩: GEMM 계산 뒤로 가중치 전송을 파이프라인화하는 명시적인 통신/계산 오버랩 전략을 포함한 TSP 어텐션 및 MLP 블록의 상세 구현.
4. 토폴로지 최적화: TSP 가 전체 모델 복제본을 단일 고속 노드 (노드 내) 에 맞출 수 있게 하여, 종종 2D TP+SP 메시에서 요구되는 느린 노드 간 링크를 가로지르는 성능 패널티를 피할 수 있음을 시연.

4. 결과

저자들은 7B 파라미터 밀도 트랜스포머 모델을 사용하여 최대 1,024 개의 AMD MI300X GPU(128 노드) 클러스터에서 TSP 를 평가했습니다.

• 메모리 효율성:
  • TSP 는 테스트된 모든 시퀀스 길이 (16K~128K 토큰) 에서 최소 피크 메모리를 일관되게 달성했습니다.
  • 짧은 컨텍스트에서는 파라미터 메모리가 지배적이어서 TSP 가 TP 와 일치했습니다.
  • 긴 컨텍스트에서는 활성화 메모리가 지배적이어서 TSP 가 TP 및 TP+SP 를 크게 능가했으며, 다른 전략들이 실패한 메모리 부족 (OOM) 오류를 방지했습니다.
• 처리량:
  • 가중치 이동으로 인한 더 높은 이론적 통신 볼륨에도 불구하고, TSP 는 TP+SP 와 경쟁력 있는 처리량을 유지했습니다.
  • 추가 통신은 파이프라이닝을 통해 계산 뒤에 성공적으로 숨겨졌습니다.
  • 처리량 이점은 병렬화 정도 ($D$) 와 시퀀스 길이가 증가함에 따라 커졌습니다.
• 확장성:
  • TSP 는 감소된 메모리 발자국으로 인해 더 큰 마이크로 배치를 허용하여 처리량을 추가로 개선했습니다.
  • 단일 노드 (8 GPU) 내에 모델 복제본을 완전히 배치할 수 있게 하여, 더 큰 클러스터에서 데이터 병렬성 (DP) 을 위한 랭크를 확보했습니다.

5. 의의

• 하드웨어 인식형 확장: 노드 내 대역폭은 높지만 노드 간 대역폭이 병목인 현대 HPC 클러스터에 대한 실용적인 솔루션을 제공합니다. 이는 값비싼 노드 내 상호 연결의 활용도를 극대화합니다.
• 긴 컨텍스트 실현 가능성: 모델이 더 긴 컨텍스트 윈도우 (10 만 개 이상의 토큰) 로 이동함에 따라 활성화 메모리가 주요 병목이 됩니다. TSP 는 2 차원 병렬화의 막대한 장치 오버헤드 없이 이를 처리할 수 있는 유일한 위치에 있습니다.
• 유연성: TSP 는 파이프라인 병렬성, 전문가 병렬성, 데이터 병렬성과 직교적으로 구성될 수 있어 시스템 설계자가 특정 하드웨어 토폴로지에 맞춰 최적화할 수 있는 유연한 "축"을 제공합니다.
• 비용 효율성: 모델 병렬성을 위해 필요한 장치 수를 줄임으로써, 더 많은 장치를 데이터 병렬성에 사용할 수 있게 하여 주어진 학습 작업에 필요한 전체 클러스터 크기를 줄이거나 유효 배치 크기를 증가시킬 수 있습니다.

결론적으로, TSP 는 관리 가능한 통신 볼륨 증가를 대가로 메모리 효율성과 하드웨어 활용도의 상당한 향상을 얻는 분산 학습 전략의 중요한 진화를 나타내며, 차세대 긴 컨텍스트 LLM 학습을 위한 핵심 도구가 됩니다.