Pathways: Asynchronous Distributed Dataflow for ML 논문 정리

논문 : https://arxiv.org/pdf/2203.12533.pdf

Introduction

• distributed ML을 위한 구글이 제시하고 실제 사용중인 새로운 시스템인 Pathways에 대한 논문입니다. 실 사용례로는 google의 llm인 palm 학습에 사용되었습니다.
• 결론부터 말하자면 Pathways는 asynchronous distributed dataflow를 사용하여 여러 개의 가속기에서 동시에 작업을 수행할 수 있으며, 이를 통해 높은 처리량과 낮은 대기 시간을 제공하는 도구라입니다.
• Very large language model이 부상하면서 SPMD computation으로는 기대하는 성능을 볼 수 없었고, MoE, Mpi-style system 등 pipeline을 적극 도입하고 있으나 이 정도로는 부족합니다.
• 현재는 “islands”와 같은 형태로 single user program이 단일한 accelerators를 점유하기 떄문에 accelerator를 busy하게 지속시키지 않으면 낭비될 수 밖에 없습니다.
• 이러한 문제를 해결하기 위해 Pathways는 최신 ML 시스템의 기능과 성능을 충족시키면서 미래의 ML 작업부하를 지원하기 위해 필요한 기능을 제공합니다.
• 클라이언트-서버 아키텍처를 사용하여 많은 클라이언트를 대신하여 system-managed islands of compute 방식으로 프로그램을 실행합니다.
• “Pods” of TPUs(https://cloud.google.com/tpu/docs/training-on-tpu-pods?hl=ko)를 포괄하는 프로그램을 투명하고 효율적으로 실행하도록 설계된 최초의 시스템이며, new dataflow execution model을 채택하여 수천 개의 accelerators에 확장할 수 있습니다.
• non-SPMD computations을 쉽게 지원하며, 중앙 집중식 자원 관리 및 가상화를 통해 accelerator utilization를 향상시킬 수 있습니다.

DESIGN MOTIVATION

• distributed ML system의 design은 종종 hardware accelerators에 의존되어 결정되지만 여기엔 몇 가지 문제가 있습니다.

Multi-controller system architecture

• 최신 SPMD 모델을 훈련하기 위해 multi-controller architecture가 자주 채택되지만. 이 아키텍처에서는 동일한 클라이언트 실행 파일이 시스템의 모든 호스트에서 직접 실행되며, 프로그램 실행 기간 동안 해당 호스트의 리소스를 독점적으로 소유합니다.
• 이 아키텍처의 예로는 MPI(Clarke et al., 1994), PyTorch(Paszke et al., 2019), JAX(Bradbury et al., 2018), TensorFlow(Shazeer et al., 2018; Agrawal et al., 2019) 등이 있습니다.
• 이 아키텍처의 주요 장점은 accelerator computations에 대한 low latency입니다. 사용자의 코드가 각 가속기 호스트에서 동일하게 실행되고, 전송은 상대적으로 빠른 PCIe link를 통해 이루어지기 때문입니다. 다른 호스트 간의 모든 통신은 전용 인터커넥트(예: NVLink(Foley and Danskin, 2017)와 ICI(Jouppi et al., 2020))를 사용하여 호스트 메모리를 거치지 않고 수행됩니다.
• 그러나 이 아키텍처는 파이프라인이나 computational sparsity를 사용하는 현대 ML 작업 부하에는 적합하지 않습니다.
• multi-controller system에서 standard collectives를 넘어서는 모든 통신은 사용자가 자체적으로 조정해야 합니다
• 또한 일반적으로 하드웨어 리소스를 독점적으로 소유하기 때문에 이비싼 가속기의 높은 활용도를 보장하는 책임을 사용자에게 전가할 뿐만 아니라, 효율적인 클러스터 전체 ML 인프라를 구축하기 위해 필요한 자원 가상화 및 다중화와 같은 기능의 설계를 복잡하게 만듭니다.

Single-controller system architecture

• 반면 TensorFlow v1(Abadi et al., 2016)과 같은 Single-controller system은 최적화된 in-graph control flow(Yu et al., 2018)를 통해 매우 일반적인 분산 데이터 흐름 모델을 제공하기도 합니다.
• TensorFlow(TF) Python 클라이언트는 computation graph를 빌드하고 이를 coordinator runtime에 전달합니다.
• coordinator runtime은 graph를 서브그래프(subgraph)로 분할하고 worker의 local runtime에 위임합니다.
• Single-controller design은 유연한 프로그래밍 모델과 리소스의 가상화를 제공하지만, 구현 과제를 제시합니다.

• 그렇다고 Single-controller system이 더 이상적이지는 않습니다.
• 첫 번째로, accelerator computations를 dispatch하기 위해 PCIe(Peripheral Component Interconnect Express)통신만을 하는 multi-controller system에 비해 Single-controller client는 더 멀리 떨어져 있습니다. PCIe보다 10배 정도 느린 DCN(datacenter network)에 의존하기 때문입니다.
• 두 번째로, MPMD(multiple-program multiple-data) 프로그램의 동시 실행을 지원하기 위해서 SPMD(single-program multiple-data) sub-computations을 포함하는 accelerators 서브셋(subset)의 가속기를 공유 클러스터에서 추출해야합니다.
• 이를 위해서는 Gang-scheduling(https://en.wikipedia.org/wiki/Gang_scheduling)이 지원되어야 하지만 TPU는 single-threaded로 동작하기 때문에 통신 계산이 일관된 순서로 queue에 삽입되지 않으면 시스템이 데드락에 빠지기 쉽습니다.
• 컴퓨터 과학에서 Gang scheduling은 parallel 시스템에서 관련된 스레드나 프로세스를 다른 프로세서에서 동시에 실행하도록 스케줄링하는 알고리즘입니다. (https://en.wikipedia.org/wiki/Gang_scheduling)
• 때문에 이를 해결하기 위한 distributed scheduling mechanism이 필요합니다.
• 현대 ML workload에서는 수천 개의 가속기에 분산되어 실행되도록 설계돼야 하고 sharded data structure를 지원해야 하나 매우 어렵습니다.

그래서 PATHWAYS는?

• PATHWAYS는 단일 컨트롤러 프레임워크의 유연성과 다중 컨트롤러의 성능을 결합합니다.
• 효율적인 ML 계산을 위해 더 많은 기회를 제공한다고 (google이) 믿는 single-controller를 채택하였습니다.
• 하지만 앞서 서술한 단점을 극복하고 multi-controller system의 성능을 따라가야 하기 때문에 asynchronous dispatch를 사용하고, Gang of SPMD 가속기 계산을 위한 중앙 자원 관리 및 스케줄링을 지원하며, efficient coordination을 위해 sharded dataflow system을 사용합니다.
• (대충 앞서 매우 어렵고 단점이라고 했던 서술들을 우리는 해냈다는 이야기)

PATHWAYS PROGRAMMING MODEL

• Pathways는 TF와 JAX를 지원하지만 해당 논문은 JAX를 중심으로 평가합니다.
• multi-controller configurations로 실행되는 JAX 프로그램이 XLA(shapes, bounded loops) 집단을 사용하여 모든 데이터를 전송하기 때문입니다.
• 하지만 PATHWAYS에서는 ICI와 DCN 모두에서 통신할 수 있기 때문에 가능합니다. PATHWAYS는 JAX 백엔드의 플러그인 대체품으로 사용할 수 있으며, JAX 코드가 수정되지 않은 상태로 실행될 수 있습니다.
• 이제 TPU pod의 local core뿐만 아니라 시스템에서 제공되는 많은 다른 pod의 코어에 액세스할 수 있습니다. JAX 프로그램은 이제 수천 개의 TPU 코어를 포함하는 여러 TPU pod까지 확장할 수 있습니다.

PATHWAYS SYSTEM ARCHITECTURE

• PATHWAYS는 XLA(TensorFlow, 2019)를 사용하여 TPU 계산을 표현하고 실행하며, TensorFlow 그래프와 실행기(Abadiet al. 2016)를 사용하여 분산 CPU 계산을 표현하고 실행합니다. 또한, JAX(Bradbury et al. 2018)와 TensorFlow API 등의 Python 프로그래밍 프레임워크를 활용하여 기존 ML 모델을 최소한의 코드 변경으로 실행할 수 있습니다. 이러한 구성 요소를 활용하여 PATHWAYS의 새로운 조정 측면에 초점을 맞출 수 있습니다

Resource Manager

• PATHWAYS backend는 DCN을 통해 서로 연결된 밀접하게 결합된 일련의 accelerator 세트가 있고 resource manager가 디바이스를 중앙에서 관리합니다.
• resource manager는 요구되는 interconnect topology, memory capacity 등을 만족하는 가상 디바이스에 대해 동적으로 물리적 디바이스를 할당합니다.
• resource manager가 사용 가능한 디바이스를 추적하면서 백엔드 계산 리소스를 동적으로 추가 및 제거할 수 있도록 합니다.

Client

• user가 program을 실행할 때 Pathways는 가상 디바이스를 할당하고 resource manager에 등록하여 computation을 background에서 compile하도록 trigger하는 client library를 call합니다.
• client는 intermediate representation(IR)을 생성하고 점진적인 compile을 거쳐 physical device간의 네트워크 연결을 고려하기 위해 physical device의 위치까지 포함하는 low-level program을 생성합니다.
• 데이터 교환이 필요한 경우 shard to shard 전송 작업을 포합하며 가상 디바이스의 위치가 변경되지 않는 경우 low-level program을 반복적으로 실행하기 때문에 효율적입니다.
• 위에서 서술한대로 기존의 single-controller system에서는 클라이언트가 수천 개의 개별 계산과 각 계산 shard에 해당하는 데이터 버퍼를 조정해야 하기 때문에 성능 병목이 있었지만, PATHWAYS 클라이언트는 여러 디바이스에 분산될 수 있는 논리적 버퍼를 나타내는 sharded buffer를 추상화하여 확장성을 향상시키는 데 도움을 줍니다.

Coordination implementation

• PATHWAYS는 DCN을 사용하는 모든 호스트 간 코디네이션에 PLAQUE를 의존합니다.
• PLAQUE는 구글에서 많은 고객 대상 서비스에 사용되는 sharded dataflow system으로, high-fanout 또는 high-fanin 통신이 필요하고 확장성과 대기 시간이 중요한 경우에 사용됩니다.
• low-level IR은 PLAQUE 프로그램으로 직접 변환되며, dataflow graph로 표현됩니다. PATHWAYS는 Coordination에 대해 엄격한 요구 사항을 가지고 있지만 이는 모두 PLAQUE에서 충족됩니다.
• coordination substrate는 스케줄링 메시지와 데이터 핸들을 전송하는 데 중요한 경로에 있는 DCN 메시지를 전송하는 데 사용되므로, 낮은 대기 시간으로 중요한 메시지를 전송하고, 높은 처리량이 필요할 때 동일한 호스트에 대한 메시지를 일괄 처리해야 합니다.

• PLAQUE 대신 Ray(Moritz et al., 2018)와 같은 다른 분산 프레임워크를 사용하여 전체 PATHWAYS 디자인을 다시 구현하는 것도 가능합니다.
• 이러한 구현에서는 PATHWAYS executors 및 schedulers가 Ray actor로 대체되며, 이는 Ray 클러스터 스케줄링 위에 PATHWAYS 스케줄링을 구현하고, executors PyTorch를 사용하여 GPU 계산 및 Collectives를 수행할 수 있습니다.
• Ray가 GPU 인터커넥트를 통해 원격 객체를 효율적으로 전송하는 능력이 부족하기 때문에, 유사한 성능을 달성하기 위해 추가적인 기능이 필요할 수는 있습니다.

Gang-scheduled dynamic dispatch

• 이전에 논의한 바와 같이, 공유된 가속기 세트에서 SPMD 계산을 지원하기 위한 한 가지 요구 사항은 효율적인 Gang scheduling을 지원하는 것입니다.
• PATHWAYS 런타임에는 island마다 중앙 집중식 스케줄러가 포함되어 있으며, island 내 모든 계산을 일관되게 정렬합니다.

• 스케줄러는 밀리초 단위의 시간 척도에서 가속기를 할당하는 정책을 구현해야 합니다. 현재 구현은 단순히 FIFO 순서로 작업을 enqueueing하지만, 더 정교한 스케줄러는 예를 들어 예상 실행 시간에 따라 계산을 다시 정렬할 수 있습니다.

Parallel asynchronous dispatch

• 가속기에서 계산을 실행할 때, 시스템은 비동기 API를 활용하여 계산과 조정을 오버랩할 수 있습니다(Kwon et al. 2020).
• three-node graph인 해당 그림의 예시로 A, B, C node라고 할 때, 모든 노드 계산은 regular한 컴파일된 함수입니다.

Sequential dispach에서

• 호스트 A는 node A를 enqueue하고 A의 출력에 대한 future를 수신하며, 호스트 B로 future를 전송합니다.
• 호스트 B는 B의 입력을 할당하고, 호스트 A로 입력 버퍼 주소를 전송하며, node B의 함수를 실행하기 위한 대부분의 준비 작업을 수행합니다.
• node A가 완료되면, 출력은 accelerator interconnect를 통해 직접 B의 입력 버퍼로 전송되며, 호스트 B는 node B를 시작합니다.
• 이러한 디자인은 computation이 스케줄링, 자원 할당, 호스트 간 조정에 소요되는 시간보다 더 오래 걸릴 때 잘 작동합니다.
• 그러나 계산 시간이 너무 짧으면, 한 노드가 완료되고 다음 노드가 시작되는 지연 시간이 데이터 전송 시간보다 높을 수 있기 때문에 비동기 파이프라인이 정체되고 호스트 측 작업이 전체 계산 순서를 실행하는 데 중요한 병목이 됩니다.
• 컴파일된 함수가 모두 regular하다면, 선행A의 계산이 완료되기 전에도 선행A의 입력 모양은 계산할 수 있습니다.

Parallel dispatch에서

• 따라서 아래 그림과 같은 보이는 새로운 병렬 비동기 디스패치 디자인을 도입합니다.

Data management

• 각 호스트는 Ray의 object store(Moritz et al., 2018)와 유사한 sharded object store를 관리합니다.
• 그러나 각 shard에서 accelerator HBM에 저장된 버퍼도 추적하도록 확장되었습니다.
• HBM(High Bandwidth Memory)은 고성능 컴퓨팅 분야에서 사용되는 고속 메모리 기술 중 하나입니다. HBM은 일반적인 DRAM(Dynamic Random Access Memory)과 달리, 메모리 칩 내부에 프로세서와 직접 연결되는 인터페이스를 가지고 있습니다. 이 때문에 HBM은 일반적인 DRAM보다 훨씬 빠른 속도로 데이터를 처리할 수 있습니다.(https://en.wikipedia.org/wiki/High_Bandwidth_Memory)
• 클라이언트 프로그램은 원격 호스트 또는 가속기 메모리에 있는 객체에 대한 참조를 유지할 수 있으며, 클라이언트와 서버는 필요한 경우 참조합니다.
• Intermediate program 값도 객체 저장소에 저장됩니다. 예를 들어 시스템이 이를 가속기로 전송하거나 다음 계산으로 전달하기를 기다리는 동안입니다. 객체는 소유권 라벨이 붙어 있어서 프로그램이나 클라이언트가 실패하면 가비지 수집될 수 있습니다. 우리는 다른 계산의 버퍼가 일시적으로 HBM을 차지하고 있기 때문에 메모리를 할당할 수 없는 경우 간단한 백-프레셔를 사용하여 계산을 지연시킬 수 있습니다.

Conclusions

지금까지 PATHWAYS를 구현하기 위한 design motivation, programming model, architecture를 살펴보았는데요. 앞서 서술된 개념들을 조합하여 google은 사용자 코드를 single-controller 모델로 되돌리며, 중앙 집중식 자원 관리 및 스케줄링 프레임워크를 클라이언트와 가속기 사이에 삽입하여 multi-controller의 성능을 제공하면서도 single-controller 프로그래밍 모델로 사용자에게 훨씬 풍부한 계산 패턴에 대한 간단한 접근을 허용할 수 있는 distributed ML train system을 구축할 수 있었다고 합니다.

사실 개념이 많이 추상화 되어 있어 이해하기는 조금 어려운 내용이었습니다. 실제 구현부나 사용 예시를 봐야 제대로 된 이해를 할 수 있을 거 같은데요. 현재 외부 client가 사용할 수 있도록 공개되어 있진 않아서 사용 해보지는 못 한 것이 조금 아쉬웠지만 구글의 distributed ML system을 간접적으로나마 접할 수 있어서 재밌는 논문이었습니다.