Google, TPU용 네이티브 PyTorch 백엔드 TorchTPU 오픈소스로 공개

TorchTPU는 Google의 TPU에서 PyTorch 모델과 연산을 실행할 수 있게 하는 백엔드임 TPU용 ATen 커널과 XLA 연산 실행을 최적화하는 컴파일 캐시를 포함함 torch.compile() 과 통합해 그래프 모드 실행을 지원함 데이터 병렬화와 텐서 병렬화 기반 분산 학습을 지원하며, Llama, Qwen, ResNet 등 여러 모델의 예제를 제공함 현재 Python 3.11 이상을 지원하며, Apache 2.0 라이선스로 배포됨 TPU에서 PyTorch 실행 TorchTPU는 Google의 Tensor Processing Unit(TPU)에서 PyTorch 모델과 연산을 실행하도록 하는 백엔드임 다양한 PyTorch 연산을 위한 TPU용 ATen 커널을 구현함 컴파일 캐시로 XLA 연산 실행을 최적화하며, torch.compile() 과 통합해 그래프 모드로 실행할 수 있음 분산 학습과 모델 예제 데이터 병렬화와 텐서 병렬화를 활용한 분산 학습을 지원함 분산 학습과 모델 추론을 위한 예제를 포함하며, 대상 모델은 Llama, Qwen, ResNet, DLRM, minGPT, UNet임 디버깅과 벤치마킹을 위한 유틸리티를 제공함 사용자 문서와 실행 환경 사용자 가이드는 대화형 노트북의 정적 버전이며, 대화형 노트북은 저장소의 docs 폴더에 있음 노트북은 Jupyter와 유사한 Marimo 기반으로, 브라우저에서 사용할 수 있음 실행에는 TPU VM이 필요하며, v6e GCE 인스턴스가 가장 간편한 선택임 PyTorch 공식 문서와 TorchTPU 공개 API 문서를 제공함 배포와 빌드 시 주의사항 현재 공개 릴리스는 0.2.0이며, pip 설치와 소스 빌드를 지원함 Python 3.11 이상을 지원하며, 버전 지원 정책은 PyTorch 릴리스 호환성 표를 따름 세부 내용은 Python 버전 지원 정책에서 확인할 수 있음 설치 시 PyTorch CPU 빌드용 추가 인덱스를 지정하거나 CPU 빌드를 별도로 설치해야 함 추가 인덱스 없이 설치하면 CUDA 버전의 torch가 설치됨 빌드와 의존성 관리에는 Bazel을 사용하며, 기본 빌드는 최적화 모드임 비최적화 빌드는 성능을 저하시킬 수 있어 일반적인 wheel 빌드에는 사용하지 않음 wheel 빌드에는 --config=wheel_common 이 필수이며, 공유 XLA/MLIR 백엔드를 하나의 libxla_base.so로 묶도록 함 이 설정이 없으...