기술 요약: PipeFusion

문제 정의

Diffusion Transformers(DiT) 는 고해상도 이미지 및 비디오 합성을 위한 선도적인 아키텍처로 부상하여 기존 U-Net 모델을 능가했습니다. 그러나 DiT 는 특히 고해상도 이미지나 긴 시퀀스를 생성할 때 추론 중 심각한 지연 시간 문제를 겪습니다. 이 지연 시간은 시퀀스 길이에 따라 attention 연산이 이차적으로 증가하는 데서 기인합니다.

기존의 병렬 추론 전략들은 멀티 GPU 환경에서 이러한 문제를 효율적으로 해결하는 데 어려움을 겪고 있습니다:

• 텐서 병렬성 (Tensor Parallelism, TP): 대규모 언어 모델 (LLM) 에서는 효과적이지만, DiT 에서는 큰 활성화 (activation) 볼륨으로 인해 비효율적입니다. 통신 오버헤드 (특히 All-Reduce 연산) 가 병렬 연산의 이점을 종종 상쇄하며, 레이어가 많은 모델의 경우 더욱 두드러집니다.
• 시퀀스 병렬성 (Sequence Parallelism, SP): DeepSpeed-Ulysses 및 Ring-Attention 과 같은 방법은 입력 시퀀스 (이미지 패치) 를 장치 간에 분할합니다. 메모리 부담을 줄이지만, 모든 레이어마다 Query, Key, Value 활성화 데이터를 광범위하게 교환해야 하므로 통신 비용이 매우 높습니다.
• DistriFusion: 연속된 확산 시간 단계 간의 입력 및 활성화 유사성인 입력 시간적 중복성을 활용하는 비동기 시퀀스 병렬성 방법입니다. 통신과 연산을 겹치지만 메모리 비효율성으로 고통받습니다. DistriFusion 은 모든 장치에서 모든 레이어에 대한 Key(K) 와 Value(V) 활성화의 전체 공간 차원을 유지해야 하므로, 메모리 비용이 시퀀스 길이에 비례하여 증가하며 Flux.1 과 같은 대규모 모델의 효과적인 확장을 방해합니다.

방법론: PipeFusion

본 논문은 DiT 추론을 위해 특별히 설계된 새로운 패치 레벨 파이프라인 병렬성 전략인 PipeFusion을 제안합니다. PipeFusion 은 모델 레이어와 입력 이미지 패치를 여러 GPU 에 걸쳐 동시에 분할함으로써 기존 방법들의 한계를 해결합니다.

핵심 메커니즘

1. 이중 분할:

   • 레이어 분할: DiT 백본을 데이터 흐름 방향으로 분할합니다. 각 GPU 에는 연속된 트랜스포머 레이어 집합이 할당됩니다.
   • 패치 분할: 입력 이미지를 $M$ 개의 겹치지 않는 패치로 나눕니다. 각 GPU 는 할당된 레이어를 통해 하나의 패치를 파이프라인 방식으로 처리합니다.

2. 시간적 중복성 활용:

   • 확산 과정은 연속된 시간 단계 간의 입력 및 활성화에서 높은 유사성을 보입니다. PipeFusion 은 이전 확산 시간 단계 ($T+1$) 의 "오래된 (stale)" 특징 맵을 재사용하여 현재 시간 단계 ($T$) 에 대한 컨텍스트를 제공함으로써 이를 활용합니다.
   • 모든 장치로부터 현재 시간 단계의 신선한 활성화를 기다리는 대신, 장치는 이전 단계의 오래된 활성화를 사용하여 계산을 시작할 수 있습니다. 이로써 초기 시작 시를 제외하고 파이프라인이 최소한의 유휴 시간 (버블) 으로 지속적으로 실행될 수 있습니다.

3. 통신 최적화:

   • 모든 레이어에 대한 활성화를 전송하는 SP 및 TP 와 달리, PipeFusion 은 특정 스테이지에 대해 초기 레이어의 입력과 최종 레이어의 출력만 전송합니다.
   • 비동기 포인트 투 포인트 (P2P) 통신을 사용하여 마이크로 스텝 패치 활성화를 후속 장치로 전송함으로써 통신과 연산을 효과적으로 겹칩니다.

4. 메모리 효율성:

   • PipeFusion 은 모델 파라미터를 장치 간에 분산하여 파라미터 메모리 사용량을 장치 수 ($N$) 배만큼 줄입니다.
   • 특히 attention 에 필요한 KV 버퍼와 관련하여: PipeFusion 은 해당 장치에 할당된 레이어에 대해 전체 공간 K 및 V 텐서의 $1/N$만 저장하면 됩니다. 반면, DistriFusion 은 모든 장치에서 모든 레이어에 대한 전체 공간 K 및 V 텐서를 필요로 합니다. 이로 인해 PipeFusion 은 대규모 모델에서 훨씬 더 메모리 효율적입니다.

주요 기여

• 새로운 병렬성 전략: DiT 의 풀-attention 메커니즘에 맞춰 레이어 단위 파이프라인 병렬성과 패치 레벨 시퀀스 분할을 결합한 PipeFusion 도입.
• 감소된 통신 비용: PipeFusion 은 레이어별 활성화가 아닌 스테이지 입력/출력만 전송하여 장치 간 통신을 줄입니다. 이론적 분석에 따르면 통신 비용은 레이어 수 ($L$) 와 무관하며 시퀀스 길이 ($p$) 와 히든 크기 ($h_s$) 에만 비례합니다.
• 향상된 메모리 확장성: 파라미터를 분산하고 필요한 KV 버퍼 크기를 $1/N$으로 줄임으로써, PipeFusion 은 다른 방법들이 메모리 부족 (OOM) 오류로 실패하는 상용 하드웨어에서 대규모 모델 (예: 12B Flux.1) 의 추론을 가능하게 합니다.
• 품질 보존: 이론적 분석과 실험 결과에 따르면 PipeFusion 은 확산 과정에서 더 넓은 영역의 "신선한" 활성화를 활용하고 DistriFusion 은 오래된 데이터에 더 많이 의존하기 때문에, DistriFusion 보다 더 높은 생성 품질을 유지합니다.

실험 결과

저자들은 8×L40 PCIe GPU 클러스터에서 PixArt(0.6B), Stable Diffusion 3(2B), Flux.1(12B) 세 가지 최첨단 DiT 모델을 사용하여 PipeFusion 을 평가했습니다.

• 지연 시간 성능:

  • PipeFusion 은 테스트된 모든 해상도 (1024px~4096px) 및 모델에서 가장 낮은 추론 지연 시간을 달성했습니다.
  • 8×L40 GPU 에서 PipeFusion 은 PixArt 및 SD3 에 대해 최상의 시퀀스 병렬성 기준 (USP) 보다 1.48 배에서 1.55 배 더 빠릅니다.
  • 12B Flux.1 모델의 경우 PipeFusion 은 최상의 시퀀스 병렬성 대비 1.23 배에서 1.25 배의 속도 향상을 달성했습니다.
  • 고해상도 작업 (4096px) 에서 DistriFusion 은 OOM 문제를 겪은 반면, PipeFusion 은 전체 지연 시간 중 통신 점유율이 4.6% 에 불과 (SP 는 17.9%) 하여 강력한 성능을 유지했습니다.

• 메모리 효율성:

  • PipeFusion 은 우수한 메모리 확장성을 입증했습니다. Flux.1 의 2048px 해상도에서 PipeFusion 의 총 메모리 사용량은 시퀀스 병렬성이 필요로 하는 양의 **32%~36%**에 불과했습니다.
  • DistriFusion 이 실패한 PixArt 의 8192px 해상도 작업도 성공적으로 처리했습니다.

• 이미지 품질:

  • Fréchet Inception Distance(FID) 를 사용한 정량적 평가에서 PipeFusion 은 DistriFusion 일관되게 우세하여, 원래 단일 장치 구현에 더 근접한 점수를 달성했습니다.
  • 시각적 비교는 PipeFusion 이 생성한 이미지들이 원래 출력과 거의 구별할 수 없음을 확인했습니다.

중요성 및 주장

본 논문은 PipeFusion 이 DiT 추론을 병렬화하는 방식에 근본적인 변화를 가져왔다고 주장합니다. 입력 시간적 중복성을 체계적으로 활용함으로써, 기존 방법들이 달성할 수 없는 통신 효율성, 메모리 풋프린트, 생성 품질 간의 균형을 이룹니다.

저자들은 PipeFusion 을 모든 병렬 전략을 대체하는 보편적인 솔루션이 아닌, 하이브리드 병렬화 계획을 위한 매우 효과적인 구성 요소로 위치시킵니다. PipeFusion 은 특히 노드 간, 대역폭이 낮은 환경에 적합하며, 시퀀스 병렬성은 노드 내, 고대역폭 (NVLink) 환경에서 효과적이라고 제안합니다.

이 연구는 확산 과정의 시간적 중복성이 새로운 형태의 병렬성을 설계하기 위한 실현 가능하고 미활용된 자원임을 강조하며, 상용 멀티 GPU 클러스터에서의 확장 가능한 고해상도 이미지 생성을 위한 길을 제시합니다. 저자들은 또한 이 접근 방식이 유사한 MM-DiT 아키텍처와 더 긴 시퀀스 길이를 공유하는 비디오 생성 모델과 자연스럽게 호환된다고 지적합니다.