기술 요약: SwiftFusion

문제 정의

Diffusion Transformers(DiT) 는 고품질 이미지 및 비디오 생성에 점점 더 널리 사용되고 있습니다. 그러나 고해상도 이미지나 장시간 비디오를 생성하면 시퀀스 길이가 길어져 활성화 텐서 (activation tensors) 가 시퀀스 길이에 비례하여 선형적으로 증가합니다. 이로 인해 단일 GPU 에서 메모리 제약과 높은 지연 시간이 발생하며, 시퀀스 병렬성 (Sequence Parallelism, SP) 을 통한 분산 추론이 필요해집니다.

기존의 SP 기술인 Ring Attention, Ulysses Attention 및 이들의 조합인 Unified Sequence Parallelism(USP) 은 여러 GPU 머신으로 확장할 때 세 가지 주요 한계에 직면합니다:

1. 비최적의 토폴로지 활용: 현재 프레임워크는 머신 간 (예: NVSwitch) 과 머신 간 (예: EFA/RDMA) 네트워크 간의 상당한 대역폭 차이를 고려하지 못하는 경우가 많습니다. 예를 들어, USP 는 머신 간 통신에 Ring Attention 을 사용하는데, 이는 머신이 늘어나도 감소하지 않는 높은 통신 볼륨을 초래하여 병목 현상을 만듭니다.
2. 중첩되지 않는 통신: Ulysses Attention 에 의존하는 기술은 원자적 연산으로 취급되는 all-to-all 연산자를 활용합니다. 이러한 연산은 계산을 차단하여 지연 시간 버블을 생성하며, 특히 대역폭이 낮은 머신 간 통신에서 문제가 됩니다.
3. 동기화 및 계산 오버헤드: 기존 구현은 엄격한 송신자 - 수신자 동기화가 필요한 양방향 통신 라이브러리 (예: NCCL) 에 의존합니다. 이는 암묵적인 장벽을 도입하고 모델 계산과 경쟁하는 커널 기반 통신을 위해 GPU 스트리밍 멀티프로세서 (SM) 를 소모하여 오버헤드를 유발합니다.

방법론

저자들은 이러한 과제를 해결하기 위해 세 가지 핵심 혁신을 통해 토폴로지 인식 효율적 DiT 서빙 엔진인 SwiftFusion을 제안합니다:

1. 토폴로지 인식 통신 스케줄링

SwiftFusion 은 네트워크 토폴로지에 기반하여 SP 기술의 할당을 재검토합니다. 머신 간 통신에 Ring Attention 을 사용하는 USP 와 달리, SwiftFusion 은 이 전략을 반대로 적용합니다:

• 머신 간: Ulysses Attention을 사용합니다. Ulysses 는 일부 컨텍스트에서 머신당 통신 볼륨이 높을 수 있지만, 머신 수가 증가함에 따라 all-to-all 볼륨이 감소하므로 확장 시 느린 머신 간 네트워크에 더 적합합니다.
• 머신 내: Ring Attention을 사용합니다. 머신 내 네트워크 (예: NVSwitch) 는 높은 대역폭과 낮은 지연 시간을 제공하므로, Ring Attention 의 더 높은 통신 볼륨은 여기서 병목 현상이 덜합니다.
  이러한 스케줄링은 총 머신 간 통신 볼륨을 최소화하고, 통신 집약적인 Ring 연산에 높은 대역폭의 머신 내 링크를 활용합니다.

2. 토러스 어텐션 (통신 - 계산 중첩)

Ulysses Attention 에 필요한 all-to-all 연산자의 지연 시간 (표준 구현에서는 중첩되지 않음) 을 완화하기 위해 SwiftFusion 은 Torus Attention을 도입합니다.

• 청킹 (Chunking): all-to-all 연산자를 여러 청크로 분할합니다.
• 파이프라이닝: all-to-all 연산의 특정 단계에서 전송이 필요 없는 텐서의 "정지" 요소를 식별합니다. 이러한 요소는 즉시 어텐션 계산에 사용됩니다.
• 중첩 (Overlap): 정지 요소가 계산되는 동안, 시스템은 백그라운드에서 나머지 청크의 전송을 시작합니다. 이는 이전 청크의 계산과 한 청크의 통신이 중첩되는 파이프라인을 만들어 머신 간 통신 지연 시간을 효과적으로 숨깁니다.

3. 원측면 통신 구현

SwiftFusion 은 NCCL 과 같은 전통적인 양방향 라이브러리를 **원측면 통신 라이브러리 (예: NVSHMEM)**로 대체합니다.

• 메커니즘: 상호 동기화가 필요한 송신 - 수신 원시 연산 대신 푸시 (push) 및 풀 (pull) 원시 연산을 사용합니다.
• 최적화: 저자들은 Torus, Ulysses, Ring Attention 을 통합하는 통합 통신 알고리즘을 설계합니다. 이 접근 방식은 동기화 장벽을 최소화하여 머신 내 동기화와 레이어의 시작 및 끝에서 두 가지 특정 머신 간 동기화만 요구합니다.
• 효율성: 전송별 동기화를 제거하고 SM 을 위해 경쟁하는 CUDA 커널 기반 통신을 피함으로써 시스템은 동기화 및 계산 오버헤드를 모두 줄입니다.

주요 기여

본 논문은 다음과 같이 기여 사항을 요약합니다:

1. 토폴로지 인식 스케줄링: Ulysses Attention 을 머신 간 통신에, Ring Attention 을 머신 내 통신에 할당함으로써 현대 GPU 머신 토폴로지를 더 잘 활용하는 새로운 전략.
2. Torus Attention: all-to-all 연산자를 분할하여 통신을 계산과 파이프라인화하고 머신 간 지연 시간을 숨기는 새로운 분산 어텐션 알고리즘.
3. 원측면 구현: 동기화 및 계산 오버헤드를 최소화하기 위해 원측면 라이브러리를 사용하여 Torus, Ulysses, Ring Attention 을 통합한 전용 구현.
4. 실증적 평가: 이미지 (Flux) 및 비디오 (CogVideoX) 생성 워크로드에 대한 포괄적인 평가를 통해 최첨단 성능을 입증.

결과

저자들은 AWS p4de.24xlarge 인스턴스 (머신당 8 개의 NVIDIA A100 GPU, NVSwitch 및 EFA 로 상호 연결) 에서 최첨단 USP 프레임워크와 비교하여 SwiftFusion 을 평가했습니다.

• 전체 성능: SwiftFusion 은 다양한 분산 구성 및 워크로드 (Flux 3072×3072/4096×4096, CogVideoX 20 초/40 초) 에서 USP 대비 평균 1.35 배의 속도 향상을 달성했으며, 최대 1.77 배의 속도 향상을 기록했습니다.
• 구성 요소 분석:
  • 2 개 이상의 GPU 머신을 사용할 때 토폴로지 인식 스케줄링 (TAS) 만으로도 USP 대비 1.27 배의 속도 향상을 제공했는데, 이는 주로 머신 간 통신 볼륨을 줄였기 때문입니다.
  • Torus Attention 과 원측면 통신을 추가한 SwiftFusion 은 TAS 대비 추가적인 1.35 배의 속도 향상을 제공했는데, 이는 통신과 계산의 중첩 및 감소된 동기화 오버헤드 때문입니다.
• 확장성: SwiftFusion 은 다양한 GPU 머신 수 (최대 4 개) 와 다양한 분산 구성 (변동하는 Ulysses 및 Ring 병렬화 정도) 에서 USP 를 일관되게 능가했습니다.
• 메모리: 이 프레임워크는 USP 대비 더 높은 메모리 소비를 도입하지 않았으며, 텐서에 대해 최대 하나의 복사 버퍼만 사용했습니다.
• 워크로드 일반화: 성능 향상은 다양한 시퀀스 길이, 배치 크기 및 헤드 차원에서 일관되었으나, 계산 포화 상태 달성이 더 쉬운 더 큰 헤드 차원에서는 향상이 약간 더 두드러졌습니다.

의의

본 논문은 SwiftFusion 이 효율적이고 확장 가능한 분산 DiT 추론을 위한 새로운 패러다임을 제공한다고 주장합니다. 네트워크 토폴로지와 통신 패턴 간의 불일치를 명시적으로 해결하고, 통신 지연 시간을 숨기며 동기화 오버헤드를 줄이는 메커니즘을 도입함으로써 SwiftFusion 은 기존 SP 기술의 통신 제한을 극복합니다. 저자들은 이러한 통찰력이 미래 생성형 AI 애플리케이션의 고해상도 및 장시간 요구 사항에 맞춰 DiT 추론을 확장하는 데 필수적이라고 주장합니다.