Pathways란 무엇인가, 단일 컨트롤러로 대규모 모델 학습을 제어하는 구조 해설
What is Pathways? Single-controller training for large models
Google은 Gemini 학습에도 사용된 분산 시스템 Pathways의 Google Cloud 일반 제공 개시를 발표하고 그 구조를 해설했다. 하나의 클라이언트가 전체 연산을 그래프로 기술하면, 분산 런타임이 배치, 텐서 이동, 실행 순서를 담당한다. 호스트 장애 시에는 가상 디바이스를 정상 칩에 재할당해 학습을 이어갈 수 있다.
대규모 모델 학습의 대부분은 SPMD 방식으로 동작하며, 모든 가속기가 같은 프로그램을 록스텝으로 실행하고 AllReduce 등의 집합 통신으로 협조한다. 반면 파이프라인 병렬에서는 각 단계가 서로를 기다리고, RL 사후 학습에서는 샘플러, 보상 모델, 트레이너라는 서로 다른 프로그램을 겹쳐서 실행해야 한다. 이런 경우 단일 집합 통신으로는 의존 관계를 관리할 수 없어 협조 자체가 과제가 된다고 설명한다. 또한 SPMD에서는 머신 한 대만 잃어도 전체가 멈추고 가장 최근의 체크포인트부터 다시 시작해야 한다. Gemini 2.5의 학습은 1시간에 여러 차례 중단되었다고 한다.
Pathways에서는 일반 CPU 머신 위의 단일 Python 클라이언트가 가상 디바이스를 대상으로 프로그램을 작성하고, 리소스 매니저가 실제 칩에 할당하며, 각 가속기 호스트의 워커가 컴파일된 프로그램을 실행한다. 호스트가 고장 나면 정상 칩으로 작업을 재할당하므로 전체 작업을 재시작하지 않고 계속할 수 있으며, 이것이 elastic training의 기반이 된다. 물리 칩의 프로비저닝 이전부터 컴파일을 시작할 수 있는 점과, Shared Pathways Service를 통해 하나의 클러스터를 여러 클라이언트가 공유할 수 있는 점도 장점으로 꼽힌다.
실행 측면에서는 블록을 노드, 블록 사이의 텐서를 엣지로 하는 분산 데이터플로 그래프로 의존 관계를 표현한다. 런타임은 future와 병렬 비동기 디스패치를 사용해, 앞 단계의 출력이 나오기 전에 뒤 단계의 버퍼 확보와 전송을 준비한다. 중앙 스케줄러가 모든 프로그램의 연산을 일관된 순서로 배열하기 때문에 통신이 수반되는 연산이 교착 상태에 빠지지 않고 동시에 시작된다. 클라이언트가 하나이므로 셸에 들어가 브레이크포인트를 설정하는 등 단일 호스트와 같은 개발 흐름도 유지할 수 있다고 한다.
Ray와의 차이에 대해서는, Ray가 클러스터 프로비저닝과 태스크·액터 스케줄링 같은 오케스트레이션을 담당하는 반면 각 워커 내부의 학습은 여전히 SPMD로 실행된다고 설명한다. Pathways는 실행 모델 자체를 바꾸어 하나의 클라이언트가 모든 칩을 조율하는 계층에서 동작하므로, 두 기술은 상호 보완적이며 원리적으로는 함께 사용할 수 있다고 한다. Pathways는 JAX 프로그램을 TPU에서 실행하며 Google Cloud에서 이용할 수 있다. 문서, 논문, MaxText GitHub 리포지토리 등 관련 자료도 함께 소개되어 있다.
출처: Google Developers Blog · developers.googleblog.com