- Google의 TorchTPU와 PyTorch/XLA는 JAX 스타일의 사고방식을 강요하지 않고도 TPU를 PyTorch의 네이티브 고성능 백엔드로 사용할 수 있도록 해줍니다.
- TPU 아키텍처, XLA 컴파일 및 StableHLO는 특히 분산 학습에서 대규모로 효율적인 고밀도 컴퓨팅 및 집합 연산을 가능하게 합니다.
- 새로운 즉시 실행 모드, 제한된 동적 특성, 그리고 easy-torch-tpu와 같은 생태계 도구는 GPU 중심의 PyTorch 코드를 TPU 클러스터로 마이그레이션할 때 발생하는 마찰을 줄여줍니다.
- Cloud TPU, GKE 및 Vertex AI는 연구 규모부터 포드 규모에 이르는 모든 PyTorch 워크로드를 TPU에서 실행할 수 있는 인프라를 제공합니다.
구글 TPU에서 PyTorch를 실행하는 것은 더 이상 소수의 전문가만 이용할 수 있는 틈새 시장의 실험적인 영역이 아닙니다.구글의 새로운 제품들 사이에 TorchTPU 스택검증된 PyTorch/XLA 프로젝트와 지속적으로 성장하는 툴링 및 프레임워크 생태계 덕분에 TPU에서 모델을 학습하고 서빙하는 것은 NVIDIA GPU에서 작업하는 것만큼이나 자연스러워지고 있습니다. 가장 큰 변화는 이제 고성능, 대규모 확장성, 그리고 훨씬 원활한 개발자 경험을 동시에 추구할 수 있게 되었다는 점입니다.
이 글에서는 PyTorch가 현재 TPU를 어떻게 활용하고 있는지, 그리고 앞으로 이 스택이 어떤 방향으로 나아갈지 심층적으로 살펴봅니다.이 글에서는 TorchTPU 아키텍처, 기존 PyTorch/XLA와의 차이점, 분산 학습, 컴파일 및 하드웨어 사양 작동 방식, 그리고 GPU 중심의 PyTorch 워크플로우를 마이그레이션할 때 이러한 내용이 실제로 어떤 의미를 갖는지 자세히 살펴보겠습니다. LLM, 확산, 또는 대규모 추천 시스템 분야에 종사하는 분이라면, 아래 내용은 TPU 실행 속도를 좌우하는 핵심적인 세부 사항입니다.
지금 TPU에서 PyTorch를 사용하는 것이 중요한 이유
현대 AI 워크로드는 이제 단순한 "머신 한 대에 몇 개의 GPU"로 처리하던 시대를 넘어섰습니다.최첨단 모델은 이제 수만 개의 가속기를 포함하는 클러스터에 걸쳐 확장되어 소프트웨어가 극한의 확장성, 안정적인 분산 실행, 그리고 다양한 칩과 공급업체에서 이식 가능한 성능을 처리하도록 요구하고 있습니다. AI 인프라.
구글의 텐서 처리 장치(TPU)는 이러한 최첨단 기술의 중심에 자리 잡고 있습니다.TPU는 Gemini 및 Veo와 같은 내부 시스템뿐만 아니라 Google Cloud 고객의 학습 및 추론 워크로드의 상당 부분을 구동합니다. 과거에는 TPU가 JAX 및 TensorFlow와 밀접하게 연관되어 있었지만, 더 넓은 생태계가 PyTorch로 크게 표준화되면서 GPU는 "PyTorch + CUDA"를, TPU는 "JAX + XLA"를 의미하는 불편한 분리가 발생했습니다.
구글의 답변은 TPU를 최고급 PyTorch 타겟처럼 느끼게 만들기 위한 총력전입니다.TorchTPU는 최고 수준의 성능으로 네이티브하고 즉각적인 PyTorch 시맨틱스를 제공하는 것을 목표로 하며, PyTorch/XLA는 이미 프로덕션 환경에서 널리 사용되고 있는 강력한 지연 컴파일 방식입니다. 이러한 스택을 기반으로 Cloud TPU, GKE, Vertex AI 및 easy-torch-tpu와 같은 커뮤니티 프레임워크는 TPU 클러스터를 1억에서 70억 개 이상의 파라미터를 가진 모델까지 처리할 수 있는 간편하고 스크립트 가능한 인프라로 전환하고 있습니다.
TPU 하드웨어 내부: 단순히 더 빠른 칩 그 이상
TPU 시스템은 기본적으로 칩, 호스트 및 인터커넥트가 긴밀하게 통합된 구조입니다.단순히 단일 가속기 카드가 아닙니다. 이러한 구성을 이해하는 것은 TorchTPU의 설계 방식과 순수 GPU 스택과 컴파일러 선택 방식이 다른 이유를 파악하는 데 필수적입니다.
각 TPU 호스트는 칩 간 상호 연결(ICI)을 통해 여러 TPU 칩에 연결됩니다.ICI는 고대역폭의 2D 또는 3D 토러스 토폴로지를 형성하여 대규모 포드가 단일 논리 가속기처럼 동작할 수 있도록 합니다. 기존 네트워킹 스택을 통해 그래디언트를 전달하는 대신, 콜렉티브는 이 토러스 상에서 직접 전송되므로 소프트웨어가 콜렉티브를 올바르게 표현하는 방법을 알게 되면 확장성이 훨씬 향상됩니다.
TPU 칩 내부에서 연산은 TensorCore와 SparseCore로 나뉘어집니다.TensorCore는 밀집 행렬 연산에 특화된 단일 스레드 엔진으로, Transformer, CNN 및 대부분의 표준 딥러닝 레이어에 사용되는 핵심 기능입니다. SparseCore는 임베딩, 개더/스캐터, 오프로드된 집합 연산과 같이 불규칙적인 메모리 접근 패턴을 가진 워크로드에 맞게 설계되었습니다.
이 아키텍처는 딥러닝에 매우 적합하지만, 입력 방식에 매우 까다롭습니다.예를 들어, 많은 트랜스포머 구현에서 어텐션 헤드 차원을 64로 하드코딩합니다. 현재 TPU 세대는 128~256 사이에서 최적의 성능을 보이는 경향이 있는데, 이는 헤드 차원을 두 배로 늘리는 것만으로도 행렬 곱셈 효율성과 TensorCore 활용도를 크게 향상시킬 수 있음을 의미합니다. 이식성이 이러한 하드웨어적 현실을 없애는 것은 아니지만, 이러한 현실에 더 쉽게 도달할 수 있도록 해줍니다.
PyTorch/XLA에서 TorchTPU까지: TPU에서 PyTorch를 실행하는 두 가지 상호 보완적인 방법
PyTorch는 PyTorch/XLA(torch_xla)를 통해 현재 TPU에서 실행될 수 있습니다.이 방식은 TPU를 표준 PyTorch 장치처럼 제공하고 내부적으로 지연 실행 방식의 XLA 그래프를 컴파일합니다. 그러나 많은 연구자들은 코드 변경 사항이 이론상으로는 사소해 보이지만, GPU의 즉시 실행 방식과 비교했을 때 동작 차이가 상당히 거슬린다는 것을 발견했습니다.
TorchTPU는 Google이 새롭게 개발한 네이티브 PyTorch 백엔드로, 단순한 래퍼가 아닌 "진짜" PyTorch처럼 느껴지도록 설계되었습니다.PyTorch를 모든 곳에 Lazy Tensor를 사용하는 JAX와 같은 모델로 억지로 맞추는 대신, TorchTPU는 PyTorch의 즉시 실행 및 최신 컴파일 API를 활용합니다. 토치.컴파일. 그것은 개인용1 PyTorch의 디바이스 메커니즘이므로, 여러분의 관점에서는 일반적인 방식으로 작업하는 것과 같습니다. 토치.텐서 TPU 상에 존재하는 객체들.
두 접근 방식의 핵심적인 차이점은 실행 방식입니다.PyTorch/XLA는 기본적으로 지연 실행 방식을 사용합니다. 연산들이 그래프를 구축한 후, 학습 루프의 단계와 같은 동기화 장벽에 부딪히면 XLA 컴파일이 트리거됩니다. 이와 대조적으로 TorchTPU는 "적극적 우선(Eager First)" 아키텍처를 채택하여, 연산들을 점진적으로 통합하고 최적화된 서브그래프를 XLA에 전달하는 추가 모드를 제공합니다. 이러한 방식은 표준 PyTorch의 사고방식을 포기할 필요가 없습니다.
클라우드 TPU, GKE 및 Vertex AI: 인프라의 핵심 기반
어떤 PyTorch-on-TPU 스택을 선택하든 그 기반에는 클라우드 TPU 플랫폼이 있습니다.이는 맞춤형 ASIC을 학습 및 추론 모두에 최적화된 확장 가능한 클라우드 리소스로 제공하는 것입니다. 이러한 가속기는 매우 다양한 워크로드에 사용됩니다. 대화형 에이전트코드 생성, 이미지 및 미디어 모델, 음성, 추천 시스템, 개인화 엔진 등이 포함됩니다.
클라우드 TPU는 Google Kubernetes Engine(GKE)과 긴밀하게 통합되어 있습니다.이를 통해 표준 Kubernetes 기본 요소를 사용하여 대규모 PyTorch 작업을 예약할 수 있습니다. 동적 워크로드 스케줄러를 사용하면 필요한 전체 가속기 플릿을 한 번에 요청할 수 있으므로 수동 조정 없이 수천 개의 TPU 칩이 동시에 온라인 상태가 되어 모델을 학습하거나 제공할 수 있습니다.
가장 간편한 도입을 원하는 팀을 위해 Vertex AI는 클러스터 관리의 대부분을 추상화합니다.TPU는 관리형 학습 및 서비스 워크플로에서 대상으로 지정할 수 있으며, 여기에는 다음 상황에서도 마찬가지입니다. PyTorch 기반 모델구글 클라우드는 이러한 유연성(TPU 또는 GPU, 관리형 또는 자체 구축형 쿠버네티스)을 기업과 연구소 모두에서 폭발적으로 증가하는 AI 인프라 수요에 대한 직접적인 해답으로 제시합니다.
TorchTPU의 핵심 철학: “PyTorch 시민의식”
TorchTPU의 핵심 설계 목표는 명확합니다. 즉, 낯선 프레임워크처럼 느껴지지 않고 PyTorch처럼 느껴져야 한다는 것입니다.CUDA GPU에서 모델을 학습시키는 방법을 이미 알고 있다면, 최소한의 코드 수정만으로, 그리고 기존의 사고방식을 바꿀 필요 없이 동일한 학습 스크립트를 TPU로 이식할 수 있을 것입니다.
실질적인 관점에서 이상적인 이민은 거의 우스꽝스러울 정도로 간단해 보인다.일반적으로 글을 쓰는 곳에 device = torch.device('cuda')대신 TorchTPU 모듈에서 TPU 장치를 얻습니다. 개념적으로는 다음과 같습니다. device = tpu.get_device()—그리고 전화하세요 모델.to(장치) GPU에서 작업할 때와 마찬가지입니다. 순방향 패스, 최적화 로직, 그리고 Hugging Face 모델을 호출하는 방식은 변경하지 않고 그대로 사용할 수 있습니다.
이전의 TPU 통합 작업에서는 PyTorch가 JAX를 모방하는 경우가 많았습니다.이전 버전들은 지연 평가 텐서(Lazy Tensor)에 크게 의존하여 정적 그래프 사고방식을 강요했습니다. 이는 PyTorch의 가장 큰 장점 중 하나를 무색하게 만들었습니다. 순방향 전달(forward pass) 도중에 print 문을 삽입하여 형태나 값을 확인할 수 없게 된 것입니다. TorchTPU는 이러한 절충안을 거부합니다. 즉시 실행(eager behavior)을 기본으로 유지하고, 이를 기반으로 성능을 향상시키는 방식을 채택하여, 사용자가 즉시 실행 방식을 포기하도록 요구하지 않습니다.
이 "PyTorch 시민의식" 원칙은 오류 처리에도 적용됩니다.XLA 스택 깊숙이 숨겨진 500줄짜리 난해한 C++ 스택 트레이스 대신, 학습 루프나 모델 정의에서 문제가 되는 줄을 직접 가리키는 깔끔한 Python 트레이스백을 표시하는 것이 목표입니다. 수십억 개의 매개변수를 가진 모델과 수천 개의 TPU를 다룰 때, 이러한 편의성 개선은 오후에 문제를 해결하는 것과 며칠 동안 목적 없이 디버깅하는 것 사이의 차이를 만들어냅니다.
TorchTPU의 실행 모드: 디버그, 엄격, 융합
대규모 융합 그래프를 위해 설계된 하드웨어에서 네이티브 즉시 실행 환경을 제공하는 것은 결코 쉬운 일이 아닙니다.TorchTPU는 공유 컴파일 및 실행 파이프라인을 기반으로 하는 여러 즉시 실행 모드를 제공하여 이 문제를 해결하므로 "작동하게 만드는 것"에서 "빠르게 만드는 것"으로 원활하게 전환할 수 있습니다.
디버그 열성 가장 느리지만 가장 투명한 모드입니다. 이 모드는 다음과 같은 작업을 수행합니다. 한 번에 한 가지 작업씩 TPU로 전송되어 각 연산 후 CPU와 동기화됩니다. 성능 저하는 의도적으로 설계되었으므로 즉각적인 피드백과 명확한 스택 추적을 통해 NaN 값, 메모리 형태 불일치 또는 메모리 부족 오류를 쉽게 추적할 수 있습니다.
엄격한 열정적인 이 단일 작업 디스패치 의미론을 유지하면서 실행 비동기 적으로TPU와 CPU는 사용자 코드가 동기화 지점에 도달할 때까지 병렬로 실행될 수 있으므로, 표준 GPU 기반 PyTorch와 훨씬 더 유사한 환경을 제공하면서도 무거운 그래프 컴파일 요구 사항은 없습니다.
Fused Eager는 성능적인 측면에서 정말 흥미로운 부분이 나오는 부분입니다.TorchTPU는 사용자가 실행하는 연산 스트림을 관찰하고 이를 자동으로 더 크고 밀도 높은 연산 덩어리로 병합한 후 XLA를 통해 TPU로 전송합니다. 이러한 동적 병합 단계는 TensorCore 활용도를 크게 높이고 메모리 대역폭 오버헤드를 줄여 일반적으로 더 나은 성능을 제공합니다. Strict Eager 대비 50~100% 이상 속도 향상 모델 코드 변경 없이.
세 가지 즉시 실행 모드 모두 공통 컴파일 캐시를 공유합니다. 단일 호스트에서 실행되거나 분산 환경에서 여러 호스트에 걸쳐 영구적으로 저장될 수 있습니다. 시간이 지남에 따라 학습 루프가 안정화되고 시스템이 동일한 패턴을 학습하게 되면 컴파일 비용이 감소하고 실행 파일을 빌드하는 대신 텐서를 처리하는 데 더 많은 실제 시간을 소비하게 됩니다.
정적 컴파일: torch.compile, XLA 및 StableHLO
TPU에서 최고의 성능이 필요할 때, TorchTPU는 최신 PyTorch 컴파일 파이프라인에 직접 연결됩니다.모델이나 함수를 래핑할 수 있습니다. 토치.컴파일()Torch Dynamo를 사용하여 FX 그래프를 캡처한 다음 일반적인 TorchInductor 백엔드를 우회하고 대신 XLA에 제어권을 넘깁니다.
XLA를 주요 백엔드로 선택한 것은 TPU의 현실을 고려한 신중한 결정입니다.XLA는 TPU 포드 전반에 걸쳐 수년간 배포되면서 성능이 향상되었으며, ICI 토러스 상에서 밀집된 수학과 집단 통신의 교차점을 깊이 이해하고 있습니다. TorchTPU는 PyTorch 연산자를 XLA에 직접 매핑합니다. StableHLOOpenXLA가 이해하는 텐서 IR을 사용하면 XLA의 최적화 패스를 통해 최적화된 TPU 바이너리를 생성할 수 있으며, 가능한 한 즉시 실행 모드와 동일한 런타임 경로를 재사용합니다.
사용자 정의 연산자의 확장성은 나중에 고려되는 사항이 아닙니다.TorchTPU는 Pallas 및 JAX로 정의된 사용자 지정 커널을 지원합니다. 예를 들어 JAX 함수에 다음과 같은 코드를 추가하여 사용할 수 있습니다. @torch_tpu.pallas.custom_jax_kernel이를 통해 전역 최적화 도구의 이점을 유지하면서 컴파일 경로에 하드웨어에 최적화된 저수준 코드를 삽입할 수 있습니다. 또한 더욱 유연한 커널 개발을 위해 Helion과 같은 추가 DSL을 지원하는 작업도 진행 중입니다.
TPU 기반 분산 PyTorch: DDP, FSDP, DTensor 및 MPMD
대규모 모델은 단일 가속기에서 학습되지 않으며, TorchTPU는 이러한 현실을 최우선으로 고려하여 설계되었습니다.이 라이브러리는 PyTorch의 표준 분산 API와 직접 통합됩니다. 분산 데이터 병렬(DDP), FSDPv2예산 및 디텐서또한 이러한 추상화를 기반으로 구축된 타사 라이브러리를 통해 검증되었습니다.
PyTorch/XLA의 가장 큰 역사적 문제점 중 하나는 엄격한 SPMD(단일 프로그램, 다중 데이터) 편향이었습니다.실제 PyTorch 학습 스크립트에서는 랭크 간에 약간의 차이가 발생하는 경우가 많습니다. 예를 들어 랭크 0은 로깅, 체크포인트 생성 또는 메트릭 처리를 담당하는 반면, 다른 랭크는 순수 연산만 수행합니다. XLA의 전역 그래프 보기에서는 이러한 동작이 어색했고, 개발자들은 종종 차이를 피하기 위해 코드를 다시 작성해야 했습니다.
TorchTPU는 MPMD(다중 프로그램, 다중 데이터) 시나리오를 명시적으로 지원합니다.이 방식은 통신 기본 요소를 신중하게 분리하고 범위를 지정하여 서로 다른 동작으로 인해 정확성이 손상되거나 성능이 저하되는 것을 방지합니다. 가능한 한 XLA가 분산 컴퓨팅의 전체적인 상황을 파악하여 통신과 연산을 연관시킬 수 있도록 하지만, 더 이상 비현실적으로 순수한 SPMD 스타일을 강요하지는 않습니다.
이것이 기존 PyTorch 분산 패러다임과 어떻게 조화를 이루는지가 특히 중요합니다.FSDP, DTensor와 같은 프레임워크 및 TorchTitan과 같은 에코시스템 도구는 다음을 기반으로 합니다. 프로세스 그룹 all-reduce, all-gather, broadcast와 같은 집합 연산을 위한 API입니다. GPU에서 이러한 호출은 일반적으로 NCCL로 해석됩니다. TorchTPU는 ProcessGroup 계층에서 이러한 집합 연산을 가로채어 StableHLO 집합 연산으로 변환하고, TPU 하드웨어와 ICI 토러스가 이를 네이티브로 실행합니다. FSDP 또는 DTensor의 관점에서는 아무것도 변경되지 않았으며, 단지 다른 백엔드를 보게 될 뿐입니다.
PyTorch/XLA: 지연 실행, 동기화 지점 및 실용적인 팁
TorchTPU가 장기적인 완전 네이티브 솔루션이지만, 현재로서는 PyTorch/XLA가 TPU에서 PyTorch를 실행하는 데 중요한 도구로 남아 있습니다.CUDA의 즉시 실행 방식에 익숙하다면, PyTorch/XLA에서 가장 큰 개념적 변화는 텐서가 게으른작업은 그래프로 기록되며, 실제 실행 및 컴파일은 명시적 또는 암묵적으로 동기화할 때 발생합니다.
동기화 지점은 PyTorch/XLA가 컴파일 및 실행을 위해 구축된 그래프를 XLA에 전달하는 지점입니다.일반적인 장애 요인으로는 다음과 같은 전화 통화가 있습니다. torch_xla.sync() 또는 다음과 같은 상위 수준 유틸리티 xm.optimizer_step(optimizer)이는 분산 환경에서 최적화 프로그램을 단계별로 실행하고 장치 간 기울기를 동기화하는 두 가지 기능을 모두 제공합니다.
이러한 게으른 모델은 성능에 중대한 영향을 미칩니다.주어진 그래프(또는 새로운 입력 형태를 가진 그래프)를 처음 실행할 때는 컴파일 비용이 발생하지만, 구조가 안정적으로 유지되는 한 이후의 반복 실행은 훨씬 빨라집니다. 이것이 바로 형태 안정성(고정된 시퀀스 길이, 일관된 배치 크기)이 PyTorch/XLA 워크로드에 매우 중요한 이유이며, 바로 이 때문입니다. 패딩 입력값을 고정 크기로 이는 매우 흔한 패턴입니다.
PyTorch/XLA의 멀티 프로세스 학습은 자체적인 편의 도구를 사용합니다.일반적으로 핵심 학습 함수(예: _mp_mnist_fn) 그리고 이를 다양한 기기에서 실행하세요 torch_xla.launch데이터 로딩은 다음을 통해 관리됩니다. torch_xla.distributed.parallel_loader.MpDeviceLoader이 함수는 표준 PyTorch DataLoader를 사용하여 각 프로세스가 고유한 데이터 조각을 볼 수 있도록 보장하는 동시에 적절한 TPU 장치로 배치를 미리 가져옵니다.
TPU에서의 데이터 로딩, 분산 실행 및 AMP
효율적인 입력 파이프라인은 GPU에서와 마찬가지로 TPU에서도 매우 중요합니다.PyTorch/XLA에서, MpDeviceLoader 호스트 측 데이터 로딩과 디바이스 측 실행이 중첩되어 배치를 TPU에 직접 공급함으로써 가속기가 새 데이터를 기다리는 동안 발생하는 장시간의 유휴 시간을 방지할 수 있습니다.
분산 학습의 경우, xm.optimizer_step(optimizer)는 일반적인 옵티마이저 단계보다 더 많은 작업을 수행합니다.이 함수는 여러 기기에서 그래디언트 전체 리듀스를 수행하고, 평균값을 계산하고, 가중치 업데이트를 적용하고, 필요한 동기화를 처리하므로 일반적으로 각 반복마다 별도의 명시적 동기화 호출이 필요하지 않습니다. 로깅 도우미 함수도 제공됩니다. xm.is_master_ordinal(local=False) 중복을 방지하기 위해 메트릭 및 체크포인트 처리를 하나의 프로세스만 수행하도록 하십시오.
자동 혼합 정밀도(AMP)는 TPU에서와 GPU에서 약간 다르게 작동합니다.TPU는 기본적으로 지원합니다. bfloat16 (BF16)BF16은 float16보다 훨씬 넓은 지수 범위를 제공하며 일반적으로 안정성을 위해 명시적인 손실 스케일링이 필요하지 않습니다. PyTorch/XLA는 PyTorch AMP를 확장하여 필요한 경우 BF16과 FP32 간의 자동 매핑을 지원하므로 TPU에서 혼합 정밀도 학습을 간편하고 안정적으로 수행할 수 있습니다.
모델 저장에도 TPU에 특화된 모범 사례가 있습니다.전화할 수는 있지만 토치.저장 디바이스 텐서의 경우 일반적으로 다음과 같은 방법이 권장됩니다. 직렬화 전에 상태 사전을 CPU로 이동합니다. PyTorch/XLA를 사용하면 일반 GPU 머신과 같은 비TPU 하드웨어에서도 더 쉽게 다시 로드할 수 있습니다.
Easy-torch-tpu 및 실제 TPU 훈련 프레임워크
공식 스택 외에도 커뮤니티는 TPU를 더 쉽게 도입할 수 있도록 고수준 프레임워크를 구축하고 있습니다.. 한 가지 예는 aklein4/easy-torch-tpu이는 Google Cloud TPU 클러스터에서 PyTorch/XLA 워크플로우를 간소화하기 위해 특별히 개발된 경량 학습 프레임워크입니다.
Easy-torch-tpu는 Hypercomputer/torchprime과 같은 크고 경직된 코드베이스에 비해 더 간단하고 유연한 대안으로 자리매김하고 있습니다.이 제품의 설계 우선순위는 명확합니다. 쉬운 설치, 간편한 맞춤 설정, 그리고 원활한 통합입니다. gcloud ssh-가상 클러스터 워크플로우. 이는 의도적으로 "학술적 규모" 실험, 즉 약 32~64개의 TPU 칩에서 10억~1억 개의 매개변수 범위를 갖는 모델을 대상으로 합니다.
확장성은 서브클래싱과 설정 파일을 통해 처리됩니다.새로운 서브클래스를 추가하면 자체 아키텍처, 학습 루프, 옵티마이저, 데이터 로더는 물론 사용자 지정 샤딩 및 재물질화 전략까지 연결할 수 있습니다. 이를 통해 프레임워크의 분산 및 로깅 스캐폴딩을 재사용하면서 자유롭게 실험할 수 있습니다.
이 프레임워크는 핵심 생태계 도구와 긴밀하게 통합됩니다.Weights & Biases 지원으로 실험 추적이 간편해졌으며, Hugging Face 통합을 통해 데이터셋 로딩, 사전 학습된 체크포인트 가져오기, 그리고 표준 GPU 기반 PyTorch에서 실행할 수 있는 모델 저장 작업이 간소화되었습니다. 저장소에는 설치 문서와 시작 예제가 포함되어 있으며, 커뮤니티 피드백을 반영하여 지속적으로 발전하고 있습니다.
제한 사항, 디버깅 및 성능 문제점
이러한 모든 개선에도 불구하고 TPU에서 PyTorch를 실행하는 것은 아직 완전히 원활하지는 않습니다.문제가 발생할 수 있는 부분을 이해하면 대규모 모델이나 동적 워크로드를 처리할 때 많은 시간을 절약할 수 있습니다.
그래프 재컴파일은 숨겨진 성능 저하 요인 중 하나로 남아 있습니다.동기화 시점 사이에 계산 그래프나 입력 형태가 변경될 때마다 XLA는 재컴파일을 해야 할 수 있으며, 이로 인해 눈에 띄는 일시 정지가 발생할 수 있습니다. 이는 특히 언어 모델링 및 생성 워크로드에서 흔히 사용되는 가변 길이 시퀀스 또는 적응형 배치 크기에서 자주 발생합니다.
지원되지 않거나 부분적으로만 지원되는 연산자는 조용히 성능을 저하시킬 수 있습니다.PyTorch/XLA와 TorchTPU는 광범위한 연산자 지원을 목표로 하지만, 일부 ATen 연산자는 아직 네이티브 XLA 최적화가 지원되지 않을 수 있습니다. 이러한 경우 실행이 CPU로 전환될 수 있는데, 이는 기술적으로는 올바르지만 속도가 몇 배나 느려질 수 있습니다. 내장 디버깅 유틸리티 및 메트릭(예: torch_xla.debug.metrics) CPU 폴백이나 예기치 않은 재컴파일이 발생하는 위치를 파악하는 데 도움이 됩니다.
Nsight나 nvprof 같은 기존 GPU 프로파일링 도구는 TPU 커널 내부를 볼 수 없습니다.대신 XLA 전용 프로파일링 후크, TPU 런타임 메트릭 및 상위 수준 로깅을 활용하여 병목 현상을 파악합니다. 많은 팀에서 모범 사례(정적 형태 유지, 신중한 데이터 로딩, 재컴파일 모니터링)를 적용하면 예측 가능한 성능을 빠르게 달성할 수 있다는 것을 알게 됩니다.
구글의 컴파일러 로드맵은 이러한 문제점들을 명시적으로 해결하고자 합니다.XLA 내의 고급 제한된 동적 처리(bounded dynamic) 작업은 모델이 새로운 컴파일을 발생시키지 않고도 다양한 시퀀스 길이와 배치 크기를 처리할 수 있도록 하는 것을 목표로 합니다. 또한, 사전 컴파일된 TPU 커널 라이브러리가 점차 확장됨에 따라 새로운 그래프의 첫 번째 반복에서 콜드 스타트 지연 시간을 대폭 줄이는 것을 목표로 합니다.
로드맵 및 생태계: TPU에서 원활한 PyTorch 사용을 향하여
앞으로 구글의 TorchTPU 로드맵은 야심차며 더 넓은 PyTorch 생태계와 긴밀하게 연계되어 있습니다.공개 GitHub 저장소가 계획되어 있으며, 여기에는 광범위한 문서, 아키텍처 튜토리얼, 학습 및 서비스 시나리오를 모두 포괄하는 재현 가능한 예제가 포함될 예정입니다.
PyTorch의 Helion DSL과의 통합이 곧 이루어질 예정입니다.이를 통해 개발자는 XLA 또는 하드웨어별 코드의 가장 깊은 계층에 접근하지 않고도 사용자 지정 TPU 커널을 작성할 수 있는 옵션이 확장될 것입니다. 동적 형상에 대한 네이티브 수준의 기본 지원이 제공됩니다. 토치.컴파일 이는 또한 우선순위이며, 현대 순차 기반 모델의 현실을 반영합니다.
다중 대기열 지원은 또 다른 핵심 중점 분야입니다.많은 프로덕션 PyTorch 코드베이스는 비동기 실행 패턴과 분리된 메모리/컴퓨팅 스트림에 크게 의존합니다. 이러한 관용구를 대대적인 리팩토링 없이 TPU에 깔끔하게 매핑할 수 있다면 규모가 크고 성숙한 프로젝트의 마이그레이션 마찰을 크게 줄일 수 있을 것입니다.
생태계 전반에 걸친 심층적인 통합은 이미 진행 중입니다.현재 TPU Pod의 최대 크기까지 강력한 확장성을 검증하고 vLLM 및 TorchTitan과 같은 주요 PyTorch 기반 시스템과 연동하기 위한 노력이 진행 중입니다. 동시에 Google은 Meta 및 PyTorch 커뮤니티와 긴밀히 협력하여 TorchTPU의 핵심 구성 요소를 오픈 소스로 공개하여 도입 및 투명성을 높이는 방안을 모색하고 있습니다.
이 모든 것은 TPU 용량이 급격히 증가하는 더 큰 비즈니스 환경 속에서 이루어지고 있습니다.구글 클라우드는 수십억 달러 규모의 AI 인프라 계약을 추가로 체결하고 있으며, 앤트로픽은 최대 백만 개의 TPU(약 1기가와트 용량)에 대한 접근을 계획하고 있고, 구글은 심지어 온프레미스 데이터 센터용 TPU를 직접 판매하고 있습니다. TPU가 틈새 시장의 구글 내부 전용 자원이었던 시대는 이미 오래전에 끝났습니다.
종합적으로 보면, TPU 기반 PyTorch는 "특이한 우회 경로"에서 "표준 옵션"으로 놀라울 정도로 빠르게 자리 잡고 있습니다.TorchTPU의 네이티브 즉시 실행 경험, PyTorch/XLA의 검증된 지연 실행, easy-torch-tpu와 같은 프레임워크, 그리고 이들을 둘러싼 풍부한 클라우드 TPU 인프라 덕분에 이제 주류 PyTorch 모델을 (대부분 장치 문자열만 변경하면) 사용 가능한 최대 규모의 AI 슈퍼컴퓨터에서 효율적으로 실행할 수 있습니다. 스택이 새로운 사고방식을 강요하는 대신 익숙한 PyTorch 관용구로 수렴할수록 하드웨어 선택을 근본적인 설계 제약 조건이 아닌 구현 세부 사항으로 간주하는 것이 더욱 현실적이게 됩니다.
