🤖 machine learning

AoiZora: Topology-Aware Auto-Parallel Optimization for Inference of Diffusion Transformers

AoiZora는 물리적 인터커넥트 인지 능력을 자동 병렬화 탐색 프로세스에 통합함으로써 TPU 서브 슬라이스 상의 저지연 비디오 확산 추론을 최적화하는 컴파일러 매개 토폴로지 플래너로, 기존 솔루션 대비 Wan 2.1에 대해 최대 1.42배의 지연 시간 감소를 달성한다.

원저자: Kaijian Wang, Yuanyuan Xu, Fanjiang Ye, Ye Cao, Jingwei Zuo, T. S. Eugene Ng, Yarong Mu, Yuke Wang

게시일 2026-06-17

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Kaijian Wang, Yuanyuan Xu, Fanjiang Ye, Ye Cao, Jingwei Zuo, T. S. Eugene Ng, Yarong Mu, Yuke Wang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 비디오 제작을 위해 거대하고 중요한 수준의 이어달리기 경주를 조직한다고 상상해 보세요. 여기서 "주자"는 컴퓨터 칩(구체적으로 구글의 TPU 칩)이며, "바톤"은 비디오의 각 프레임을 생성하는 데 필요한 데이터입니다.

비디오 생성 과정은 마치 주자들이 이미지가 선명해질 때까지 반복해서 노이즈를 제거(이를 "디노이징"이라고 합니다)하는 길고 반복적인 루프와 같습니다. 이 비디오는 매우 방대하고 과정이 느리기 때문에, 단 하나의 칩만으로는 수행할 수 없습니다. 따라서 여러 칩이 팀을 이루어 함께 협력해야 합니다.

문제점: "좌석 배치도"의 실수
현재 엔지니어들이 이 팀을 구성할 때, 그들은 논리적인 계획에 집중합니다: "주자 A는 배경을 담당하고, 주자 B는 하늘을, 주자 C는 사람을 담당한다." 이것은 마치 화이트보드에 할 일을 적어두는 것과 같습니다.

하지만 그들은 물리적인 현실, 즉 **이 주자들이 경기장의 정확히 어디에 앉아 있는가?**를 간과하곤 합니다.

실제 경기장(물리적 TPU 네트워크)에서는 어떤 주자들은 서로 바로 옆에 앉아 있어 바톤을 즉시 전달할 수 있습니다. 반면, 어떤 주자들은 경기장 반대편에 있어 바톤을 전달하는 데 시간이 걸리고 공유된 경로에서 교통 체증을 유발합니다.

이 논문은 설령 주자들에게 똑같은 과업(동일한 논리적 계획)을 부여하더라도, 단순히 주자들이 경기장의 어디에 앉느냐에 따라 경주 속도가 16.6% 더 느려지거나 빨라질 수 있음을 발견했습니다. 주요 원인은 주자가 달리는 거리가 아니라, 일부 주자들이 동일한 좁은 통로(공유 네트워크 링크)를 동시에 사용해야 해서 발생하는 교통 체증입니다.

해결책: AoiZora (스마트한 코치)
저자들은 AoiZora(일본어로 '푸른 하늘'이라는 뜻)라고 불리는 새로운 도구를 만들었습니다. AoiZora를 단순한 경기 계획을 짜는 것을 넘어, 실제 경기장 배치에 기반하여 좌석 배치도까지 설계하는 초스마트한 코치라고 생각해보세요.

AoiZola는 모든 가능한 좌석 배치 조합을 확인하는 것이 너무 느리고 비용이 많이 든다는 까다로운 문제를 해결합니다. 만약 모든 좌석 배치 방식을 테스트하려고 한다면, 실제 경주를 하는 것보다 계획을 세우는 데 더 많은 시간을 쓰게 될 것입니다.

AoiZora의 작동 방식 (2단계 전략)
AoiZora는 마치 채용 과정과 같은 영리한 "필터링 및 정교화(filter and refine)" 접근 방식을 사용합니다.

1단계: 빠른 서류 심사 (Pre-Compilation)
먼저, AoiZora는 저렴하고 빠른 방법을 사용하여 "이력서"(논리적 과업 계획)를 검토합니다. 아직 전체 경기장을 구축하는 단계가 아닙니다. 그저 계획이 서류상으로 타당한지만 확인합니다. 예를 들어, 나눌 수 없는 과업을 나누려고 시도하는 것과 같이 명백히 실패할 계획들을 빠르게 걸러냅니다. 이는 30초짜리 짧은 면접처럼 빠르고 저렴합니다.
2단계: 전체 시뮬레이션 (Post-Compilation)
첫 번째 라운드를 통과한 소수의 우수 후보들에 대해, AoiZora는 본격적인 작업을 수행합니다. 실제 경기장(물질적 TPU 칩)의 상세한 시뮬레이션을 구축하고 전체 테스트를 실행합니다. 이를 통해 특정 네트워크 링크에서 정확히 어디에 교통 체증이 발생하는지 확인합니다. 그런 다음, 이러한 정체를 피할 수 있는 단 하나의 최적의 좌석 배치를 선택합니다.

결과
이 2단계 방법을 통해 AoiZora는 "무엇을 할 것인가"와 "어디에서 할 것인가"의 완벽한 조합을 찾아냅니다.

비유: 당신이 저녁 식사 모임을 준비한다고 상상해 보세요.
- 기존 방식: 누가 어떤 음식을 가져올지 결정하고(논리적 계획), 아무 데나 앉으라고 말합니다. 가끔 무거운 수프를 가져오는 사람이 섬세한 케이크를 가져오는 사람 옆을 지나가야 해서, 서로 부딪혀 음식을 쏟는 일이 발생합니다.
- AoiZora 방식: 먼저 음식을 가져올 수 없는 손님들을 빠르게 탈락시킵니다. 그다음, 남은 손님들을 위해 실제 테이블 배치를 살펴보고, 수프를 가져오는 사람은 주방 문 바로 옆에, 케이크를 가져오는 사람은 반대편에 앉혀서 아무도 서로 부딪히지 않도록 합니다.

핵심 요약
구글의 TPU 칩에서 AoiZora는 기존 방식보다 비디오 생성을 1.42배 더 빠르게 만들었습니다. AoiZora는 비디오 게임 자체나 경기의 규칙을 변경한 것이 아닙니다. 단지 주자들이 서로 발을 헛디디지 않고 바톤을 전달할 수 있도록 가장 스마트한 좌석 배치 방법을 찾아냈을 뿐입니다. 이를 통해 기초 소프트웨어를 다시 작성할 필요 없이, 더 빠른 속도와 낮은 지연 시간의 비디오 생성이 가능해졌습니다.

기술 요약: AoiZora

문제 정의
Wan 2.1과 같은 모델을 위한 비디오 확산(diffusion) 추론은 매우 중요한 생성 워크로드로 부상했습니다. 그러나 추론 과정은 수십만 개의 시공간 토큰과 거대한 활성화 텐서에 대한 반복적인 디노이징 단계를 필요로 하므로 계산 집약적입니다. 이는 여러 가속기 간의 분산 추론을 필수로 만듭니다. 현재의 오토 패러럴(auto-parallel) 시스템(예: JAX/XLA 기반)은 최적의 논리적(logical) 장치 메쉬(텐서를 어떻게 샤딩할지 결정하는 과정)를 찾는 데는 효과적이지만, 이러한 논리적 랭크(logical ranks)를 하드웨어 상의 물리적(physical) 위치에 배치하는 문제는 거의 고려하지 않습니다.

TPU 서브 슬라이스(TPU 포드의 연속적인 직사각형 부분 집합)에서 인터-칩 인터커넥트(ICI)는 모든 차원에서 랩-어라운드 링크(wrap-around links)가 없는 근접 이웃 메쉬 구조를 가집니다. 결과적으로, 동일한 논리적 샤딩 전략이라도 논리적 랭크가 물리적 칩에 어떻게 매핑되느냐에 따라 최종 엔드 투 엔드 지연 시간이 크게 달라질 수 있습니다. 본 논문은 이러한 "배치 효과(placement effect)"만으로도 v5e-8 슬라이스에서 1단계 디노이징 지연 시간에 최대 16.6%의 변동을 일으킬 수 있음을 관찰했습니다. 이는 주로 홉 수(hop count)보다는 공유 링크의 경합(shared-link contention) 때문입니다. 기존 시스템은 컴파일러가 구체적인 집합 통신 그래프(collective communication graph)를 노출하기 전(대개 기본 순서를 통해)에 논리적 랭크를 물리적 장치에 조기에 바인딩하므로, 토폴로지 의존적인 성능 최-적화를 놓치게 됩니다.

방법론: 2단계 컴파일러 매개 플래너(Two-Stage Compiler-Mediated Planner)
AoiZora는 컴파일 분석과 실행 사이의 경계에 토폴로지 인식 계획(topology-aware planning)을 배치함으로써 이 격차를 해결합니다. 이 시스템은 JAX/XLA 컴파일 파이프라인의 "비용/충실도 비대칭성(cost/fidelity asymmetry)"을 활용합니다. 즉, 초기 중간 표현(IR)은 저렴하지만 구체적인 통신 세부 사항이 부족한 반면, 완전히 컴파일된 고수준 최적화(HLO)는 비용이 많이 들지만 정확한 집합 그래프와 장치 그룹을 노출한다는 점을 이용합니다.

플래너는 두 가지 별도의 단계로 작동합니다:

1단계: 프리-컴파일 논리적 샤딩 프루닝 (Pre-Compilation Logical Sharding Pruning)
- 목표: 비효율적인 논리적 샤딩 후보를 효율적으로 제거합니다.
- 메커니즘: AoiZora는 가능한 워크로드 축 할당(예: 텐서, 시퀀스, 데이터 병렬성)을 열거하고 이를 StableHLO/Shardy IR로 낮춥니다.
- 분석: 명시적 통신과 합성된 암시적 통신 노드(예: 리샤딩, 샤딩된 리덕션, 불투명 커널 경계)를 포함하는 예측 DAG를 구축하여 논리적 통신량을 추정합니다.
- 스코어링: 배치 불가지론적 처리량 상한( $Q_1$ )을 사용하여 후보를 평가합니다. 이 단계에서는 메모리 풋프린트가 나쁘거나 논리적 통신이 과도한 샤딩을 걸러내어, 소수의 "Top-K" 생존자만을 남깁니다.
2단계: 포스트-컴파일 토폴로지 랭킹 (Post-Compilation Topology Ranking)
- 목표: 생존한 논리적 샤딩들에 대해 최적의 물리적 배치를 선택합니다.
- 메커니즘: Top-K 논리적 후보들은 구체적인 HLO를 생성하기 위해 XLA에 의해 완전히 컴파일됩니다.
- 분석: 시스템은 서로 다른 물리적 축 순서(논리적 축을 TPU 서브 슬라이스 좌표에 매핑)를 열거합니다. 이후 컴파일된 HLO를 파싱하여 복제 그룹, 페이로드 크기, 의존성을 포함한 정확한 집합 그래프를 추출합니다.
- 스코어링: 토폴로지 인식 비용 모델( $Q_2$ )이 각 샤딩-배치 쌍의 순위를 매깁니다. 이 모델은 집합 연산을 위한 물리적 경로를 재구성하고, 링크 부하를 추정하며, 공유 ICI 링크에서 경합을 유발하는 배치를 패널티를 부여합니다.
- 선택: 가장 낮은 추정 비용을 가진 쌍이 선택됩니다.

주요 기여

배치 민감도 특성화: 저자들은 TPU 서브 슬라이스에서 물리적 배치가 일급 최적화 변수임을 입증했습니다. 동일한 논리적 샤딩이라도 공유 링크 경합으로 인해 최대 16.6%의 지연 시간 차이가 발생할 수 있음을 보여줍니다.
컴파일러 충실도 인식 탐색: 본 논문은 저렴한 프리-컴파일 IR(고재현율 프루닝에 적합)과 값비싼 컴파일된 HLO(정확한 토폴로지 랭킹에 필요) 사이의 비대칭성을 식별하고 활용합니다.
AoiZora 시스템 설계: 논리적 프루닝과 물리적 랭킹을 분리하는 컴파일러 매개 플래너입니다. 결정적으로, 이 시스템은 모델 코드, 로잉 규칙(lowering rules), 집합 커널 구현 또는 네트워크 라우팅을 수정하지 않고도 표준 JAX/XLA 스택과 통합됩니다. 이는 샤딩 전략과 물리적 축 순서를 출력하며, 이후 표준 컴파일러가 이를 실행합니다.
평가: 다양한 비디오 해상도(480p/720p) 및 프레임 수에 대해 TPU v5e 서브 슬라이스(4, 8, 16개 칩) 상의 Wan 2.1을 대상으로 광범위한 실험을 수행했습니다.

실험 결과
Wan 2.1을 대상으로 평가한 결과, AoiZora는 다음 두 베이스라인보다 일관되게 우수한 성능을 보였습니다:

기존 배포 방식: 고정된 물리적 배치를 사용하는 전문가 구성의 핸드-튜닝된 병렬 전략(FSDP, 시퀀스 및 텐서 병렬성 사용).
논리적 탐색 베이스라인: 추상적 비용 모델을 사용하여 논리적 샤딩을 최적화하지만, 포스트-컴파일 물리적 배치 랭킹을 수행하지 않는 Alpa 기반 시스템.

성능 향상: AoiZora는 기존 솔루션 대비 최대 1.42배, 논리적 탐색 베이스라인 대비 최대 1.37배의 속도 향상을 달립다. 기하 평균 속도 향상은 각각 1.09배 및 1.11배입니다.
충실도: 프리-컴파일 단계는 테스트된 모든 설정에서 측정된 최적의 논리적 샤딩을 성공적으로 유지(Top-1 재현율 100%)했으며, 포스트-컴파일 단계는 예측된 지연 시간과 실제 측정된 지연 시간 사이에 높은 순위 상관관계( $\rho_s$ 최대 0.985)를 달성했습니다.
오버헤드: 오프라인 계획 비용은 v5e-4의 경우 약 47초, v5e-16의 경우 약 376초로 실용적인 수준이며, 이는 동일한 형태의 반복적인 추론 요청 과정에서 상쇄(amortized)됩니다.

의의 및 주장
본 논문은 AoiZora가 물리적 TPU 토폴로지에 대한 변경 없이도 논리적 병렬성과 물리적 토폴로지 사이의 간극을 메운다고 주장합니다. 컴파일러가 구체적인 집합 그래프를 실체화할 때까지 물리적 배치 결정을 미룸으로써, AoiZora는 논리적 메쉬를 완전한 추상화로 취급하는 기존 오토 패러럴 시스템의 "간과(oversight)"를 방지합니다.

저자들은 랩-어라운드 링크의 부재로 인해 토폴로지 인식 배치가 매우 중요한 TPU 서브 슬라이스의 특정 제약 조건에 대한 실질적인 해결책으로 AoiZora를 제시합니다. 이들의 접근 방식은 모델 개발자가 토폴로지 세부 사항에 무관하면서도 비디오 확산 워크로드에 대해 최적에 가까운 추론 지연 시간을 달립 수 있게 해준다고 주장합니다. 이 연구는 정적인 텐서 형태를 가진 밀집 비디오 DiT 아키텍처의 경우, 일회성 오프라인 계획 단계가 고정된 물리적 배치에서 발생하는 경합 문제를 해결함으로써 서빙 성능을 크게 향상시킬 수 있음을 시사합니다.

유사한 논문