🤖 AI

Fine-Tuning and Serving Gemma 4 31B on Google Cloud TPU: A Technical Comparison with GPU Baselines

본 논문은 구글 클라우드 TPU 에서 Gemma 4 31B 의 파인튜닝과 서빙을 처음 종단간으로 구현한 사례를 제시하며, GPU 네이티브 워크플로우를 JAX 스택으로 포팅하기 위해 필요한 코드 적응 사항을 상세히 설명하고, H100 GPU 베이스라인과 비교해 추론 처리량은 동등하게 유지하면서 지연 시간은 크게 줄이고 1.61 배 빠른 학습 속도와 1.82 배 낮은 비용 달성을 위한 TPU 구성의 성과를 입증합니다.

원저자: Jatin Kishnani, Mayank Goel, Amit Singh, Pulkit Agrawal, Sairanjan Mishra

게시일 2026-05-26

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Jatin Kishnani, Mayank Goel, Amit Singh, Pulkit Agrawal, Sairanjan Mishra

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 전자 회로 (Verilog) 를 위한 컴퓨터 코드를 작성하는 새로운 기술을 가르치려고 하는 천재적이지만 매우 비싼 로봇 (AI 모델) 을 상상해 보십시오. 이 훈련이 이루어지는 "학교"에는 두 가지 옵션이 있습니다:

GPU 학교: 엔비디아 (NVIDIA) 가 운영하며, 유명한 H100 칩을 사용합니다. 이는 많은 경험 많은 교사들이 있는 표준적이고 널리 알려진 길입니다.
TPU 학교: 구글 (Google) 이 운영하며, 자체 TPU 칩을 사용합니다. 이는 더 빠르고 저렴하지만, 입학하려면 완전히 새로운 언어를 배워야 하는 새롭고 매우 전문화된 캠퍼스입니다.

h2loop.ai의 이 보고서는 구글 TPU 학교를 시도하기로 결정한 엔지니어들이 작성한 "현장 가이드"입니다. 그들은 310 억 파라미터 AI 모델 (Gemma 4) 을 가져와 코드를 작성하도록 가르친 후, 표준 NVIDIA 학교와 비교하여 얼마나 잘 작동하는지 평가했습니다.

다음은 그들의 여정을 간단히 설명한 내용입니다:

1. 언어 장벽 (이식 문제)

가장 큰 장애물은 하드웨어 자체가 아니라 소프트웨어였습니다.

GPU 학교는 대부분의 AI 개발자가 이미 알고 있는 인기 있는 프로그래밍 언어인 PyTorch를 사용합니다.
TPU 학교는 다른 사고방식을 요구하는 다른 언어인 JAX를 사용합니다.

저자들은 전체 "수업 계획"을 PyTorch 에서 JAX 로 번역해야 했습니다. 그들은 다음을 수행해야 했습니다:

여러 칩에 걸쳐 분할된 모델의 뇌를 재배치하기 (2 개가 아닌 4 개의 방에 책이 흩어져 있는 도서관을 정리하는 것과 같은 방식).
TPU 학교가 다른 용어를 사용하기 때문에 모델의 특정 부분을 이름 변경하기 (예: "q_proj"를 "q_einsum"으로 부름).
TPU 학교는 세계 나머지 부분 (Safetensors) 과 다른 파일 시스템 (Orbax) 을 사용하므로 작업을 저장하기 위한 맞춤형 "다리"를 구축하기.

비유: 집을 옮기는 것과 같습니다. 가구 (AI 모델) 는 동일하지만, 새로운 집 (TPU) 은 다른 출입구와 평면도를 가지고 있습니다. 가구를 분해하여 새로운 문을 통해 운반하고 다시 조립하지 않으면 들어맞지 않습니다.

2. 훈련 레이스 (로봇 가르치기)

모델을 설정한 후, 그들은 훈련 레이스를 시작했습니다.

속도: TPU 학교는 1.6 배 더 빨랐습니다. 훈련 과정을 3.3 시간 만에 완료한 반면, GPU 학교는 5.4 시간이 걸렸습니다.
비용: TPU 학교는 시간당 요금이 더 저렴할 뿐만 아니라 더 빠르게 완료하므로 총 비용은 2.1 배 더 저렴했습니다.
- GPU 비용: 약 $119
- TPU 비용: 약 $56

왜 TPU 가 더 빨랐을까요?
TPU 칩을 데이터 (봉투) 를 서로 매우 빠르게 전달하는 초고속 내부 고속도로 (ICI) 를 통해 서로 연결된 4 명의 주자 팀으로 생각하십시오. GPU 학교에는 2 명의 주자만 있었습니다. 개별 GPU 주자가 약간 더 빨랐지만, TPU 팀의 데이터 조정 및 이동 능력 덕분에 전체 팀이 승리했습니다.

3. 시험 (로봇이 배웠는가?)

훈련 후, 그들은 로봇들을 Verilog 코딩 시험으로 테스트했습니다.

결과: GPU 에서 훈련된 로봇이 가장 어려운 질문에서 약간 더 높은 점수 (약 5% 더 좋음) 를 받았습니다.
주의할 점: 저자들은 이 차이가 한 로봇이 본질적으로 더 똑똑해서가 아니라 두 학교가 시험을 약간 다르게 채점했기 때문일 가능성이 높다고 지적했습니다 (하나는 질문의 "프롬프트" 부분을 무시했고 다른 하나는 무시하지 않았습니다). 이 차이는 통계적으로 크지 않았습니다.

4. 취업 면접 (모델 서비스)

훈련 후, 모델은 실시간으로 사용자의 질문에 답하도록 배치되어야 합니다. 이를 "추론"이라고 합니다. 그들은 모델이 다양한 길이의 질문에 얼마나 빠르게 답할 수 있는지 테스트했습니다.

짧은 질문 (간단한 작업):
- GPU 학교는 매우 짧은 질문에 답하는 데 약간 더 빨랐습니다.
- 또한 이러한 빠른 작업에 대한 답변당 비용도 약간 더 저렴했습니다.
- 비유: 커피 한 잔만 사야 한다면, 지역 상점 (GPU) 이 약간 더 효율적입니다.
긴 질문 (복잡한 작업):
- 질문이 길어지면 (4,000 단어 이상), TPU 학교가 압도적으로 우위를 점했습니다.
- 속도: TPU 는 긴 질문에 대한 답변을 시작하는 속도 (Time-to-First-Token) 가 23 배 더 빨랐습니다.
- 용량: TPU 는 속도가 느려지지 않으면서 동시에 4 배 더 많은 사람들이 긴 질문을 할 수 있었습니다.
- 비유: 전체 소설을 써야 한다면, TPU 학교는 매끄럽게 협력할 수 있는 팀과 거대한 도서관을 가지고 있습니다. GPU 학교는 압도되어 책을 떨어뜨리기 시작합니다.

5. 최종 판결

저자들은 특정 요구 사항 (대규모 모델 훈련 및 사용자에게 서비스 제공) 에 대해 구글의 TPU 가 승리자라고 결론지었습니다.

훈련: TPU 는 명확한 챔피언입니다 (더 빠르고 훨씬 저렴함).
추론: TPU 는 복잡하거나 긴 모든 작업의 챔피언입니다. 매우 간단하고 짧은 작업의 경우 GPU 가 여전히 약간 더 좋습니다.
총 비용: 훈련 비용과 하루 동안 사용자에게 서비스를 제공하는 비용을 합치면, TPU 설정은 전체적으로 1.8 배 더 저렴했습니다.

핵심 요약:
TPU 학교로 전환하려면 많은 초기 노력이 필요했습니다 (새로운 언어 학습, 가구 수리, 다리 건설). 하지만 정착한 후에는 학교가 더 빠르게 운영되었고, 비용이 더 적게 들었으며, 전통적인 GPU 학교보다 훨씬 더 크고 복잡한 작업을 처리했습니다. 중대한 AI 작업을 수행하는 회사에게는 전환을 위한 추가 노력이 비용 절감과 성능 향상으로 충분히 가치가 있었습니다.

기술 요약: Gemma 4 31B LoRA SFT 를 위한 TPU 대 GPU

문제 정의

대형 언어 모델 (LLM) 의 미세 조정 (fine-tuning) 과 서비스 제공은 데이터 가용성보다는 연산 비용과 반복 속도에 의해 increasingly 제약받고 있습니다. 구글 클라우드 TPU 는 높은 이론적 성능을 제공하지만, 성숙한 PyTorch 기반 GPU 생태계에 비해 상당한 엔지니어링 노력이 필요하다는 인식이 지배적입니다. 본 보고서는 Gemma 4 31B를 위한 표준 GPU 네이티브 학습 레시피 (PyTorch, HuggingFace TRL, FSDP) 를 JAX 기반 TPU 스택 (Tunix, Qwix, Orbax) 으로 포팅하는 엔드 - 투 - 엔드 프로세스를 문서화함으로써 오픈 툴링의 격차를 해소합니다. 본 연구는 Verilog 코드 생성을 위한 **LoRA 감독 미세 조정 (SFT)**과 후속 추론 서비스 제공을 구체적으로 타겟팅하여, 동일한 하이퍼파라미터 하에서 TPU 와 GPU 하드웨어 간의 경험적 비교를 제공합니다.

방법론

저자들은 두 가지 다른 하드웨어 구성에서 Gemma 4 31B 의 미세 조정 및 서비스 제공을 병렬 비교했습니다:

학습: TPU v5p-8(4 칩) 대 2× NVIDIA H100 80GB GPU.
추론: TPU v6e-8(Trillium, 8 칩) 대 2× NVIDIA H100 80GB GPU.
데이터셋: 학습용 CodeV-R1(10k 자연어 -Verilog 쌍); 평가용 NVlabs/verilog-eval(156 명세 -RTL 문제).

엔지니어링 적응 (GPU 에서 TPU 로)

방법론의 핵심은 PyTorch 에서 JAX 로 학습 및 추론 파이프라인을 포팅하기 위한 종합적인 작업이었습니다:

프레임워크 및 메시 구성: PyTorch FSDP 에서 명시적 장치 메시 구성 (fsdp=1, tp=4, v5p-8 기준) 을 갖춘 JAX 로 전환했습니다. 저자들은 Gemma 4 의 어텐션 헤드 수가 8 로 나누어지지 않아 tp=8이 실패했다고 지적했습니다.
LoRA 주입 및 명명: HuggingFace(PyTorch) 와 Tunix(JAX) 간의 모듈 명명 불일치를 해결했습니다. 특히 Tunix 는 K 와 V 프로젝션을 단일 kv_einsum 텐서로 융합하므로, 체크포인트 병합 시 커스텀 처리가 필요했습니다.
샤딩 및 최적화: Qwix 가 주입한 파라미터가 기본 가중치와 일치하지 않는 PartitionSpec 랭크를 상속받는 LoRA 샤딩 주석에 대한 수정을 구현했습니다. 옵티마이저를 Optax AdamW로 전환하여 커스텀 상태 샤딩 로직이 필요했습니다.
데이터 파이프라인: HuggingFace DataLoader 를 구글의 Grain 라이브러리로 교체하여 손실 마스킹 (시스템 프롬프트 및 추론 체인 제거) 과 시퀀스 처리를 위한 커스텀 로직을 구현했습니다.
체크포인팅 및 변환: 텐서 모양 전치와 융합된 kv_einsum 분할을 고려하여 Orbax(JAX 형식) 에서의 LoRA 어댑터를 추론을 위해 Safetensors(PyTorch/HF 형식) 로 병합하는 커스텀 파이프라인을 개발했습니다.
추론 설정: 특정 Docker 이미지 (vllm/vllm-tpu:gemma4) 를 사용하여 v6e-8 에 vLLM-TPU를 배포했습니다. 시작 시 XLA 컴파일, HBM 할당, 그리고 Gemma 4 의 이질적인 어텐션 차원으로 인한 분할된 멀티모달 입력 비활성화와 같은 제약 사항을 처리했습니다.

주요 결과

학습 성능

속도: TPU v5p-8 은 2×H100 기준 (5.39 시간) 보다 1.61 배 빠른 (3.34 시간) 학습을 완료했습니다.
처리량: TPU 는 763 토큰/초를 달성한 반면 GPU 는 486 토큰/초였습니다.
비용: 시간당 요금 ( $16.80 대$ 22.12) 이 낮고 실행 속도가 빨라 TPU 학습 비용은 2.12 배 저렴했습니다 ( $56.11 대$ 119.23).
손실 수렴: 두 모델 모두 유사한 수렴 경로를 따랐습니다. 절대 손실 차이 (TPU 0.072 대 GPU 2.258) 는 다른 손실 계산 범위 (TPU 는 어시스턴트 전용 손실 사용, GPU 는 전체 시퀀스 손실 사용) 로 인해 발생했습니다.

추론 성능

짧은 컨텍스트 (≤2048 토큰): GPU 는 피크 출력 처리량에서 +6% 의 미미한 우위를 보였으며, 칩당 HBM 대역폭이 더 높고 XLA 버킷 패딩 오버헤드가 부재하여 P99 꼬리 지연 시간이 약간 더 좋았습니다.
긴 컨텍스트 (≥4096 토큰): TPU v6e-8 은 우세함을 입증했습니다:
- 처리량: 4096 토큰에서 출력 토큰/초가 66% 더 높았습니다.
- 첫 번째 토큰 도달 시간 (TTFT): 4096 토큰에서 23.6 배 빠르며 (61ms 대 1,443ms), 8192 토큰에서는 7.1 배 빠릅니다.
- 확장성: TPU 는 fp8 KV 캐시(실제 용량 두 배) 와 더 높은 집계 연산 능력 덕분에 중간 컨텍스트에서 최대 **32 개의 동시 요청 (QPS)**을 유지한 반면, GPU 는 약 8 개에서 포화 상태에 도달했습니다.
비용: 긴 컨텍스트 워크로드의 경우 TPU 는 출력 토큰당 41% 더 저렴했습니다.

평가 품질

Verilog 생성 벤치마크에서 GPU 학습 모델은 **pass@1(0.697 대 0.641)**과 **pass@5(0.814 대 0.795)**에서 약간 더 높은 성능을 달성했습니다. 저자들은 이 미미한 차이가 하드웨어 능력보다는 학습 목표 변이 (손실 마스킹) 와 확률적 샘플링 변동에 기인한다고 설명합니다.

중요성과 주장

본 논문은 오픈 툴링 생태계의 중요한 격차를 효과적으로 해소하며, TPU 하드웨어에서 Gemma 4 31B 의 미세 조정 및 서비스 제공을 보여주는 최초의 엔드 - 투 - 엔드 시연을 제공한다고 주장합니다.

생산 준비 상태: 본 작업은 특정 Docker 구성, 메시 설정, 체크포인트 변환 스크립트를 포함하여 TPU 인프라에 Gemma 4 를 배포할 수 있는 재현 가능하고 생산 준비가 된 레시피를 제공합니다.
총 소유 비용 (TCO): 본 연구는 대표적인 학습 및 서비스 제공 워크로드의 경우 TPU 구성이 GPU 기준보다 1.82 배 저렴하다고 결론지었습니다. 이 우위는 주로 학습 비용 절감에 기인하며, 추론 비용은 긴 컨텍스트 애플리케이션의 경우 크게 누적됩니다.
하드웨어 적합성: 저자들은 TPU 를 긴 컨텍스트 워크로드와 비용에 민감한 생산 파이프라인에 선호되는 인프라로 위치시키면서, 짧은 컨텍스트 처리량과 생태계의 성숙도(단순한 프로토타이핑) 에서의 GPU 우위를 인정합니다.

본 보고서는 JAX/Tunix 로 포팅하는 엔지니어링 오버헤드가 무시할 수 없으며 (메시 구성, 샤딩 수정, 파이프라인 재구성에 약 1 주일 소요 추정) 라고 강조하지만, 이는 많은 학습 실행에 걸쳐 상각되는 일회성 비용이므로 대규모 LM 적응을 위한 TPU 가 실현 가능하고 더 우수한 대안임을 강조합니다.