GPU 클러스터를 위한 효율적인 디퓨전 트랜스포머 시퀀스 병렬 엔진

초록

본 논문은 대규모 이미지·비디오 생성에 사용되는 디퓨전 트랜스포머(DiT)의 분산 추론을 위해, 기존 시퀀스 병렬 방식의 통신 비효율성을 극복한 ‘StreamFusion’ 엔진을 제안한다. 토폴로지 인식 통신 스케줄링, Torus Attention 기반 통신·연산 겹침, 그리고 일방향(NVSHMEM) 통신 구현을 결합해 평균 1.35배(최대 1.77배)까지 속도 향상을 달성한다.

상세 분석

StreamFusion은 세 가지 핵심 혁신을 통해 기존 시퀀스 병렬(USP, Ring, Ulysses) 방식의 한계를 극복한다. 첫째, 토폴로지 인식 통신 스케줄링은 현대 GPU 클러스터의 이질적인 대역폭( intra‑machine NVSwitch/NVLink vs. inter‑machine Ethernet/InfiniBand) 차이를 명시적으로 모델링한다. 이를 위해 각 GPU가 담당하는 시퀀스 조각을 intra‑machine에서는 Ring Attention(저대역폭, 저지연)으로, inter‑machine에서는 Ulysses Attention(고대역폭, 대용량 all‑to‑all)으로 처리하도록 재배치한다. 이러한 매핑은 기존 USP가 Ring Attention을 inter‑machine에 고정시킨 설계와 달리, 통신량이 기계 수에 따라 감소하지 않는 Ring Attention의 병목을 완화한다.

둘째, Torus Attention은 Ulysses Attention의 all‑to‑all 연산을 청크 단위로 분할하고, 현재 청크의 연산을 수행하면서 다음 청크의 all‑to‑all 전송을 파이프라인한다. 논문은 “정적 요소(stationary elements)”가 all‑to‑all 전·후에 존재한다는 관찰을 기반으로, 연산‑통신 겹침을 구현한다. 결과적으로 inter‑machine all‑to‑all에 소요되는 대기 시간이 연산 단계와 겹쳐져, 전체 지연이 크게 감소한다.

셋째, 기존 NCCL 기반 두‑방향 통신은 송·수신 동기화를 강제해 불필요한 스톰을 유발한다. StreamFusion은 NVSHMEM 같은 일방향 메모리 복사 방식을 채택해, 프로그래머가 명시적으로 동기화 시점을 제어한다. 이를 통해 GPU 간 sender‑receiver 동기화 오버헤드를 최소화하고, 연산 파이프라인을 더욱 매끄럽게 유지한다.

실험에서는 Stable Diffusion, CogVideoX 등 최신 DiT 모델을 다양한 GPU 머신 구성(28 머신, 각 48 GPU)에서 평가하였다. 특히 inter‑machine 대역폭이 제한적인 100 Gbps Ethernet 환경에서도 평균 1.35×, 최고 1.77×의 속도 향상을 기록했으며, 통신 프로파일링 결과 Torus Attention이 all‑to‑all 대기 시간을 45 % 이상 감소시켰음을 확인했다. 또한, 메모리 사용량은 기존 USP와 동등하거나 약간 감소해, 대형 시퀀스(예: 10 초 768×1360 비디오)에서도 OOM 없이 추론이 가능했다.

전반적으로 StreamFusion은 토폴로지‑감지 스케줄링, 청크‑기반 겹침, 일방향 통신이라는 세 축을 결합해, DiT와 같은 대규모 트랜스포머 모델의 분산 추론을 실용적인 수준으로 끌어올렸다. 향후 NVLink‑기반 다중 노드 클러스터나, 더 높은 차원의 비디오 생성 작업에도 동일한 설계 원칙을 적용할 수 있을 것으로 기대된다.