Ray 2.55 TPU 지원: 대규모 AI 워크로드를 GPU 전제에서 벗어나게 만드는 변화

Google 개발자 블로그에서 소개된 Ray 2.55의 TPU 지원은 단순한 프레임워크 업데이트가 아닙니다. 대규모 AI 워크로드를 돌리는 팀에게는 ‘GPU 클러스터를 어떻게 확보할 것인가’라는 전제에서 벗어나, TPU slice를 Ray의 task·actor·Serve·Train 흐름 안에서 다룰 수 있게 되는 변화입니다.

Ray는 Python 기반 분산 실행 프레임워크로, 모델 학습, 데이터 처리, 서빙, 에이전트 실행 파이프라인에서 널리 쓰입니다. TPU는 Google의 AI 가속기입니다. 둘은 원래 성격이 다릅니다. Ray는 유연한 분산 실행을 잘하고, TPU는 정해진 토폴로지와 고속 interconnect를 전제로 큰 행렬 연산을 빠르게 처리합니다. 이번 변화의 의미는 이 둘 사이의 운영 간극을 줄였다는 데 있습니다.

무엇이 달라졌나

핵심은 Ray 2.55가 Google Cloud TPU를 공식적으로 지원한다는 점입니다. 특히 KubeRay Operator on GKE가 TPU slice의 하드웨어 배치를 프로비저닝하고 라벨링하며, Ray Core는 slice_placement_group() 같은 primitive를 사용해 완전한 slice를 원자적으로 예약할 수 있습니다.

TPU는 일반 VM 몇 대를 아무렇게나 붙여 쓰는 방식과 다릅니다. 여러 호스트가 하나의 slice로 묶이고, Inter-Chip Interconnect 구조를 유지해야 성능이 나옵니다. 일부 노드만 잡거나 토폴로지를 깨면 작업이 느려지거나 아예 정상 실행이 어렵습니다. 그래서 TPU 운영에서 중요한 것은 ‘필요한 가속기를 몇 개 빌린다’가 아니라 ‘하나의 논리적 slice를 온전히 확보한다’입니다.

Ray가 이 부분을 추상화하면 개발자는 하드웨어 배치 세부사항을 매번 직접 다루지 않아도 됩니다. 작업에서는 4x4 같은 토폴로지를 선언하고, KubeRay와 Ray가 slice 배치와 예약을 맡습니다. 이 구조는 특히 실험이 잦은 ML 팀에게 유용합니다.

왜 실무 개발자에게 중요할까

첫째, TPU 도입 장벽이 낮아집니다. TPU는 성능이 좋아도 운영 난도가 높다는 인식이 있습니다. JAX나 XLA에 익숙한 팀은 잘 쓰지만, 일반 Python 분산 워크로드를 운영하던 팀은 GPU 쪽으로 기울기 쉽습니다. Ray 지원이 좋아지면 기존 Ray 코드와 운영 경험을 상당 부분 유지하면서 TPU를 검토할 수 있습니다.

둘째, 학습과 서빙의 경계가 줄어듭니다. 모델을 학습할 때는 Ray Train, 대량 데이터를 처리할 때는 Ray Data, 추론 서비스를 띄울 때는 Ray Serve를 쓸 수 있습니다. 같은 클러스터 운영 모델 안에서 데이터 로딩, 분산 학습, 온라인 서빙을 이어 붙이면 파이프라인 단절이 줄어듭니다.

셋째, 비용 최적화 선택지가 늘어납니다. GPU 수급이 불안정하거나 특정 워크로드가 TPU에서 더 잘 맞는 경우, 클라우드 비용과 처리량을 비교할 수 있습니다. GPU만 전제로 설계된 팀은 가격이 올라가도 대안이 적지만, TPU까지 선택지에 넣으면 워크로드별로 더 현실적인 계산이 가능합니다.

TPU에 맞는 워크로드와 맞지 않는 워크로드

TPU가 모든 문제를 해결하지는 않습니다. 잘 맞는 작업은 대규모 행렬 연산 비중이 높고, 배치 처리나 분산 학습 구조가 명확한 워크로드입니다. LLM 파인튜닝, 대규모 임베딩 생성, JAX 기반 학습, 일정한 패턴의 추론 처리처럼 하드웨어 활용률을 높이기 쉬운 작업이 여기에 들어갑니다.

반대로 작고 불규칙한 작업을 많이 던지는 구조는 TPU 효율이 낮을 수 있습니다. 예를 들어 요청마다 모델 크기, 입력 길이, 후처리 로직이 크게 달라지는 에이전트 서비스는 단순히 TPU로 옮긴다고 빨라지지 않습니다. 이런 경우에는 세션 관리, 큐잉, batching, 캐시 전략이 먼저입니다.

Ray의 장점은 이 중간 지대를 다룰 수 있다는 데 있습니다. 모든 것을 하나의 거대한 학습 job으로 만들지 않아도 task와 actor로 분산 실행을 구성할 수 있고, Serve로 서빙 레이어를 만들 수 있습니다. 다만 TPU에서는 placement와 데이터 이동이 성능을 크게 좌우하므로, Ray가 추상화해준다고 해서 측정이 필요 없어지는 것은 아닙니다.

도입 전에 확인해야 할 운영 포인트

가장 먼저 볼 것은 데이터 파이프라인입니다. TPU가 빠르더라도 데이터 로딩이 느리면 전체 처리량은 오르지 않습니다. Ray Data나 별도 데이터 로더가 TPU에 충분히 빠르게 배치를 공급할 수 있는지 봐야 합니다. Cloud Storage, 네트워크, 전처리 CPU 병목이 숨어 있는 경우가 많습니다.

두 번째는 체크포인트와 장애 복구입니다. TPU slice 단위로 큰 작업을 돌리면 실패 비용도 커집니다. 긴 학습 작업에서는 체크포인트 주기, 저장 위치, 재시작 절차를 정해야 합니다. Ray가 fault tolerance를 제공하더라도 모델 상태와 optimizer 상태를 어떻게 복원할지는 애플리케이션 설계 문제입니다.

세 번째는 관측성입니다. GPU 운영에서는 GPU utilization, memory, kernel time을 봅니다. TPU에서도 compute, memory bandwidth, network, host transfer 같은 지표를 봐야 합니다. Google이 별도로 TPU microbenchmark를 강조하는 이유도 여기 있습니다. 하드웨어를 바꿨는데 병목이 그대로이면 비용만 늘어납니다.

마이그레이션 접근법

처음부터 프로덕션 전체를 옮기면 위험합니다. 가장 좋은 시작점은 배치성 워크로드입니다. 예를 들어 야간 임베딩 생성, 대량 평가, 오프라인 파인튜닝처럼 사용자 트래픽과 분리된 작업을 먼저 TPU slice에서 실행해봅니다.

그 다음에는 동일한 작업을 GPU와 TPU에서 각각 돌려 비교합니다. 단순 실행 시간만 보면 안 됩니다. 총 비용, 실패율, 재시도 비용, 개발자 운영 시간, 데이터 이동 비용까지 합산해야 합니다. TPU가 30% 빠르더라도 운영 난도가 두 배라면 초기에는 GPU가 나을 수 있습니다. 반대로 안정화 후 처리량이 크게 오른다면 TPU가 장기적으로 유리합니다.

실행 체크리스트

• 현재 Ray 사용 여부와 분산 실행 구조를 확인한다.
• TPU에 적합한 배치성 워크로드를 하나 고른다.
• KubeRay on GKE와 TPU slice 예약 방식을 검토한다.
• 데이터 로딩 병목을 먼저 측정한다.
• GPU 대비 TPU의 처리량, 비용, 실패율을 같은 기준으로 비교한다.
• 체크포인트·재시작·로그 수집 절차를 문서화한다.
• 서빙으로 확장하기 전 오프라인 job에서 안정성을 확인한다.

Ray 2.55의 TPU 지원은 ‘TPU가 무조건 답’이라는 이야기가 아닙니다. 더 중요한 메시지는 AI 인프라 선택지가 넓어지고 있다는 점입니다. GPU만 전제로 아키텍처를 짜면 비용과 수급에 묶입니다. Ray와 TPU 조합은 대규모 AI 워크로드를 더 유연하게 배치할 수 있는 현실적인 대안으로 봐야 합니다.