Meta RecIR, PyTorch 추천 모델로 GPU와 TPU 학습을 단일 코드로 처리

• Meta가 PyTorch 기반 추천 모델을 GPU와 TPU에서 학습할 수 있는 RecIR 컴파일러를 알파 버전으로 공개함
• PyTorch로 작성한 학습 단계를 FX 그래프로 캡처해 GPU(TorchRec, Inductor) 또는 TPU(torchax, JAX/XLA)용으로 변환함
• 현재까지 지원되는 레퍼런스는 generative-recommenders 저장소의 DLRMv3 모델 학습 하나뿐이며 추론 경로는 실험 단계임
• GPU와 TPU 간 수치 동등성(numerical parity)을 검증하는 도구를 제공하며, 성능 수치(처리량)는 미공개 상태임

기술 아키텍처 및 특징

• 모델은 일반 torch.nn.Module, 옵티마이저는 PyTorch 표준을 유지하며 RecIR는 캡처, 그래프 변환, 상태 관리, 분산, 백엔드 변환을 담당함
• TPU에서는 Pallas 링 시프트(ring shift) 임베딩 조회와 행 단위 샤딩을 사용하며, GPU에서는 FBGEMM TBE와 Inductor를 활용함
• XLA의 고정 모양 요구와 가변 길이 UIH(User Interaction History) 문제를 해결하기 위해 토큰 예산 배치(Token-budget batching) 방식을 적용함

설치 및 제약 사항

• PyPI 패키지 없이 Python 3.12 이상에서 소스 체크아웃으로 설치하며, DLRMv3 학습에는 PyTorch 2.10, TorchRec 1.4, FBGEMM 1.5 조합이 필요함
• 파이썬 데이터 의존 제어 흐름은 자동 캡처가 불가하며, 배치 크기나 테이블 모양 변경 시 재컴파일이 필요함
• 임의의 추천 모델을 즉시 운영에 올리려는 팀에는 API 불안정성 등으로 부적절하며, HSTU/DLRMv3 계열 대규모 모델 비교 학습에 적합함