Google이 내부에서 쓰던 대규모 AI 학습 기술, 이제 GKE에서? – Pathways 완전 정리

 

복잡한 TPU 분산 학습, 더 이상 어렵지 않아야 하지 않을까요?

AI 모델이 커지면서 학습에 필요한 연산량도 폭발적으로 늘어나고 있습니다. 특히 TPUs를 활용한 대규모 분산 학습 환경에서는 “어떻게 효율적으로 리소스를 나누고 통제할 것인가?”가 핵심 과제로 떠오르고 있습니다.

Google은 이러한 문제를 해결하기 위해 Pathways라는 내부 기술을 활용해왔고, 지금은 이 기술을 Google Cloud 환경에서도 사용할 수 있게 했습니다. 기존 JAX 기반의 멀티 컨트롤러 방식보다 훨씬 간단하고 확장성 있는 방식으로, GKE(Google Kubernetes Engine) 환경에서 수천 개의 TPU를 유연하게 관리하고 훈련에 활용할 수 있습니다.

이 글에서는 Pathways가 어떤 기술인지, 기존 방식과 무엇이 다른지, 그리고 실제로 어떻게 GKE에서 활용할 수 있는지를 상세히 정리해 드립니다.

반응형

---

🔎 Pathways란 무엇인가?

Pathways는 Google이 개발한 대규모 분산 머신러닝 시스템입니다. 핵심 목적은 간단합니다:

“수천 개의 TPU를 유연하게 조합해, 하나의 거대한 머신처럼 활용할 수 있도록 하자.”

기존 방식에서는 작업마다 고정된 TPU 수를 할당하고 동일한 연산을 실행해야 했지만, Pathways는 작업의 복잡도나 필요에 따라 TPU 리소스를 동적으로 조절할 수 있습니다. 이 덕분에 복잡한 분산 훈련 패턴(SPMD를 넘는 비대칭 파이프라인, 희소 연산 등) 도 가능해졌습니다.

또한 Google은 이 Pathways 기술을 내부에서 Gemini 모델 같은 대형 모델 훈련에 실제로 사용하고 있으며, 이제 이 기능을 Google Cloud 사용자도 사용할 수 있게 되었습니다.

---

🧱 주요 아키텍처 구성 요소

Pathways는 여러 컴포넌트가 협업하여 작동합니다. 각 구성 요소는 역할이 명확히 나눠져 있으며, TPU 연산을 위한 전체 흐름을 아래와 같이 구성합니다.

1. Pathways Resource Manager

• 기능: 전체 리소스를 통제하는 중앙 컨트롤러
• 리소스 요구: CPU만 사용
• 역할:
  • TPU 할당 및 모니터링
  • 작업 스케줄링 및 상태 관리
  • 시스템 에러 감지 및 리포트

2. Pathways Client

• 기능: 사용자 프로그램이 시스템에 명령을 내리는 진입점
• 리소스 요구: CPU만 사용
• 역할:
  • 사용자 코드에서 나오는 HLO(High Level Operations)를 수신
  • 실행 위치 결정 및 Pathways와의 인터페이스 수행

3. Pathways Worker

• 기능: 실질적으로 TPU에서 연산을 수행하는 프로세스
• 리소스 요구: TPU 필요
• 역할:
  • 컴파일된 프로그램을 실행
  • IFRT 서버와 통신하며 연산 결과 반환

4. IFRT Proxy Server / Client

• 기능: 사용자 코드와 Pathways 사이의 통신 브리지
• 리소스 요구: CPU
• 역할:
  • 프록시 서버는 클라이언트로부터 요청을 받아 Pathways로 전달
  • 사용자 코드가 내부 시스템을 몰라도 실행 가능하도록 함

5. Sidecar Server

• 기능: latency 최소화를 위한 보조 서버
• 역할:
  • Pathways Worker와 동일한 VM에 위치
  • Python 코드를 Accelerator VM에서 직접 실행

---

🆚 기존 방식과 무엇이 다른가? (Single Controller vs Multi Controller)

항목	기존 JAX 방식 (Multi-controller)	Pathways (Single-controller)
제어 방식	각기 다른 Python 인스턴스	하나의 클라이언트에서 전체 제어
시야	각 프로세스가 제한된 장치만 인식	전체 장치를 하나의 시스템처럼 인식
프로그래밍	SPMD 기반 동일 프로그램 실행	다양한 연산 방식 및 구조 지원
유연성	상대적으로 제한적	복잡한 분산 구조에도 유연하게 대응 가능

결론적으로 Pathways는 마치 모든 TPU를 하나의 슈퍼머신처럼 사용하는 경험을 제공합니다.

---

⚙️ GKE에서 Pathways 사용하기

Pathways는 GKE 환경에 컨테이너 형태로 배포됩니다. 직접 YAML을 다루기보다는 PathwaysJob API를 통해 관리하는 구조입니다. 다음은 대표적인 구성 예시입니다.

1. Resource Manager 컨테이너 예시

- name: pathways-rm
  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
  args:
    - --server_port=29001
    - --node_type=resource_manager
    - --instance_count=WORKLOAD_NODEPOOL_COUNT
    - --instance_type=tpuv5e:4x4
    - --gcs_scratch_location=gs://BUCKET_NAME

2. IFRT Proxy Server 예시

- name: pathways-proxy
  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
  args:
    - --resource_manager_address=$(PATHWAYS_HEAD):29001
    - --server_port=29000
    - --gcs_scratch_location=gs://BUCKET_NAME

3. Pathways Worker 예시

- name: worker
  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
  args:
    - --server_port=29001
    - --resource_manager_address=$(PATHWAYS_HEAD):29001
    - --gcs_scratch_location=gs://BUCKET_NAME
  resources:
    limits:
      google.com/tpu: "4"

이처럼 각 구성 요소는 명확한 역할을 가지고 컨테이너화되어 있으며, GKE 상에서 서로 통신하며 작업을 수행합니다.

---

🌟 Pathways의 장점 요약

• ✅ 중앙 집중 제어: 복잡한 분산 시스템을 단일 클라이언트에서 제어
• ✅ 연산 유연성: SPMD를 넘는 다양한 연산 패턴 가능
• ✅ 초대형 TPU 지원: 수천 개의 TPU를 마치 하나의 머신처럼 활용
• ✅ GKE 완전 통합: 컨테이너 기반 배포로 클라우드 인프라에 쉽게 접목
• ✅ 코드 이식성 강화: IFRT API를 통한 유연한 실행 구조

---

728x90

Pathways는 Google이 직접 내부에서 대형 AI 모델을 훈련할 때 사용하던 시스템을 Google Cloud 사용자에게 개방한 것입니다. 기존 JAX 환경에서 멀티 컨트롤러 방식으로 분산 TPU를 관리하던 방식은 설정이 복잡하고 유연성이 부족했지만, Pathways는 이를 더 간단하고 확장 가능하게 바꿔줍니다.

AI 인프라의 확장을 고민하고 있다면, Pathways는 분명 한 번은 검토해볼 가치가 있는 솔루션입니다. 특히 GKE 환경에서 수천 개의 TPU를 효율적으로 운영하고 싶은 팀이라면, 지금이 바로 그 시작점입니다.

https://cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro?fbclid=IwY2xjawJj9BZleHRuA2FlbQIxMAABHsvuXG9yJSGxrcfbBQXFe12EzdvUX_wrvLIoxe_P9yyc7C0vOlEvBkpoh68s_aem_9EaUuIQoJQaTCfa0h914Tw

Introduction to Pathways on Cloud  |  AI Hypercomputer  |  Google Cloud