Google이 TorchAX 라이브러리로 PyTorch 모델을 JAX로 변환해 TPU 네이티브 실행과 JIT 컴파일 최적화 지원

Context

2024년 HuggingFace가 transformers 라이브러리에서 JAX와 TensorFlow 네이티브 지원을 제거하고 PyTorch에 집중했다. Google Cloud TPU를 사용하는 개발자들이 HuggingFace 모델 생태계를 JAX 기반으로 실행할 직접적인 방법을 잃게 되었다.

Technical Solution

• PyTorch 텐서를 JAX 배열로 변환: torch.Tensor 서브클래스인 torchax.Tensor를 만들어 내부적으로 jax.Array를 보유하도록 구현
• 연산 가로채기(interception) 방식 도입: PyTorch 연산을 호출하면 torchax가 이를 가로채 JAX 동등 연산으로 실행
• Pytree 등록 메커니즘 추가: HuggingFace의 CausalLMOutputWithPast 같은 커스텀 출력 타입을 JAX pytree로 등록하여 JIT 컴파일 지원
• 모델 변경 없이 장치 전환: model.to("jax") 한 줄로 PyTorch 모델을 JAX 실행으로 전환
• JIT 컴파일 및 자동 병렬화 활성화: jax.jit 래퍼와 gSPMD 기반 자동 샤딩으로 TPU/GPU 병렬 실행 제공

Impact

JIT 컴파일을 통해 10~100배 속도 향상을 달성할 수 있다.

Key Takeaway

TorchAX는 모델 코드 수정 없이 PyTorch 생태계를 JAX의 고성능 백엔드와 연결하는 어댑터 패턴으로, JAX 전문성이 없는 개발자도 TPU와 JIT 컴파일 이점을 활용할 수 있게 한다. Pytree 등록 같은 타입 호환성 문제 해결이 프레임워크 간 통합의 핵심 과제임을 보여준다.