구글 AI 데이터센터는 어떻게 생겼을까?
바야흐로 AI 시대가 열렸습니다. AlphaGo, ChatGPT 등 세상을 놀라게 하는 AI 가 계속해서 개발되고 있습니다. 벤더들은 AI 시대에 패권을 잡기 위해 AI 데이터센터 구축에 박차를 가하고 있습니다. 구글도 마찬가지 입니다. 이번 글에서는 구글은 어떻게 AI 데이터센터를 구축하고 있는지 엿볼 수 있는 논문인 Resiliency at Scale: Managing Google’s TPUv4 Machine Learning Supercomputer (NSDI’24) 를 리뷰합니다.
Large Language Model (LLM) 의 등장으로 Machine Learning (ML) 모델의 크기와 복잡도는 계속해서 증가하고 있습니다. 더 큰 LLM 을 학습하기 위해서는 AI 슈퍼컴퓨터를 확장해야 합니다. AI 슈퍼컴퓨터를 확장하는 것에는 여러가지 어려운 점이 있는데 그 중 하나는 컴퓨팅 자원의 고가용성을 유지하는 것입니다. 그 이유는 ML 학습이 Map-Reduce 같은 기존 슈퍼컴퓨팅 워크로드와 다른 특성을 갖기 때문입니다. 기존 워크로드는 결합도가 낮은 작은 작업으로 분산시켜서 수행했기 때문에 자원 할당이나 실패 처리에 용이했습니다. 반면 ML 학습은 주로 컴파일 타임에 정적으로 자원을 할당하고 모든 컴퓨팅 자원이 실패없이 동작해야 계속해서 작업을 수행할 수 있습니다. 정적으로 할당된 자원중 하나라도 실패하면 ML 학습이 중단됩니다. 그래서 LLM 의 경우 Mean Time Between Failures (MTBF) 가 시간 혹은 분 단위까지 내려갑니다. 또한 모델의 크기나 형태가 다양하기 때문에 사용자는 다양한 설정으로 자원을 할당 받으려 합니다.
구글은 이런 문제를 해결할 수 있도록 AI 슈퍼컴퓨터를 디자인 했습니다.
구글의 AI 슈퍼컴퓨터는 Tensor Processing Unit (TPUv4), TPU 끼리 이어주는 Inter-Chip Interconnect (ICI), TPU 를 동적으로 연결하는 Optical Circuit Switches (OCS), TPU 를 스케줄링하는 Borg, TPU 연결을 관리하는 Pod Manager, ICI 네트워크 연결을 수행하는 libtpunet, 머신 상태를 체크하기 위한 healthd 로 구성되어 있습니다. 하드웨어(TPU, ICI, OCS)와 소프트웨어(Borg, Pod Manager, libtpunet, healthd)가 같이 디자인 되었습니다.
정적 구성만 가능한 TPUv3 의 한계
구글 AI 슈퍼컴퓨터 가용성의 핵심은 programmable OCS 에 기반한 reconfigurable ICI fabric topology 입니다. 동적으로 네트워크 형상을 재구성할 수 없다면 많은 리소스를 사용할수록 가용성이 급격히 떨어지게 됩니다. 장비 고장 및 점검, 다양한 워크로드 구성은 정적 구성 TPUv3 의 가용성을 떨어트리게 됩니다. Figure 1 은 OCS 가 없기 때문에 정적 구성만 할 수 있는 TPUv3 의 가용성이 학습 크기(필요 TPU 수)에 따라 떨어지는 것을 보여줍니다. 반면 TPUv4 는 3000 개가 넘도록 가용성을 94% 이상 유지하는 것을 볼 수 있습니다.
동적 재구성 가능한 OCS 기반의 TPUv4
TPUv4 는 reconfigurability 를 위해 OCS 를 사용합니다. OCS 는 동적으로 configure 가능하며, 프로그램 가능한 cross-connect (xconnect) 로 서로 다른 TPU machine 의 포트를 연결할 수 있습니다. 두 포트가 xconnect 되면 ICI 가 연결되고 두 포트가 통신할 수 있도록 라우팅됩니다.
1 TPUv4 Supercomputer = 64 cubes = 64×16 machines = 64×16×4 chips
TPU machine 은 4개의 TPU chip (2×2×1) 으로 구성되어 있습니다. 16개의 TPU machine (2×2×4) 가 하나의 rack 에 들어가고, 이렇게 구성된 4×4×4 TPU chip 은 TPU cube 라고 부릅니다. TPU cube 의 한 면 당 16개의 ICI 가 연결되어있어서 총 96개의 ICI 가 연결됩니다. TPUv4 슈퍼컴퓨터는 64 개의 cube 와 48 개의 OCS 로 구성됩니다. 이렇게 OCS 로 연결된 TPUv4 는 reconfigurable 하기 때문에 TPU cube 가 물리적으로 인접하지 않더라도 인접한 노드처럼 학습에 사용될 수 있습니다.
TPUv4 의 ICI 프로토콜은 프로그래밍 가능하며 Transaction, Routing, Data, Physical 4개의 layer 로 나뉘어있고 소프트웨어에 의해 제어됩니다. Table 1 에서 각 layer 의 기능과 layer 를 컨트롤하는 소프트웨어를 보여줍니다. 소프트웨어로 프로그램 가능한 ICI 프로토콜은 4096 TPU chip 의 네트워크 형상을 동적으로 재구성할 수 있게 해줍니다.
- Request: TPUv4 슈퍼컴퓨터는 3D 네트워크 형상(torus, twisted-torus)을 채택하고 있습니다. 따라서 사용자는 (4x, 4y, 4z) 형식의 네트워크 형상을 명시해서 training job 을 요청해야 합니다.
- Scheduling: Borg 는 이런 요청을 받아서 스케쥴링합니다. Borg 는 이 형상으로 재구성 가능한 cube 를 찾아서 xconnect 요청을 발행합니다. Pod Manager 는 주기적으로 xconnect 요청을 polling 합니다. xconnect 요청이 발생하면 OCS 의 MEMS mirror 를 회전시켜서 물리적인 ICI 채널을 수립합니다. xconnect 가 성공적으로 완료되면 Pod Manager 는 Borg 에 확인 메시지를 전달합니다. 그러면 Borg 는 job 실행파일을 선택된 TPU machine 에 보내고, 각 TPU machine 의 borglet daemon 은 container 를 만들고 그 안에서 학습을 실행합니다.
- libtpunet: ICI 채널이 수립되면 libtpunet 이 실행됩니다. libtpunet 은 먼저 요청된 구성되로 TPU 가 연결되었는지 확인하기 위해 topology discovery 를 수행합니다. 그리고 각 TPU 의 link-level flow-control (data layer) 설정과 packet forwarding table (routing layer) 을 프로그램합니다. 마지막으로 ICI 세션을 시작하고 link enabling 을 합니다. 이 과정에서 에러가 발생하면 Borg 에 알리고 다시 스케쥴링 되도록 합니다.
- Preflight check: 학습 실행에 앞서 preflight 체크를 수행합니다. 두 가지 체크를 하는데, 먼저 end-to-end 체크는 작은 sample 워크로드를 실행해서 TPU 가 정상 상태인지 확인합니다. 그리고 intent-driven 체크는 여러가지 하드웨어 성능 지표를 확인하여 상태를 확인합니다. preflight 에 실패하면 Borg 에 알리고 다시 스케쥴링 되도록 합니다.
- healthd: 모든 TPU machine 은 healthd daemon 이 실행됩니다. healthd 는 하드웨어 상태를 실시간 모니터링합니다. ICI link, PCIe channel, TPU chip 등이 모니터링됩니다. 이상 징후가 관측되면 Borg 에 알리고 다시 스케쥴링 되도록 합니다.
Fault-tolerant ICI Routing
일반 torus 형상으로 설정되면 ICI 는 Dimension-order Routing (DOR) 을 사용합니다. DOR 은 최단경로를 따라 한 축으로 한번씩만 route 합니다. (X 축으로 먼저 이동하고 Y 축으로 이동하면 XY 로 표기합니다.) 그런데 중간 link 가 끊어지면 DOR 을 사용할 수 없습니다. 이런 경우 Wild-first Routing (WFR) 을 사용합니다. WFR 은 각 축마다 최대 한 번의 wild hop 을 허용합니다. 예를들어 Figure 11에서 XY 순서로 DOR 이 불가능하면(path A) y 축으로 wild hop 을 해서 (path B & C) route 하게됩니다. (yXY로 표기합니다.) WFR 은 deadlock-free 를 보장하기 위해 한가지 제약이 있습니다. Wild hop 순서는 DOR 의 축 순서의 반대여야 합니다. 예를들어 xyZYX 는 deadlock-free 하지만 yxZYX 는 deadlock-free 하지 않습니다.
구글 TPUv4 슈퍼컴퓨터 운영 통계
구글 TPUv4 슈퍼컴퓨터는 하루에 수천 건의 학습이 요청됩니다. Figure 13 은 두 달 동안 요청된 학습과 xconnect 변경 건 수의 샘플입니다. 학습이 요청되면 OCS 는 학습에 맞게 xconnect 를 변경하여 요청된 topology 를 만들기 때문에 학습 요청과 xconnect 가 비례하는 것을 볼 수 있습니다.
장애도 발생합니다. 하루 평균 TPU machine 의 0.08%, ICI cable 의 0.005%, OCS 의 0.04% 가 장애를 겪습니다. 하지만 TPUv4 슈퍼컴퓨터는 reconfigurable 한 TPU 와 fault-tolerant 한 ICI routing 이 있기 때문에 장애를 견딜 수 있습니다.
Fault-tolerant ICI routing 은 장애 link 주변에 더 많은 트래픽을 발생 시키기 때문에 성능 저하를 야기합니다. Table 3 은 Recommendation Model (RM), Large Language Model (LLM), BERT 모델의 성능 저하를 보여줍니다.
이 논문을 통해 재구성 가능한 TPUv4 슈퍼컴퓨터의 특성 덕분에 구글의 AI 데이터센터는 높은 가용성을 유지하며 운영되고 있다는 것을 알게됐습니다. 기회가 된다면 TPU, OCS, Borg 등 각 구성요소들을 좀 더 깊게 알아보면 재밌을 것 같습니다.
