[논문 리뷰] Pathways: Asynchronous Distributed Dataflow for ML

Paul Barham, Aakanksha Chowdhery외 14명|arXiv (Cornell University)|2022. 03. 23.

Parallel Computing and Optimization Techniques인용 수 37

한 줄 요약

Pathways는 샤딩된 데이터흐름 그래프와 중앙집중식 스케줄링을 통해 수천 대의 가속기에서 고활용을 달성하고 이질적, MPMD 패턴을 지원하는 단일 컨트롤러의 비동기 분산 데이터 흐름 시스템을 제공합니다.

ABSTRACT

We present the design of a new large scale orchestration layer for accelerators. Our system, Pathways, is explicitly designed to enable exploration of new systems and ML research ideas, while retaining state of the art performance for current models. Pathways uses a sharded dataflow graph of asynchronous operators that consume and produce futures, and efficiently gang-schedules heterogeneous parallel computations on thousands of accelerators while coordinating data transfers over their dedicated interconnects. Pathways makes use of a novel asynchronous distributed dataflow design that lets the control plane execute in parallel despite dependencies in the data plane. This design, with careful engineering, allows Pathways to adopt a single-controller model that makes it easier to express complex new parallelism patterns. We demonstrate that Pathways can achieve performance parity (~100% accelerator utilization) with state-of-the-art systems when running SPMD computations over 2048 TPUs, while also delivering throughput comparable to the SPMD case for Transformer models that are pipelined across 16 stages, or sharded across two islands of accelerators connected over a data center network.

연구 동기 및 목표

• 전통적인 SPMD MPI 스타일 모델을 넘어서 이질적이고 희소한 워크로드를 지원하는 프로그래머블하고 확장 가능한 ML 시스템의 필요성을 제시한다.
• 샤딩된 데이터흐름과 비동기 디스패치를 갖춘 단일 컨트롤러 아키텍처를 제안하여 다중 컨트롤러 성능에 맞춘다.
• 가용성 향상과 새로운 병렬 패턴 지원을 위해 중앙 집중식 자원 관리 및 갱 스케줄링을 가능하게 한다.
• 비-SPMD 연산을 수용하고 JAX/TensorFlow와 같은 기존 프레임워크와의 용이한 통합을 지원하는 프로그래밍 모델을 제공한다.

제안 방법

• 샤딩된 데이터흐름 그래프를 도입하여 각 노드가 컴파일된 함수이며 간선은 데이터 전송을 표현한다.
• 가속기 섬에 걸쳐 가상 디바이스 슬라이스를 할당하고 이를 물리적 장치에 매핑하는 중앙 집중식 자원 관리자를 구현한다.
• Pathways IR로부터 희소하고 고처리량의 DCN 통신을 가능하게 하는 Plaque 기반의 교차 호스트 조정을 사용해 저수준 데이터흐름 프로그램을 형성한다.
• 일반적인(예측 가능한) 계산을 위한 스케줄링 및 데이터 전송과 호스트 측 작업을 중첩시키기 위해 병렬 비동기 디스패치를 통합한다.
• 섬마다 하나의 중앙 스케줄러를 두어 실행 순서를 정하고 다수의 샤드에 걸친 SPMD 연산의 갱 스케줄링을 가능하게 한다.
• CPU, GPU, 가속기 메모리 전반의 버퍼를 관리하기 위한 샤딩된 오브젝트 스토어를 유지하고 역압(back-pressure) 및 GC를 포함한다.

실험 결과

연구 질문

• RQ1 단일 컨트롤러 비동기 데이터흐름 시스템이 대규모 가속기에서 다중 컨트롤러 ML 시스템의 성능에 맞출 수 있는가?
• RQ2 중앙 집중식 자원 관리와 갱 스케줄링이 수천 대의 가속기에 걸친 효율적인 MPMD 및 이질적 워크로드를 어떻게 가능하게 할 수 있는가?
• RQ3 성능을 유지하면서 비-SPMD 워크로드를 지원하게 하는 프로그래밍 모델 및 IR 설계는 무엇인가?
• RQ4 IPIs, DCN, 인터커넥트 전반에서 높은 활용도를 달성하기 위한 데이터 이동 및 조정 메커니즘은 무엇인가?
• RQ5 다양한 구성에서 JAX, TF, Ray와 비교해 Pathways가 마이크로 벤치마크 및 엔드투엔드 ML 워크로드에서 어떤 성능을 보이는가?

주요 결과

• Pathways는 2048 TPU에서 SPMD 연산을 실행할 때 약 100%의 가속기 활용도를 달성한다.
• Pathways는 데이터 센터 네트워크로 연결된 두 가속기 섬에 걸쳐 샤딩되거나 16단계로 파이프라인화된 Transformer 모델에 대해 SPMD 기준선과 유사한 처리량을 제공합니다.
• 병렬 비동기 디스패치는 호스트 측 병목을 줄이고 파이프라인 단계가 증가할수록 순차적 디스패치보다 더 나은 처리량을 얻는다.
• Pathways는 충분히 큰 연산에 대해 JAX 처리량에 일치하고 수천 대의 가속기로 확장되며, 마이크로 벤치마크에서 일부 단일 컨트롤러 프레임워크를 능가한다.
• 이 시스템은 다중 테넌시를 지원하여 높은 총처리량과 효율적인 자원 공유로 동시 프로그램 실행을 가능하게 한다.