구글은 자사의 연구와 서비스에 활용하기 위해 TPU를 비롯한 AI 인프라를 직접 설계해 왔습니다. 최근 오픈AI가 추론 비용 절감과 인프라 다변화를 목적으로 구글 TPU를 테스트하고 있다는 보도가 나오면서, TPU를 포함한 구글의 전체 AI 전략에도 시장의 관심이 높아지고 있습니다. (자료: 로이터)
구글은 오래전부터 소프트웨어와 하드웨어를 포함하는 수직 통합적 인공지능 플랫폼을 꾸준히 준비해왔습니다.
- 2015년: 머신러닝, 딥러닝 소프트웨어 프레임워크 텐서플로우(TensorFlow) 공개
- 2016년: AI 전용 하드웨어 가속기 TPU(Tensor Processing Unit) 발표
- 2018년: 고성능 수치 계산을 강화한 머신러닝 프레임워크 JAX 공개
그동안 인공지능 업계에서는 엔비디아(NVIDIA)의 GPU가 사실상 표준으로 활용되어 왔고, 구글의 TPU는 메일, 포토, 검색 등 자사 서비스에 한정적으로 적용되어 온 것으로 알려져 있습니다. 그러나 최근 구글은 기존 서비스뿐 아니라 최신 인공지능 제품군인 제미니(Gemini) 역시 TPU 기반으로 학습 및 구동하고 있다고 공개했습니다.
특히 제미니를 비롯한 구글의 최신 인공지능 모델들은 머신러닝 프레임워크 JAX와 TPU의 조합으로 학습되었으며, 이로 인해 시장은 구글의 수직 통합형 인공지능 전략에 관심을 기울이고 있습니다.
이번 글에서는 다음과 같은 흐름으로 구글의 기술 전략에 대해 살펴보겠습니다.
- 구글 TPU의 구조와 특징
- 딥러닝 프레임워크 이해를 위한 핵심 알고리즘 개요
- 구글의 딥러닝 프레임워크 텐서플로우
- 구글의 최신 머신러닝 프레임워크 JAX
이제부터 구글 TPU의 구조와 특징을 알아보겠습니다.
구글 TPU란?
구글 TPU의 정의와 특징
구글 TPU는 구글이 인공지능 모델의 학습과 추론을 위해 자체적으로 설계, 개발한 인공지능 가속기(AI Accelerator)입니다. 여기서 ‘가속기’란 AI 모델을 처리하는 과정에서 필요한 행렬 곱셈, 벡터 연산 등 대규모 연산을 병렬로 처리해 속도를 높이는 전용 하드웨어를 말합니다. 구글 TPU뿐 아니라, 우리가 익히 알고 있는 GPU 역시 인공지능 가속기의 한 종류입니다.
TPU는 Tensor Processing Unit의 약자로, 말 그대로 텐서(Tensor)를 처리하는 칩입니다. 텐서는 수학적으로 n차원의 다차원 배열을 의미하며, 프로그래밍 관점에서는 n차원 데이터를 표현하는 구조로 이해할 수 있습니다.
인공지능 분야, 그중에서도 딥러닝에서 다루는 대부분의 데이터 – 예를 들어 수치 데이터, 언어 데이터, 이미지, 비디오 등 – 는 모두 n 차원 텐서 형태로 표현됩니다. 결국, 딥러닝에서 사용하는 데이터는 구조화된 숫자 형태라는 의미인데요. 따라서 딥러닝 학습 및 추론 과정에서 발생하는 방대한 연산 역시 대부분 텐서 연산으로 표현되며, 구글 TPU는 이러한 텐서 연산을 빠르고 효율적으로 처리하는 데 최적화 된 칩입니다.
행렬 곱셈(Matrix Multiplication)
행렬 곱셈의 개념
딥러닝에서 수행되는 다양한 텐서 연산 중에서도 행렬 곱셉(Matrix Multiplication)은 가장 널리 사용되는 핵심 연산입니다.
행렬 곱셈은 두 입력 행렬을 곱해 새로운 결과 행렬을 만드는 연산으로, 예를 들어 5 x 3 행렬 A(왼쪽 초록)와 3 x 4 행렬 B(중간 분홍)를 곱하면 5 x 4 크기의 행렬 C(오른쪽 파랑)가 생성됩니다. (자료: 행렬 곱셈, Wikipedia)
행렬 곱셈의 활용
행렬 곱셈을 사용하는 예시로 원본 이미지를 픽셀화 하는 과정을 살펴보겠습니다. 보통 행렬 A를 입력(Input) 데이터, 행렬 B를 가중치(Weight) 또는 필터라고 부르며, 이 둘을 곱해 새로운 행렬 C를 생성하여 예측 결과를 얻습니다.
예를 들어, 입력 행렬 A를 아래 그림의 다운캣 로고 이미지로, 필터 행렬 B를 픽셀화 필터로 표현한다면, 행렬 A와 B를 곱해 생성한 새로운 행렬 C는 픽셀화된 다운캣 이미지로 나타납니다. (실제 픽셀화는 격자 분할, 평균 색상 계산 등 여러 단계를 거치지만, 이해를 돕기 위해 단순화하였습니다.)
그림: 다운캣 로고(좌), 다운캣 로고 픽셀화(우)
위와 같이 원본 데이터에 이미 알고 있는, 즉 기존에 학습된 필터를 적용하여 결과를 생성하는 과정을 딥러닝 추론이라고 합니다. 이러한 추론 과정에서는 모델의 크기가 커질수록 행렬 곱셈을 포함한 텐서 연산량이 큰 폭으로 급증하게 됩니다.
TPU 아키텍처
TPU의 행렬 곱셈 처리 방식
구글 TPU는 기본적으로 딥러닝에서 대규모 행렬 곱셈 연산을 효율적으로 처리하도록 설계된 전용 하드웨어 칩입니다.
TPU의 핵심 특징은 행렬 곱셈을 반복적으로 누적 수행하는 구조에 있습니다. 입력 행렬 A와 필터 행렬 B를 곱해 결과 행렬 C를 생성하고, 이 C를 다시 새로운 입력으로 활용해 다음 필터 행렬과 연산을 이어가는 방식입니다. 이러한 구조는 중간 연산 결과를 재사용할 수 있어 연산 속도와 효율성을 높이는 데 유리합니다.
시스톨릭 배열(Systolic Array) 구조
TPU의 연산 효율성을 가능하게 하는 핵심 아키텍처가 바로 시스톨릭 배열(Systolic Array)입니다.
TPU 내부에는 수천 개의 누산기(MAC: Multiply and Accumulate Unit)가 바둑판 격자 형태로 연결되어 있습니다. 이 구조를 통해 TPU는 마치 칩 위에 거대한 행렬을 물리적으로 구현한 것처럼 작동합니다.
시스톨릭 배열에서는 한 번 메모리에서 데이터를 불러온 후, 행렬 곱셈의 연산 결과를 인접한 누산기로 직접 전달해 다음 연산에 활용합니다. 이 방식에서는 중간 연산 결과를 재사용하기 때문에 메모리 접근을 최소화할 수 있습니다. 따라서, TPU는 대규모 행렬 연산을 더 빠른 속도와 더 낮은 전력으로 수행할 수 있습니다
딥러닝 학습 알고리즘
구글의 머신러닝 프레임워크인 텐서플로우와 JAX에 대해 알아보기 전에 딥러닝 모델의 학습 과정 원리를 살펴보고, 학습 과정에서 왜 방대한 연산이 필요한 지 간단하게 알아보겠습니다.
다운캣 픽셀화 예시로 보는 딥러닝 학습 원리
예를 들어, 다운캣 로고를 픽셀화한다고 가정해 보겠습니다.
입력 이미지 A에 픽셀화 필터 B를 적용해 결과 이미지 C를 만드는 과정은 단순한 행렬 연산입니다. 이 때 필터 B를 이미 알고 있다면, A x B = C 연산은 계산이 가능합니다.
하지만, 딥러닝 학습은 이와 반대의 과정을 수행합니다.
입력 이미지 A와 정답인 결과 이미지 C는 알고 있지만, 필터 B가 어떤 값으로 구성되어 있는 지 모르는 상태에서 알고리즘을 통해 최적의 필터 B(가중치)를 찾아야 합니다. 이렇게 최적의 필터 B를 찾는 과정을 학습(Training)이라고 하고, 학습한 필터 B를 새로운 입력 A’에 적용해 새로운 결과 C’를 예측하는 것을 추론(Inference)이라고 합니다.
학습 과정의 직관적 이해
딥러닝 학습은 입력 이미지 A와 목표로 하는 정답 이미지 C를 가지고 있는 상태에서, 최적화 된 필터 B를 찾는 작업입니다. 이를 간단히 수학적 개념으로 표현해 보겠습니다.
- 입력 이미지 A에 임의의 필터 B를 적용해 예측 이미지 C0를 만듭니다.
- 예측 이미지 C0와 원래의 정답 이미지 C의 차이를 계산합니다. 이 차이를 모델의 손실(Loss)이라고 합니다.
- 필터 B의 각 원소가 손실에 얼마나 영향을 미쳤는지 거꾸로 추적해 계산합니다. 이 과정을 기울기(Gradient) 계산이라고 하며, 기울기는 손실이 얼마나 빠르게 변화하는 지를 나타내는 값입니다.
- 계산된 기울기를 바탕으로 손실을 줄일 수 있는 방향으로 필터 B의 값을 조금씩 조정합니다.
위의 과정을 딥러닝 전문 용어로 바꾸면 다음과 같습니다
- 순전파 (Forward Propagation)
- 손실 계산 (Loss Calculation)
- 역전파 (Backpropagation)
- 경사하강법 (Gradient Descent)
딥러닝 학습 과정에서는 위와 같은 과정을 수천 ~ 수만 번 반복하면서 모델에 최적화 된 필터 B를 얻게 됩니다.
이 때 모델이 커질수록, 즉 필터 B에 해당하는 파라미터 개수가 많아질수록 연산량은 기하급수적으로 증가합니다. 특히 GPT와 같은 대규모 언어 모델(LLM)에서는 수십억 개의 파라미터가 존재하기 때문에, 학습 과정에서 필요한 연산량이 폭발적으로 증가하게 됩니다.
다음에는 구글의 머신러닝 프레임워크인 텐서플로우와 JAX에 대해 차례로 알아보겠습니다.
텐서플로우(TensorFlow)란? – 구글의 인공지능 소프트웨어 프레임워크
텐서플로우의 정의 및 특징
텐서플로우(TensorFlow)는 다양한 딥러닝 모델과 학습 알고리즘을 지원하는 종합 인공지능 프레임워크입니다.프레임워크라는 단어가 다소 생소할 수 있는데, 쉽게 말해 인공지능 모델을 설계, 학습, 추론하는 데 필요한 다양한 기능을 통합적으로 제공하는 개발 환경을 의미합니다.
텐서플로우는 이미지, 음성 처리에 사용되는 CNN, 시계열 데이터와 자연어 처리에 활용되는 RNN, 그리고 최신 대규모 언어 모델에서 널리 사용되는 Transformer 까지 다양한 딥러닝 모델을 지원합니다. 또한, 앞서 다룬 역전파, 경사하강법 등의 주요 학습 알고리즘을 손쉽게 적용할 수 있도록 제공합니다.
텐서플로우의 주요 특징은 다음과 같습니다.
- 다양한 딥러닝 모델 및 학습 알고리즘을 지원하며, 대규모 텐서 연산을 효율적으로 수행
- 전체 연산 과정을 계산 그래프(Computation Graph) 형태로 구조화하고 이를 최적화
- CPU, GPU, TPU 등 다양한 하드웨어 환경에서 최적화 된 방법으로 연산 자동 처리
최종적으로, 이렇게 생성된 최적화 코드를 각 하드웨어 가속기(GPU, TPU 등)에 전달해 효율적으로 연산을 수행할 수 있도록 지원합니다.
텐서플로우의 계산 그래프 (Computation Graph)
계산 그래프의 정의
텐서플로우는 연산 과정을 전체적으로 표현하고 효율적으로 관리하기 위해 계산 그래프(Computation Graph)라는 방식을 사용합니다.
텐서플로우의 계산 그래프는 모델 학습에 필요한 모든 연산 과정을 그래프 형태로 구조화한 것입니다.
- 노드(Node): 하나의 연산(Operation)을 의미 (예: 행렬 곱셈, 합성곱, 활성화 함수 등)
- 엣지(Edge): 노드 간 데이터 흐름을 의미 (연산 결과가 다음 연산으로 전달되는 경로)
즉, 모델 내부에서 이루어지는 복잡한 연산 과정을 지도처럼 시각화한 구조라고 이해할 수 있습니다.
계산 그래프의 장점
텐서플로우는 계산 그래프를 통해 모델 전체의 연산을 최적화하고, GPU, TPU와 같은 하드웨어 가속기의 병렬 처리 성능을 극대화합니다.
연산 최적화
- 불필요한 연산을 제거하고, 중복되는 연산을 병합하여 효율적으로 실행
- 연산 순서를 재배치해 전체 학습 속도를 높임
예를 들어, 행렬 A x 상수 b x 행렬 C라는 연산이 있다고 가정해 보겠습니다. 수학적으로 교환법칙이 성립하기 때문에 상수 b를 행렬 A 또는 행렬 C 중 어느 쪽과 먼저 곱하더라도 결과는 동일합니다.
하지만, 만약 행렬 A의 크기가 100만이고 행렬 C의 크기가 1만이라면 상황이 달라집니다.
- 상수 b x 행렬 A → 연산량 100만
- 상수 b x 행렬 C → 연산량 1만
즉, 상수 b를 어떤 행렬과 먼저 곱하느냐에 따라 연산량이 100배나 차이나게 됩니다.
계산 그래프는 이러한 연산 구조를 분석하여 전체적으로 연산량을 최적화할 수 있도록 구성됩니다.
병렬 처리 극대화
- 연산 간의 의존 관계를 분석해, 서로 독립적인 연산을 분리하여 병렬 처리 실행
- GPU, TPU와 같은 하드웨어 가속기에서 병렬성을 극대화
XLA를 활용한 하드웨어 수준의 최적화
텐서플로우는 XLA(Accelerated Linear Algebra)라는 전용 컴파일러를 제공하여, 계산 그래프를 하드웨어에 최적화 된 기계 코드로 변환합니다. 즉, 원래의 계산 그래프를 분석해, GPU, TPU 와 같은 가속기에서 가장 효율적으로 실행될 수 있도록 재구성하는 역할을 수행합니다.
- 최적화 된 기계 코드로 변환
- GPU, TPU 연산 코어와 메모리 간 불필요한 데이터 이동 최소화
- 여러 개의 작은 연산을 하나의 커널(하드웨어 연산 단위)로 병합하여 실행 → 효율 극대화
텐서플로우는 XLA를 활용하여 특히 TPU에서 연산 속도와 메모리 효율을 극대화할 수 있습니다.
JAX – 구글의 최신 머신러닝 프레임워크
JAX는 구글이 2018년 공개한 고성능 머신러닝 프레임워크로, 특히 자동 미분(Autograd: Automatic Differentiation, AD)을 중심으로 행렬 연산 기반의 수치 계산 성능을 강화한 것이 특징입니다. 자동 미분은 딥러닝에서 모델의 가중치를 학습하기 위해 기울기를 자동으로 계산하는 기술로, 앞서 설명한 역전파는 자동미분을 구현하는 한 가지 방법입니다.
또한, JAX는 기존 텐서플로우의 일부 한계를 보완하면서 행렬 연산을 최적화하고, 구글의 XLA 컴파일러를 활용하여 TPU 에서의 연산 속도를 극대화하는 장점을 가지고 있습니다. 이러한 이유로 JAX는 학계와 업계 모두에서 빠르게 주목받고 있습니다.
구글의 End-to-End 인공지능 전략
구글은 최신 제미니(Gemini) 제품군을 JAX + TPU 조합으로 구현했다고 공식 발표했습니다. (자료: 구글 클라우드) 제미니 모델의 뛰어난 성능이 입증되면서, JAX + TPU의 조합은 초거대 언어 모델의 학습과 추론 부문에서의 역량이 시장으로부터 인정받고 있습니다. 이에 따라 시장의 관심은 자연스럽게 구글의 인공지능 전략으로 집중되고 있습니다.
최근 구글은 자사 내부 제품과 서비스에 주로 사용해왔던 TPU를 외부 기업에도 공급하기 시작했습니다. 이는 오랫동안 엔비디아 GPU 체제가 지배해 온 AI 인프라 시장에서 구글 자체의 경쟁력을 강화하려는 행보로 보입니다.
기존에는 구글 클라우드 환경에서만 TPU를 제한적으로 사용할 수 있었지만, 최근 구글은 중소 클라우드 서비스 업체들과 협력하여 데이터 센터에 TPU를 설치하는 방안을 추진 중이라고 합니다. 또한, 지난 해 공개한 6세대 TPU인 트릴리엄(Trillim)에 이어, 올해 말 출시 예정인 대규모 추론 전용 TPU 아이언우도(Ironwood) 역시 업계의 큰 관심을 받고 있습니다. (자료: 구글 블로그)
과연 파이토치(PyTorch)와 엔비디아 GPU가 장악해 온 인공지능 시장에서, 구글이 JAX + TPU 중심의 End-to-End 전략으로 얼마나 영향력을 확대할 수 있을 지 주목됩니다.