구글, 파이토치 TPU 네이티브 지원 'TorchTPU' 공개...기존 코드 최소 수정으로 TPU 슈퍼컴퓨팅 활용 가능
구글이 파이토치 개발자들이 기존 코드를 거의 수정하지 않고 TPU를 활용할 수 있는 TorchTPU를 공개했다.
[한국정보기술신문] 구글이 파이토치 프레임워크가 텐서 처리 장치에서 네이티브하게 실행될 수 있도록 하는 TorchTPU 프로젝트를 24일 공개했다. 이 프로젝트는 파이토치 개발자와 TPU 슈퍼컴퓨팅 인프라 간의 진입 장벽을 제거하는 것을 목표로 한다.
파이토치 친화적 설계 철학
TorchTPU의 핵심 설계 철학은 간단명료하다. "파이토치처럼 느껴져야 한다"는 것이다. 개발자들은 이론적으로 기존 파이토치 스크립트에서 디바이스 초기화를 'tpu'로 변경하기만 하면 핵심 로직을 수정하지 않고도 훈련 루프를 실행할 수 있다.
이러한 설계 철학은 모든 아키텍처 결정의 기준이 되고 있다. 기존 파이토치 워크로드를 최소한의 코드 수정으로 TPU 환경에서 실행할 수 있도록 하여, 개발자들의 학습 곡선을 대폭 줄였다.
3단계 실행 모드 제공
TorchTPU는 세 가지 이거 실행 모드를 구현한다. 첫 번째는 디버그 이거 모드로, 형상 불일치와 메모리 문제를 해결하기 위한 동기식 단일 연산 디스패치를 제공한다.
두 번째는 스트릭트 이거 모드로, 표준 파이토치 동작을 반영한 비동기 단일 연산 실행을 지원한다. 세 번째는 퓨즈드 이거 모드로, 연산 스트림을 더 큰 청크로 자동 융합하여 "스트릭트 이거 모드 대비 50%에서 100% 이상의 성능 향상"을 사용자 설정 없이 제공한다.
최적화된 컴파일 경로
최대 성능을 위해 개발자들은 'torch.compile' 통합을 활용할 수 있다. 시스템은 Torch Dynamo를 통해 함수 그래프를 캡처한 다음, 주요 컴파일러 백엔드로 XLA를 활용한다. 연산들은 XLA의 중간 표현인 StableHLO에 직접 매핑되어 최적화된 TPU 바이너리를 생성한다.
TPU 특화 하드웨어 최적화
아키텍처는 TPU 고유의 특성을 고려한다. TPU는 밀도 높은 행렬 연산을 위한 TensorCore와 임베딩 같은 불규칙한 메모리 패턴을 위한 SparseCore를 특징으로 한다. 칩들은 2D/3D 토러스 토폴로지에서 Inter-Chip Interconnect를 통해 연결된다.
최적의 모델 설계는 CPU/GPU 접근법과 다를 수 있다. 예를 들어 표준 64차원 설정 대신 128-256차원의 어텐션 구성이 TPU 효율성을 더 잘 활용할 수 있다고 구글은 설명했다.
분산 훈련 지원
TorchTPU는 여러 파이토치 분산 패러다임을 지원한다. 분산 데이터 병렬, 완전 샤드 데이터 병렬 v2, DTensor 등이 포함된다. 특히 서로 다른 랭크가 약간 다른 코드를 실행하는 Mixed-Program Multi-Data 실행을 처리하면서도 단일 프로그램 다중 데이터 최적화 이점을 유지한다.
올해 추가 기능 예고
구글은 올해 추가로 공개 GitHub 저장소 출시와 문서화, torch.compile을 통한 동적 형상 지원, 초기 실행 지연을 줄이는 사전 컴파일된 커널 라이브러리 등을 계획하고 있다고 밝혔다. 또한 사용자 정의 커널을 위한 Helion DSL 통합과 vLLM, TorchTitan과의 생태계 통합도 예정되어 있다.
한국정보기술신문 클라우드분과 이준호 기자 news@kitpa.org
